Dualformer: Controllable Fast and Slow Thinking by Learning with Randomized Reasoning Traces

## Abstract

Abstract In human cognition theory, human thinking is governed by two systems: the fast and intuitive System 1 and the slower but more deliberative System 2. Recent studies have shown that incorporating System 2 process into Transformers including large language models (LLMs), significantly enhances their reasoning capabilities. Nevertheless, models that purely resemble System 2 thinking require substantially higher computational costs and are much slower to respond. To address this challenge, we present Dualformer , a single Transformer model that seamlessly integrates both the fast and slow reasoning modes. Dualformer is obtained by training on data with randomized reasoning traces, where different parts of the traces are dropped during training. The dropping strategies are specifically tailored according to the trace structure, analogous to analyzing our thinking process and creating shortcuts with patterns. At inference time, our model can be configured to output only the solutions ( fast mode ) or both the reasoning chain and the final solution ( slow mode ), or automatically decide which mode to engage ( auto mode ). In all cases, Dualformer outperforms the corresponding baseline models in both performance and computational efficiency: (1) in slow mode, Dualformer optimally solves unseen 30 × 30 30\times 30 maze navigation tasks 97.6 % 97.6\% of the time, surpassing the Searchformer (trained on data with complete reasoning traces) baseline performance of 93.3%, while only using 45.5 % 45.5\% fewer reasoning steps; (2) in fast mode, Dualformer completes those tasks with an 80 % 80\% optimal rate, significantly outperforming the Solution-Only model (trained on solution-only data), which has an optimal rate of only 30%; (3) when operating in auto mode, Dualformer achieves an optimal rate of 96.6% while utilizing 59.9 % 59.9\% fewer reasoning steps compared to Searchformer . For math problems, our techniques have also achieved improved performance with LLM fine-tuning, showing its generalization beyond task-specific models.

## 1 Introduction

In psychology, the dual process theory Wason and Evans () explains that thought can emerge through two distinct processes. Typically, one process is implicit (automatic and unconscious) and the other process is explicit (controlled and conscious). The famous book *Thinking, Fast and Slow* Kahneman () explores the concept of dual process theory extensively. It describes two different systems of thinking. System 1, corresponding to the implicit process, is fast, intuitive, and emotional. System 2, corresponding to the explicit process, is slower, more deliberative, and more logical. Kahneman applies this theory to a wide range of cognitive and behavioral activities to illustrate how these two systems influence our decision making and reasoning. As outlined in the book, System 1 can lead to biases and is prone to systematic errors due to the lack of voluntary control. As a counterbalance, System 2 is more suitable for tasks that require careful analysis, consideration, and planning.

Transformers Vaswani et al. (), the sequence modeling tool that serves as the cornerstone of foundation models in various domains including Large Language models (LLMs) Dosovitskiy (); Baevski et al. (); Radford et al. (); Touvron et al. (); Hsu et al. (); Touvron et al. (); Dubey et al. (), have been widely used in many works to approach reasoning and planning problems, see e.g., Zhou et al. (); Kojima et al. (); Pallagani et al. (); Valmeekam et al. (); Chen et al. (); Gundawar et al. (); Wang and Zhou (). Interestingly, the dual process theory generalizes to artificial intelligence, where we can categorize the reasoning modes of Transformer into *fast* and *slow*.
In fast mode, a Transformer will directly output a final solution without any reasoning steps,
whereas the intermediate steps of thinking, such as a search trace for finding a short path, will be generated along with the plan in slow mode.
The two inference modes share similar pros and cons with the two thinking systems inherent in us. As discussed in previous works Wei et al. (); Valmeekam et al. (); Lehnert et al. (); Gandhi et al. (); Saha et al. (), models operating in fast mode have a lower computational cost and allow for a quicker response,
yet they fall short in accuracy and optimality compared with slow mode models.
This raises the question:

> *Can we integrate both the fast and the slow modes into Transformer-based reasoning agents, similar to how humans possess two distinct thinking systems, and let them complement each other?*

Multiple approaches have been proposed. A popular paradigm is to start with a pure System 2, and then finetune to make it more efficient like System 1, e.g., to distill System 2 output into System 1 Yu et al. (); Deng et al. (), or to improve the reasoning efficiency of existing system 2 by distillation Wang et al. (); Shridhar et al. () or bootstrapping from symbolic systems (e.g., Searchformer Lehnert et al. (), Stream of Search Gandhi et al. ()). However, in these cases, further fine-tuning is needed, which is computationally expensive, and it is non-trivial to adapt the resulting system to be more like System 1 or System 2 on the fly. To address this problem, Saha et al. () design an explicit meta-controller to switch between two different systems.

In this work, we demonstrate a surprising finding: *a simple data recipe suffices to achieve on-the-fly System 1 and System 2 configuration in solving reasoning tasks*. The resulting model, Dualformer, can be easily configured to execute in either fast or slow mode during inference, and determines which mode to use by itself if not specified. More specifically, to imitate the System 2 reasoning process, our Transformer is trained on data that contains both the reasoning trace and the final solution.
Leveraging the structure of reasoning steps, we design specific trace dropping strategies such that the resulting traces resemble the shortcuts taken by System 1 in the thinking process. In the extreme case, we drop the entire trace and encourage the Transformer to output a final solution directly, bypassing all the intermediate steps. We randomly choose those structured trace-dropping strategies at training time. See Section 4 for the details.

We first apply our framework to train an encoder-decoder Transformer model to solve pathfinding problems, where the trace is generated by the A* search algorithm. We consider two domains: the Maze navigation and the Sokoban game as in Lehnert et al. (), where we use the same tokenization scheme. Interestingly, we have found that these problems are challenging for state-of-the-art LLMs like o1-preview and o1-mini, where the output path often breaks into the walls (See Section˜13 for an example).
In each reasoning mode, Dualformer outperforms the established baselines, achieving stronger results in both the solved rate and optimal rate. Moreover, Dualformer significantly enhances the diversity of the generated plans by identifying a greater variety of unique paths that reach the goal.
Notably, Dualformer is also efficient even when working in slow mode, generating much shorter reasoning traces than the baseline model.
Next, we apply our framework to finetune LLMs for answering math questions. Following Yu et al. (), the training examples are taken from the MATH dataset Hendrycks et al. () where answers are rewritten by a Llama-3.1-70B-Instruct model to include detailed intermediate steps. Likewise, the obtained LLMs demonstrate enhanced efficacy and efficiency.

## 2 Related Work

##### Cognitive Science & Psychology

Wason and Evans () first introduces the dual-process theory, which is then extensively explored and explained in Kahneman ().
This framework has underpinned extensive research in decision-making across various fields, including psychology and behavioral economics Thaler and Sunstein (); Povey et al. (); Tversky and Kahneman (); Kahneman (). In cognitive science, the distinction between System 1 and System 2 has facilitated a deeper understanding of how humans process information and make decisions. Research has shown that System 1 operates automatically and quickly, with little or no effort and no sense of voluntary control, often relying on heuristics that lead to biases Tversky and Kahneman (). Conversely, System 2 allocates attention to effortful mental activities and is associated with the subjective experience of agency, choice, and concentration.

##### Learning to Plan and Reason

Tremendous efforts have been made to enhance the capability of Transformer-based models to plan and reason over a long horizon. Two main types of approaches have been developed. The first type leverages existing LLMs. For instance, researchers taught LLMs to call external existing symbolic solvers, such as those found in Ahn et al. (); Besta et al. (); Sel et al. (); He-Yueya et al. (); Liu et al. (); Silver et al. ().
Pallagani et al. (, ) investigate the finetuning of LLMs to use a symbolic solver, ; Hao et al. () fine-tune LLMs to use external tools, and propose to finetune LLMs to do reasoning and planning within a world model.
Among them, several works try to integrate System-1 and System-2 thinking into LLM reasoning. Weston and Sukhbaatar () proposed the System 2 Attention (S2A), a more deliberate attention mechanism aimed at improving the reasoning capabilities of large language models by reducing their susceptibility to spurious correlations in the context. Yu et al. () distill System 2 output into a System 1 model, which aims to compile higher-quality outputs from System 2 techniques back into LLM generations without intermediate reasoning token sequences.

The second type aims to train Transformers from scratch
to plan and reason independently Lehnert et al. (); Saha et al. (); Gandhi et al. (), using task-specific language.
Both Lehnert et al. () and Gandhi et al. () teach LLMs to search by representing the search process in language,
Lehnert et al. () developed Searchformer that is trained to mimic the A* search process for pathfinding problems. Gandhi et al. () applied their model to the Countdown game. In the concurrent work Saha et al. (), the authors trained two separate models and manage them by an external meta-controller. There, one model generates rapid traceless response, while another one responds slower but with trace.
Our research is closely related to the above work, but has some key differences. Unlike Lehnert et al. (); Gandhi et al. () which do not modify the reasoning trace in the training data, our approach involves randomizing the reasoning trace.
Unlike Saha et al. (),
we do not use any explicit controller nor use two networks for each mode. Instead, we integrate both the fast and the slow mode functionalities into a single model.

##### Synthetic Data Generation using LLM

Large language models (LLMs) have been utilized for synthetic data generation in various domains. For instance, Wei et al. (); Longpre et al. (); Taori et al. () introduced a synthetic instruction dataset by sampling from diverse templates containing natural language instructions that outline specific tasks. This approach has also been applied to the visual domain Liu et al. (, ); Zhu et al. (); Brooks et al. (); Peng et al. ().
To improve the performance of LLM to answer math questions, Yu et al. () developed a method to rewrite, verify, and augment the original MATH dataset Hendrycks et al. () using specialized prompts. Similar methodologies have been further explored in other studies, including those by Yuan et al. (); Luo et al. (); Lee et al. (); Yue et al. ().

## 3 Preliminaries

Figure: Figure 3.1: An example Maze navigation task and one optimal plan obtained by the A* algorithm. The task and the plan are expressed as a prompt token sequence and a plan token sequence.
Refer to caption: x1.png

Our work builds upon the work of Lehnert et al. ().
To perform planning, we train a Transformer to model a sequence of tokens that sequentially represents
the planning task, the computation of A* algorithm, and the optimal solution derived from the A* search.
The tokenization method is illustrated in Figure 3.1. As a toy example, we consider a navigation task in a $3\times 3$
Maze where the goal is to find one shortest path from the start cell to goal cell, without hitting a wall cell.
The A* algorithm has successfully determined an optimal plan. We use a sequence of tokens to express both the task and the Maze structure, which also serves as the *prompt* for Dualformer. The solution is depicted by the plan token sequence that describes the path using coordinates.
The A* algorithm generates a search trace sequence that records the search dynamics performed, displayed in Figure 4.1. Recall that the A* algorithm is a pathfinding algorithm on a weighted graph. The create clause adds the node (represented by the subsequent coordinates) into the search frontier, and the close clause adds the node to the closed set. Each clause, either create or close, is followed by the tokens x, y, c0, and c1, representing the node’s coordinates, cost-since-start value, and heuristic values, respectively. For details of A* and the tokenization approach, we refer the readers to Russell and Norvig () and Lehnert et al. (), respectively.

## 4 Structured Trace Dropping and Randomized Training

Searchformer proposed by Lehnert et al. () has proven efficacy in addressing a variety of complex decision-making tasks. However, it still suffers from two important limitations.
Firstly, the model only operates in slow mode and outputs lengthy reasoning chains, which significantly increases the inference time.
While this can be reduced by bootstrapping Lehnert et al. (), an iterative refining technique that consists of cycles of rollouts followed by fine-tuning,
such a procedure incurs significant extra demand on computational resources.
Secondly, Searchformer struggles to generate diverse solutions, where identical rollouts are frequently sampled(^1^11We note that the lack of diversity is not due to the limitation of training data. For all our experiments, the search traces in the training datasets are generated by non-deterministic A* algorithm that randomly breaks cost ties.).
For example, across 1000 $30\times 30$ maze problems we have tested, Searchformer’s reasoning chain contains more than 1500 tokens on average,
and can only find $7.6$ unique feasible paths out of $64$ responses (see Section˜5.1.2).

To address these challenges, we propose a training framework that utilizes randomized reasoning traces.
Our approach is inspired by two lines of work.
First, we have noticed that even though Searchformer is trained on complete A* search traces, it generates shorter traces that are sketching the search process. Second, research has shown that humans often rely on shortcuts and patterns when making decisions, a concept known as System 1 thinking Kahneman ().
These observations, combined with the success of dropout technique Hinton (); Srivastava et al. () that randomly drop units from the neural network during training, motivated us to explore the use of randomized reasoning traces in our framework, and we aim to simplify the A* search trace by exploiting its structured elements and selectively dropping certain parts for each training example. The details of our approach are described below.

Figure: Figure 4.1: An illustration of the structured trace dropping strategies, which consists of four levels. Each level employs a progressively more aggressive dropping approach than the previous one.
Refer to caption: x3.png

As shown in Figure˜4.1, the A* search trace contains both the create and the close clauses, and each clause includes the node’s coordinates and its (estimated) cost to reach the start and the goal locations. To derive Dualformer, we exploit the structure of the search trace and drop certain parts of it for each training example.
There are three natural types of dropping:

- •
D1: drop a close clause
- •
D2: drop the cost tokens in a clause
- •
D3: drop a create clause

Based on them, we develop four levels of dropping strategies, each building upon the previous one. Specifically,

- •
the Level 1 strategy eliminates all the close clauses from a search trace.
- •
The Level 2 strategy goes further by additionally dropping all the cost tokens.
- •
The Level 3 strategy is even more aggressive. It further randomly drops $30\%$ of the create clauses.
- •
Ultimately, the Level 4 strategy drops the entire trace.

Figure 4.1 illustrates our strategies using the previously mentioned Maze task. Intuitively, the Level 1 dropping instructs Dualformer to effectively bypass the close-set computation of A* search,
the Level 2 dropping promotes Dualformer to bypass both the close-set and the cost computation. The Level 3 and Level 4 dropping encourage Dualformer to omit certain or all of the search steps. As we will show in Section 5, these strategies effectively guide Dualformer in learning a more concise and efficient search and reasoning process.

To promote diversity in the training data, we do not perform dropping as a data preprocessing step. Instead, at training time, for each training example within a batch, we randomly sample the dropping strategy
from a categorical distribution $\texttt{Cat}(p_{0},p_{1},p_{2},p_{3},p_{4})$, where $p_{1},\ldots,p_{4}$ are the probabilities of performing Level 1-4 droppings, and $p_{0}$
is the probability of maintaining a complete trace. This training framework enables Dualformer to learn from multiple reduced traces even for a single training example, as the same example might appear in multiple batches.

##### Comparison with Token Masking

Curious readers might already be wondering that whether our training framework resembles the token masking techniques used by famous LLMs including BERT Devlin (); Liu (); Song (); Gauthier and Levy (); Sinha et al. (); Kitouni et al. (). However, there are significant differences that distinguish our approach from those masking techniques. First, standard masking techniques usually mask the input tokens of a sequence uniformly in random. In contrast, our dropping strategies only apply to the search trace. Second, while masked LLMs generally employ bidirectional attention layers and predict the masked tokens, Dualformer uses causal attention layers, and our training objective solely focuses on the next token prediction with the overall goal of improving its reasoning and planning capability. Computationally, our training procedure is also more efficient. Dropping tokens shortens the input sequence and saves computation.
For example, it takes 30 hours to train Dualformer for the $30\times 30$ maze task on 8 Tesla V100 32GB GPUs, while it requires 36 hours if we use the full reasoning trace. We defer the training details to Section 5 and Appendix 7.

### 4.1 Controllable Generation

One appealing property of Dualformer is that it can be easily prompted to operate in either fast or slow generation mode at inference time. The control mechanism is extremely simple: we append bos and a *control* token to the standard prompt (which includes environment and task description),
where the control token is either plan or create.
If we use plan, Dualformer will operate in fast mode and directly output the plan, bypassing the reasoning steps. On the other hand, if we inject create after bos, Dualformer will work in slow mode and generate both the reasoning trace and the final plan. See Appendix 8 for concrete examples. If we only use the standard prompt, Dualformer will mimic the dual process of human decision making—depending on the situation, it generates either types of responses, which correspond to System 1 and System 2 reasoning.

## 5 Experiments

Our experiments are designed to answer the following questions:

- 1.
Does Dualformer outperform corresponding baselines in fast, slow and auto mode? Does it generate more diverse plans?
- 2.
In slow model, does Dualformer lead to faster reasoning, i.e., output a shorter trace?
- 3.
Does the structured trace dropping technique generalize to LLMs trained on natural language datasets?

We answer questions 1 and 2 in Section 5.1, where we train Transformers to solve Maze navigation tasks and the closely related Sokoban games, similar to SearchFormer Lehnert et al. (). To answer questions 3, we finetune LLama-3.1-8B and Mistral-7B models to solve math problems in Section 5.2.

### 5.1 Navigation Tasks: Maze & Sokoban

Following Lehnert et al. (), we consider the Maze and the Sokoban tasks and use the same dataset.
Both datasets contain $10^{5}$ training examples with complete A* search trace. The A* implementation is non-deterministic
where it breaks cost ties randomly and randomizes the order in which child nodes are expanded.
The size of Maze varies from $15\times 15$ to $30\times 30$. For all the Maze tasks, we randomly generate $30\%-50\%$ percentage of wall cells as obstacles, and randomly sample the goal and start locations.
The Sokoban map is of size $7\times 7$ with randomly placed two docks, two boxes, and the worker location.
We also randomly add two additional wall cells to the interior of the map.
For the detailed map generation procedure and example figures, we refer the readers to Appendix 7.1.

We first demonstrate that Dualformer can be explicitly controlled to operate in either fast or slow mode in Section 4.1. It will only output the final plan in fast mode, while a reasoning trace will be generated in slow mode. In Section 5.1.1-5.1.2, we compare Dualformer with the corresponding baselines in each mode, respectively. A variety of metrics, as listed below, are used to systematically evaluate the performance, including the correctness, optimality and diversity of generated plans, length of the reasoning trace, etc. Last, we ablate our design choices in Section 5.1.4.

##### Hyperparameters

We instantiate Dualformer using the same encoder-decoder architecture as in Lehnert et al. (). The encoder is an adaptation of the T5 architecture Raffel et al. () with rotary embeddings,
while the decoder is a GPT style architecture.
We employ a model size of 15M and 46M for the Maze and the Sokoban environments, respectively.
All models are trained on 100k training examples. We trained the model for $4\times 10^{5}$ iterations for the Maze environment and $8\times 10^{5}$ iterations for the Sokoban environment.
We defer the details of model architectures and the other hyperparameters to Appendix 7.

We train Dualformer using the structured trace dropping strategies described in Section 4. For choosing the dropping strategies for each training example, we sweep 3 sets of probabilities (1) $\{p_{0}=0.45$, $p_{1}=p_{2}=p_{3}=1/6$, $p_{4}=0.05\}$, (2) $\{p_{0}=0.8,p_{1}=p_{2}=p_{3}=p_{4}=0.05\}$, (3)$\{p_{0}=0.7,p_{1}=0.05$, $p_{2}=p_{3}=0.1$,$p_{4}=0.05\}$ and choose the set that yields the lowest validation error. The final choices are (1) for Maze and (3) for Sokoban.

##### Baselines

For the fast mode, our baseline is the *Solution-Only* model, which uses the same architecture as Dualformer but
trained on sequence data that only include the optimal final solution, without any reasoning trace. The slow mode baseline is the *Complete-Trace* model trained on data with complete A* search traces. It is referred to as the
*search augmented* model in Lehnert et al. (), and is also the base Searchformer model without search dynamics bootstrapping. We will also compare Dualformer with bootstrapped models in Section 5.1.2. All these models have the same amount of parameters, 15M for Maze problems and 46M for Sokoban problems.

##### Metrics

We evaluate whether a model generates correct and optimal plans using two metrics: *1-Solved-64* and *1-Optimal-64*.
Namely, for each evaluation task (e.g. one maze or one Sokoban game), we randomly sample 64 responses from a trained model.
Each response is parsed and evaluated regardless of the generated trace part.
If any of the 64 plans is correct, i.e. is feasible and reaches the goal location, this task is labelled as success for the *1-Solved-64* metric. If any of the 64 plans
is optimal, this task is labelled as success for the *1-Optimal-64* metric.
We repeat such process for 1000 unseen evaluation tasks and report the average success rate.
To investigate the robustness of each method, we also report metrics *3-Solved-64* and *3-Optimal-64* where a task
is labelled as success if at least 3 plans are correct or optimal. Additionally, we consider the *Success Weighted by Cost* (Wu et al., , SWC) that measures the quality of the resulting plans in terms of their costs, aggregated over individual responses. More precisely: $\text{SWC}=\tfrac{1}{64}\sum_{i=1}^{64}\mathbb{I}(\text{plan}\,i\,\text{is correct})\cdot\tfrac{c^{*}}{c_{i}}$, where $\mathbb{I}$ is the indicator function, $c^{*}$ is the cost of an optimal plan, and $c_{i}$ is the cost of the $i$-th plan. Clearly, the higher the SWC score is, the more optimal the resulting plans are. If all the generated plans are optimal, the SWC score reaches its maximum value 1. Last, to quantify the diversity of the generated plans, we check the number of unique correct plans out of the 64 responses for each task, and report the average number across 1000 evaluation tasks.

#### 5.1.1 Fast Mode

Table 5.3 reports the performance of Dualformer and the baseline Solution-Only model on Maze and Sokoban tasks, respectively.
In terms of generating correct and optimal plans, Dualformer significantly outperforms the baseline in both 1-Solved-64 and 1-Optimal-64 criteria.
It also notably surpasses the baseline in terms of the 3-Solved-64 and 3-Optimal-64 rates, which demonstrates the robustness of Dualformer in plan generation. In particular,
the performance gap increases as the task difficulty increases. For the largest $30\times 30$ Maze, the 1-Optimal-64 rate of Dualformer is $2.8\times$ that of Solution-Only model, and the 3-Optimal-64 rate is $2.97\times$.
Dualformer also achieves much higher SWC score than the baseline, which is above $0.9$ for every environment. This demonstrates that each individual plan Dualformer generated is of high quality, whose cost is very close to the optimal plan.

Dualformer also consistently generates more diverse plans for all the considered problems. In Appendix 9, we show one Maze example and plot the unique correct plans generated by Dualformer and the baseline model.
One interesting observation is that the diversity score of Dualformer, i.e., the average number of unique correct plans out of 64 responses, increases as the Maze size goes up. Intuitively, as the Maze becomes larger, there are more possible routes to reach a single goal location.
This suggests that Dualformer learns the Maze structure, while the Solution-Only model is potentially memorizing the optimal plans as its diversity score is close to 1 for all the Maze sizes.

**Table 5.1: Evaluation performance of fast mode Dualformer. The baseline model is the same architecture trained on the solution only data.**
|  | Method | 1-Optimal-64 / 3-Optimal-64 | 1-Solved-64 / 3-Solved-64 | SWC | Diversity |
| --- | --- | --- | --- | --- | --- |
| Maze 15x15 | Dualformer (fast) | 91.8 / 87.6 | 97.1 / 94.8 | 0.960 | 9.05 |
| Solution-Only | 72.0 / 68.9 | 82.7 / 80.1 | 0.610 | 1.52 |  |
| Maze 20x20 | Dualformer (fast) | 90.9 / 84.0 | 97.0 / 94.0 | 0.960 | 17.27 |
| Solution-Only | 56.3 / 52.0 | 71.9 / 67.5 | 0.690 | 1.52 |  |
| Maze 25x25 | Dualformer (fast) | 83.9 / 72.9 | 95.5 / 90.6 | 0.940 | 21.23 |
| Solution-Only | 39.7 / 34.7 | 60.3 / 55.4 | 0.570 | 1.9 |  |
| Maze 30x30 | Dualformer (fast) | 80.0 / 66.0 | 91.8 / 85.7 | 0.906 | 18.23 |
| Solution-Only | 30.0 / 26.0 | 54.1 / 47.8 | 0.500 | 1.86 |  |
| Sokoban | Dualformer (fast) | 97.3 / 94.4 | 94.8 / 90.0 | 0.970 | 4.92 |
| Solution-Only | 86.8 / 83.4 | 92.8 / 90.0 | 0.919 | 1.24 |  |

#### 5.1.2 Slow Mode

Table 5.3 reports the results when Dualformer operates in slow mode. The corresponding baseline is the Complete-Trace model, which uses the same architecture and is trained on data with complete A* search traces. In addition to the metrics reported before, we report the average length of reasoning traces across the 64 responses, aggregated over all the 1000 evaluation tasks.
The results show that Dualformer achieves both enhanced planning power and reasoning speed. It outperforms the Complete-Trace model for all the correctness and optimality metrics: solved rates, optimal rates,
and SWC. Moreover, the reasoning trace yielded by Dualformer is notably shorter than the baseline model. On average, Dualformer reduces the trace length
by 49.4% across the five tasks. As before, Dualformer also generates more diverse plans compared with the baseline. We refer the readers to Appendix 9 for concrete examples.

##### Comparison with Search Dynamics Bootstrapping

The Complete-Trace model is the base
Searchformer model in Lehnert et al. (), which has also proposed a search dynamics bootstrapping method to enhance
its performance on the Sokoban task, similar to those in Anthony et al. (); Zelikman et al. ().
After training the Searchformer model, we fine-tune it on a newly created self-bootstrapped dataset.
For each Sokoban game in the original dataset,
we generate 32 responses and include the shortest optimal response into the new dataset.
We can repeat this process multiple times. In this way, the Searchformer learns to generate shorter responses.
Table 5.4 compares
Dualformer with Searchformer models fine-tuned up to 3 steps. Dualformer is
comparable or better than bootstrapped models in most of the metrics, while only using fewer than 45.1% reasoning steps.
We note that each bootstrapping step requires rollouting $3.2\times 10^{6}$ total responses and extra fine-tuning of $10^{4}$ iterations,
This means, including the $8\times 10^{5}$ pretraining iterations, Searchformer step 3 requires a total of $8.3\times 10^{5}$ training iterations and $9.6\times 10^{6}$ rollouts, which is expensive in computation.
In comparison, Dualformer only needs a single training stage that consists of $8\times 10^{5}$ iterations, with no additional rollout requirements.

**Table 5.4: Performance for the Sokoban game. Searchformer is fine-tuned up to 3 steps using the search dynamics bootstrapping method.**
| Method | Avg Trace Length | 1-Optimal-64 / 3-Optimal-64 | 1-Solved-64 / 3-Solved-64 | SWC | Diversity |
| --- | --- | --- | --- | --- | --- |
| Dualformer (slow) | 1482 | 94.5 / 87.6 | 97.4 / 94.1 | 0.970 | 4.66 |
| Searchformer Step 1 | 3785 | 94.4 / 91.2 | 95.9 / 92.4 | 0.957 | 2.19 |
| Searchformer Step 2 | 3507 | 94.9 / 91.5 | 96.7 / 92.9 | 0.965 | 2.27 |
| Searchformer Step 3 | 3283 | 94.5 / 91.4 | 96.6 / 94.4 | 0.964 | 2.48 |

#### 5.1.3 Auto Mode: Dual Process

Instead of controlling the inference mode of Dualformer by injecting a control token after bos, we can also sample from it directly, allowing it to freely determine
the mode of operation, similar to the dual process of human decision making. We call this *auto mode* for Dualformer.
Table 5.3 reported the results. The *auto mode* Dualformer also outperforms both the Complete-Trace and Solution-Only model for all the tasks we consider.

#### 5.1.4 Ablation Study

As discussed in Section 4, the randomized traces for training Dualformer results from different trace dropping strategies, and there are numerous ways to combine them.
We hereby ablate the design choices we made.
First, to enable execution in both fast and slow modes, a naïve alternative approach is to train Dualformer using a mixture of solution-only and complete-trace data, i.e. $p_{1}=p_{2}=p_{3}=0$ in our randomization strategy. We refer to such variants as *Mix-$p$* models, where $p$ is the fraction of solution-only data in the training dataset. Note that the Solution-Only model is essentially a Mix-$1$ model, and the Complete-Trace model is a Mix-0 model. Below, we compare Dualformer to Mix-$p$ models for both inference modes.
Second, our dropping strategies are structured in a hierarchical manner. For instance, Level 2 dropping is developed based on the Level 1 strategy. We investigate how the performance changes when we halt the process at a specific level.

##### Comparison with Mix- p p Models

Figure 5.1 compares the 1-Optimal-64 rate of Dualformer and Mix-$p$ models.
We test 8 values of $p$: 0 (equivalent to the Complete-Trace model), $0.05,0.1,0.2,0.4,0.6,0.8$, and $1.0$ (equivalent to the Solution-Only model).
In both inference modes, Dualformer beats all the Mix-$p$ models for all the five tasks we consider.
In fact, Dualformer also outperforms Mix-$p$ models in the other metrics we consider. In particular,
compared with all the other models that can execute in slow mode ($p<1$), it is the “fastest” one:
it generates the shortest reasoning trace. Due to the space limit,
we only plot the 1-Optimal-64 rate here, and refer the readers to Appendix 10
for the detailed discussion.

Figure: (a) Slow mode comparison with Mix-$p$ models.
Refer to caption: plot/slow-mode2.png

**Table 5.5: Comparison of different combinations of trace randomization strategies.**
|  | Method | Avg Trace Length | 1-Optimal-64 / 3-Optimal-64 | 1-Solved-64 / 3-Solved-64 | SWC | Diversity |
| --- | --- | --- | --- | --- | --- | --- |
| Maze 15 x 15 | Dropping Level 1 | 396 | 99.5 / 98.4 | 99.9 / 99.4 | 0.998 | 11.94 |
| Dropping Level 1 + 2 | 324 | 98.9 / 98.2 | 99.7 / 99.1 | 0.996 | 11.87 |  |
| Dropping Level 1 + 2 + 3 | 287 | 99.8 / 99.3 | 100 / 99.9 | 0.999 | 12.71 |  |
| Dualformer (slow) | 278 | 99.6 / 99.2 | 99.9 / 99.9 | 0.999 | 12.54 |  |
| Maze 20 x 20 | Dropping Level 1 | 563 | 99.1 / 96.7 | 99.9 / 97.9 | 0.996 | 15.79 |
| Dropping Level 1 + 2 | 549 | 98.7 / 96.7 | 99.2 / 98.5 | 0.991 | 17.51 |  |
| Dropping Level 1 + 2 + 3 | 542 | 99.5 / 99.9 | 100 / 99.7 | 1.00 | 19.51 |  |
| Dualformer (slow) | 439 | 98.9 / 97.8 | 99.9 / 99.7 | 0.998 | 18.86 |  |
| Maze 25 x 25 | Dropping Level 1 | 735 | 97.9 / 93.4 | 99.0 / 97.1 | 0.989 | 19.46 |
| Dropping Level 1 + 2 | 750 | 99.2 / 96.6 | 99.8 / 99.1 | 0.997 | 23.18 |  |
| Dropping Level 1 + 2 + 3 | 594 | 98.2 / 95.6 | 99.5 / 98.7 | 0.994 | 23.1 |  |
| Dualformer (slow) | 589 | 99.9 / 97.2 | 99.7 / 99.3 | 0.997 | 25.05 |  |
| Maze 30 x 30 | Dropping Level 1 | 1183 | 97.5 / 93.3 | 98.9 / 97.4 | 0.988 | 22.61 |
| Dropping Level 1 + 2 | 966 | 98.1 / 93.4 | 99.6 / 97.7 | 0.995 | 25.47 |  |
| Dropping Level 1 + 2 + 3 | 759 | 96.8 / 92.2 | 99.8 / 98.4 | 0.996 | 24.94 |  |
| Dualformer (slow) | 854 | 97.6 / 93.2 | 99.5 / 98.2 | 0.993 | 25.77 |  |
| Sokoban | Dropping Level 1 | 3000 | 94.2 / 86.6 | 95.6 / 90.9 | 0.950 | 3.90 |
| Dropping Level 1 + 2 | 2272 | 93.9 / 86.7 | 96.5 / 89.8 | 0.960 | 4.41 |  |
| Dropping Level 1 + 2 + 3 | 1638 | 95.5 / 90.7 | 97.7 / 95.0 | 0.970 | 4.39 |  |
| Dualformer (slow) | 1482 | 94.5 / 87.6 | 97.4 / 94.1 | 0.970 | 4.66 |  |

##### Combination of Randomization Strategies

We compare Dualformer with its variants where we halt the dropping strategies at specific levels.
For the Maze environments,
we fix the probability of using the complete reasoning trace for a training example to be $p_{0}=0.5$ for all the variants and vary the other probabilities.
For the Sokoban environments,
we fix the probability of Level 1 dropping to be $p_{0}=0.05$.
Table 11.1 lists the probabilities of different dropping strategies we use for all the models, and
Table 5.5 shows the results. It should be noted that the variants cannot operate in fast mode, since they are not trained on any Level 4 data. Therefore, we only report the slow mode performance. As we increase the strategy level, the length of the reasoning trace decreases. In regard to the other metrics, the performance of Level 1+2+3 model and Dualformer are comparable. However, Dualformer enjoys the advantage of shorter reasoning trace and the capability to function in fast mode.

### 5.2 Application to LLM Training: Math Reasoning

In this section, we show the efficacy of structured trace dropping techniques for training large-scale LLMs to solve math problems. Particularly, we finetune Llama-3-8B and Mistral-7B models using a dataset that contains a variety of math questions and answers with detailed reasoning steps, where we utilize a trace dropping technique that leverages the specific structure of the reasoning trace for math problems too. We benchmark the resulting models against corresponding base models finetuned on the dataset directly.

##### Dataset

We evaluate all the models using a dataset named Aug-MATH, which is derived from the MATH dataset Hendrycks et al. () that contains 7500 training examples of math questions and solutions, and 5000 testing examples.
Following Yu et al. (), we query the Llama-3.1-70B-Instruct model to rewrite the solutions to include more detailed intermediate steps, following a given format. To encourage the diversity of reasoning traces, we sample 4 LLM responses for each problem using a temperature of 0.7 and top-$p=0.9$. The resulting dataset then contains $30000$ training examples and $5000$ testing examples.
In Appendix 12.1, we show the prompt template that we use for solution rewriting, and a concrete training example before and after rewriting.

##### Structured Trace Dropping and Randomization

The answer of the math questions rewritten by Llama-3.1-70B-Instruct contains 6-10 intermediate reasoning steps on average, where every step might contain multiple sentences.
We use a randomized training procedure similar to the framework proposed in Section 4.
For each of the training examples sampled in a batch, we randomly drop each intermediate reasoning step with probability $p$. Below, we report the results when varying $p$.

##### Hyperparameters

We finetune two base models Mistral-7B and LLama-3-8B for two epochs using a batch size of $32$.
We use the AdamW optimizer Loshchilov and Hutter () with a learning rate of $8e{-6}$ for the Mistral models, and $5e{-6}$ for the Llama-3 models.
Full training details, including how we select the learning rates and other hyperparameters, are deferred to Appendix 7.2.
For fine-tuning the models, we use the same prompt as in Yu et al. (), which is displayed in Appendix 12.2.

##### Evaluation

Similar to Dualformer, we evaluate the models in both the fast and the slow mode, where the LLM is required to output the final solution directly or to
solve the problem step by step. Following Yu et al. (), we use zero-shot prompting when evaluating the models. See Appendix 12.2 for the prompts we use in each mode.
We consider the *Greedy@1* metric Dubey et al. (); Yu et al. (), where for each question we generate 1 response using a temperature of 0 and verify the correctness.
We also report the *pass@20* metric Chen et al. (), where we randomly sample 20 responses using a temperature of 0.5.
For reference, we also report the results of fine-tuning those models on the original MATH dataset.

##### Results

The results are presented in Table 5.6. We test four values of $p$: $0.1,0.2,0.3$ and $0.4$. Our results show that the proposed training strategy makes both LLMs more effective and efficient. We first inspect the results of the Mistral-7B models. For the slow mode inference, fine-tuning the model with trace dropping and randomized training improves upon the baseline model that is directly finetuned on Aug-MATH dataset.
The absolute Greedy@1 metric improved by 1.7% when $p=0.1$ (which amounts to 10% relative performance improvement), and 0.9% with $p=0.2$ and $0.3$, and 0.1% when $p=0.4$.
Our models are also outperforming the baseline model for the Pass@20 metric when $p=0.1,0.2$ and $0.3$, where the absolute correct rate
increases to $61.9\%$. Under both evaluation schemes, the average length of the reasoning trace goes down as $p$ goes up.
Similarly, for inference in the fast mode, our models also achieve higher correct rate.
A similar trend of performance improvement also holds for the Llama-3-8B models.
Finally, for reader’s reference, we include the results of both the Mistral-7B and the Llama-3-8B models finetuned on the original MATH dataset. The performance of those models clearly lags behind(^2^22The trace generated by models trained on MATH is much shorter, because the answers in MATH do not contain as many intermediate steps as those in Aug-MATH.).

**Table 5.6: The performance of Llama-3 and Mistral models finetuned to solve math problems.**
| Model | Dataset & Dropping Prob | Greedy@1(%) (slow / fast) | Trace Length | Pass@20(%) (slow / fast) | Trace Length |
| --- | --- | --- | --- | --- | --- |
| Mistral-7B | Aug-MATH (baseline) | 16.9 / 9.6 | 527 / - | 59.6 / 29.8 | 521 / - |
| Aug-MATH (p=0.1) | 18.6 / 11.3 | 508 / - | 61.6 / 32.0 | 479 / - |  |
| Aug-MATH (p=0.2) | 17.8 / 11.2 | 477 / - | 61.4 / 31.9 | 470 / - |  |
| Aug-MATH (p=0.3) | 17.8 / 11.8 | 497 / - | 61.9 / 31.7 | 466 / - |  |
| Aug-MATH (p=0.4) | 17.0 / 11.0 | 434 / - | 56.4 / 28.9 | 397 / - |  |
|  | MATH | 13.1 / 8.5 | 290 / - | 53.0 / 29.4 | 227 / - |
| Llama-3-8B | Aug-MATH (baseline) | 19.7 / 13.1 | 548 / - | 62.7 / 35.6 | 535 / - |
| Aug-MATH (p=0.1) | 20.1 / 13.3 | 544 / - | 63.4 / 36.2 | 522 / - |  |
| Aug-MATH (p=0.2) | 20.5 / 13.8 | 525 / - | 63.9 / 36.7 | 497 / - |  |
| Aug-MATH (p=0.3) | 20.5 / 13.5 | 515 / - | 63.4 / 37.5 | 474 / - |  |
| Aug-MATH (p=0.4) | 20.4 / 13.5 | 490 / - | 63.4 / 37.2 | 450 / - |  |
|  | MATH | 13.3 / 12.6 | 432 / - | 52.8 / 35.5 | 332 / - |

## 6 Conclusion

We present a simple and easy-to-implement framework for training Transformers to solve reasoning and planning tasks.
We carefully probe the structure of the reasoning traces and design corresponding dropping strategies that imitate the shortcuts in human thinking process.
By randomly applying the dropping strategies to training examples,
the resulting model, Dualformer,
can be controlled to execute in either fast or slow reasoning mode, or in auto mode where it decides the mode to engage. Dualformer achieves
enhanced performance for the maze navigation tasks and the Sokoban game, while reducing the number of reasoning steps required.
Remarkably, our approach is not limited to training task-specific models from scratch. We apply techniques in the same spirit to finetune LLMs to answer math questions and obtain improved performance.
Last, the proposed framework also reduces computation consumption as the input sequences are shortened after trace dropping.
Future work could investigate whether our approach helps models scale, and explore methodologies such as curriculum learning and hierarchical planning to adapt Dualformer for more complex tasks.

## 7 Network Architecture and Hyperparameters

We use the same encoder-decoder Transformer architecture as in Lehnert et al. () for Dualformer. It first converts each token into a one-hot vector, which is then transformed into a set of vectors through the embedding layer.
The embedding vectors then go through the subsequent layers shown in Lehnert et al. (, Figure 4).
We use RoPE embeddings for positional encoding, and no dropout is used in our architecture.

For the model size, architecture parameters, batch size, we use the same setup as in as Lehnert et al. (). Specifically, for the Maze tasks, we use a 15M-parameter model which consists of 3 attention heads and 6 layers with hidden size 64.
We optimize the model by the AdamW (Loshchilov and Hutter, ) optimizer with learning rate 2.5e-4 and batch size 16, where $\beta_{0}$ and $\beta_{1}$ set to 0.9 and 0.99, respectively.
A linear warm-up strategy was employed for the first 2000 gradient updates. Afterward, we use the cosine learning rate scheduler (Loshchilov and Hutter, ).

For the Sokoban tasks, we use a 46M-parameter model which consists of 4 attention heads and 8 layers with hidden size 96. We use batch size 64 and learning rate
7.5e-5, while the other hyperparameters are the same as above.

### 7.1 Maze and Sokoban

We used the same dataset as in Lehnert et al. (), which is available at [https://github.com/facebookresearch/searchformer](https://github.com/facebookresearch/searchformer).
The Maze dataset contains $10M$ examples and the Sokoban dataset contains $10M$ examples and we use the first $100k$ example sorted in accordance to their "id" field in mongodb (following same approach as Lehnert et al. ()).
For reader’s reference, the dataset is generated as follows. For the Maze tasks, 30-50% of the cells were first randomly designated as walls. Next, a start and a goal location were chosen randomly. Then, the $A^{*}$ algorithm was applied to generate an optimal plan.
For the Sokoban tasks, we use a $7\times 7$ grid map where two wall cells were randomly inserted as obstacles. Moreover, two docks, two boxes, and two worker locations were placed randomly. Once a game is generated, it is only added to the dataset if it could be solved by the $A^{*}$ algorithm.

Figure: Figure 7.1: Map of an example Sokoban task in a 7 $\times$ 7 grid world, with two boxes and two docks (figure taken from Lehnert et al. ()).The worker needs to push (and not pull) the boxes from the start locations to the destinations.
Refer to caption: plot/sokoban-7-7-2-2-level.png

Figure: Figure 7.2: Example prompt and response token sequences for the Sokoban task depicted in Figure 7.1 (example taken from Lehnert et al. ()).

### 7.2 Math Reasoning

We use the implementation provided at [https://github.com/meta-llama/llama-recipes](https://github.com/meta-llama/llama-recipes) for finetuning the models.

We train all the models for 2 epochs, using a batch size of 32.
We use the AdamW optimizer with a learning rate of $5\times 10^{-6}$ for the Llama model and $8\times 10^{-6}$ for the Mistral models.
The learning rate is selected as follows. We sweep over 3 values $2\times 10^{-6},5\times 10^{-6},8\times 10^{-6}$ and choose the learning rate that yields the lowest validation loss. We then retrain the models using the selected learning rates on the full training dataset and report the results.
We use the default values for the other hyperparameters. More specifically, we do not use linear rate warmup, weight decay, nor multistep gradient accumulation.
We use betas=$(0.9,0.999)$, eps=$1e-8$, $\gamma=0.85$ (multiplicative step-wise learning rate decay) for AdamW and “packing” as the batching strategy.

## 8 Controllable Generation of Dualformer

Dualformer offers flexible output options for users to choose. In this section, we demonstrate this by an example navigation task in a $15\times 15$ maze, see Figure˜8.1.
The start location is (9, 10) and the goal location is (3, 6). The location coordinates and the maze structure is encoded in the original prompt below.

To configure Dualformer’s operation mode, we only need to append bos and a control token to the original prompt. To output solution only (fast mode), we inject plan. On the other hand, we inject create to let Dualformer output both trace and solution.

### 8.1 Fast Mode

The modified prompt for fast mode is displayed below. The extra tokens are highlighted in purple.

The following box shows the generated tokens, and Figure˜1(a) plot the corresponding path.

Figure: (a) Dualformer operates in fast-mode, directly outputting the final plan.
Refer to caption: plot/control_figA/control_fast_.png

### 8.2 Slow Mode

Similarly, we present the modified prompt for slow mode below. Comparing with the fast mode prompt, the only change is the control token becomes create.

The following box shows the generated output, and Figure˜1(b) plots the corresponding trace and final path.

## 9 Diversity of Generated Plans

The Dualformer outperforms baseline models in discovering unique feasible solutions.
To illustrate this visually, we select one example maze task and generate 64 responses using Dualformer in fast mode.
Figure˜9.1 plots all the unique feasible paths discovered by fast mode Dualformer alongside those found by the Solution-Only baseline (64 responses).
Dualformer (fast mode) identified 42 unique feasible paths, while the Solution-Only model only found 3.
Similarly, Figure˜9.2 compares slow mode Dualformer and the Complete-Trace (Searchformer) baseline. Dualformer (slow mode) discovered 39 unique feasible paths, whereas the Complete-Trace model only found 17.

Figure: (a) Unique feasible paths found by the Solution-Only Model
Refer to caption: x4.png

Figure: (a) Unique feasible paths found by the Complete-Trace Model (Searchformer)
Refer to caption: plot/aa/slow-1_qq.png

## 10 Comparison with Mix- p p Models

**Table 10.1: Comparison of Dualformer and Mix-$p$ models in the slow mode.**
|  | Method | Avg Trace Length | 1-Optimal-64/ 5-Optimal-64 | 1-Solved-64 / 5-Solved-64 | SWC | Diversity |
| --- | --- | --- | --- | --- | --- | --- |
| Maze 15 x 15 | Complete-Trace | 495 | 94.6 / 90.1 | 96.7 /93.0 | 0.964 | 7.60 |
| Dualformer (slow) | 278 | 99.6 / 99.2 | 99.9 / 99.9 | 0.999 | 12.54 |  |
| Mix-5% | 506 | 97.1 / 94.8 | 97.6 / 96.6 | 0.98 | 10.25 |  |
| Mix-10% | 542 | 99.0 / 97.4 | 99.4 / 98.7 | 0.99 | 11.51 |  |
| Mix-20% | 525 | 99.1 / 97.7 | 99.7 / 99.1 | 1.00 | 12.10 |  |
| Mix-40% | 556 | 98.4 / 96.9 | 99.5 / 98.8 | 0.99 | 11.77 |  |
| Mix-50% | 527 | 94.2 / 92.1 | 96.7 / 95.0 | 0.96 | 10.75 |  |
| Mix-60% | 577 | 88.1 / 85.6 | 94.2 / 92.1 | 0.93 | 6.48 |  |
| Mix-80% | 540 | 81.7 / 78.8 | 89.9 / 88.1 | 0.88 | 2.37 |  |
| Maze 20 x 20 | Complete-Trace | 851 | 98.3 / 95.5 | 98.8 / 93.0 | 0.987 | 14.53 |
| Dualformer (slow) | 439 | 98.9 / 97.8 | 99.9 / 99.7 | 0.998 | 18.86 |  |
| Mix-5% | 789 | 98.6 / 98.1 | 99.5 / 99.1 | 0.99 | 17.71 |  |
| Mix-10% | 866 | 98.9 / 98.2 | 99.8 / 99.3 | 1.00 | 18.52 |  |
| Mix-20% | 838 | 97.9 / 95.9 | 99.4 / 98.6 | 0.99 | 17.57 |  |
| Mix-40% | 900 | 93.5 / 90.4 | 97.4 / 95.1 | 0.97 | 16.51 |  |
| Mix-50% | 905 | 92.3 / 89.1 | 96.3 / 94.3 | 0.96 | 16.30 |  |
| Mix-60% | 891 | 72.6 / 67.9 | 86.5 / 82.7 | 0.84 | 3.92 |  |
| Mix-80% | 943 | 62.4 / 56.4 | 80.9 / 75.4 | 0.77 | 2.62 |  |
| Maze 25 x 25 | Complete-Trace | 1208 | 95.2 / 85.7 | 97.0 / 90.4 | 0.968 | 18.85 |
| Dualformer (slow) | 589 | 99.9 / 97.2 | 99.7 / 99.3 | 0.997 | 25.05 |  |
| Mix-5% | 1109 | 97.9 / 95.3 | 99.0 / 97.2 | 0.99 | 21.00 |  |
| Mix-10% | 1264 | 98.7 / 96.3 | 99.5 / 98.0 | 0.99 | 22.33 |  |
| Mix-20% | 1283 | 97.9 / 95.0 | 99.1 / 97.9 | 0.99 | 23.68 |  |
| Mix-40% | 1278 | 90.8 / 87.4 | 96.3 / 93.9 | 0.96 | 22.30 |  |
| Mix-50% | 1334 | 81.0 / 75.5 | 93.3 / 88.9 | 0.91 | 18.02 |  |
| Mix-60% | 1266 | 61.0 / 53.4 | 80.3 / 74.0 | 0.77 | 4.63 |  |
| Mix-80% | 1501 | 57.2 / 50.3 | 78.5 / 70.3 | 0.75 | 3.76 |  |
| Maze 30 x 30 | Complete-Trace | 1538 | 93.3 / 82.4 | 95.9 / 88.1 | 0.964 | 7.60 |
| Dualformer (slow) | 854 | 97.6 / 93.2 | 99.5 / 98.2 | 0.993 | 25.77 |  |
| Mix-5% | 2022 | 96.3 / 92.0 | 98.4 / 96.0 | 0.98 | 24.50 |  |
| Mix-10% | 1720 | 97.7 / 95.0 | 99.2 / 98.2 | 0.99 | 27.17 |  |
| Mix-20% | 1851 | 95.0 / 91.4 | 98.3 / 96.2 | 0.98 | 25.02 |  |
| Mix-40% | 1854 | 64.3 / 56.2 | 83.1 / 76.0 | 0.81 | 12.60 |  |
| Mix-50% | 1652 | 58.0 / 50.8 | 78.6 / 71.9 | 0.76 | 9.66 |  |
| Mix-60% | 1983 | 43.0 / 35.4 | 68.4 / 58.1 | 0.65 | 3.21 |  |
| Mix-80% | 1648 | 47.8 / 38.4 | 71.9 / 61.9 | 0.68 | 3.98 |  |
| Sokoban | Complete-Trace | 3600 | 92.9 / 84.4 | 94.7 / 89.0 | 0.944 | 2.91 |
| Dualformer (slow) | 1482 | 94.5 / 87.6 | 97.4 / 94.1 | 0.97 | 4.66 |  |
| Mix-5% | 3278 | 85.3 / 72.7 | 91.0 / 80.9 | 0.90 | 3.18 |  |
| Mix-10% | 3402 | 88.6 / 77.2 | 94.1 / 87.9 | 0.93 | 4.07 |  |
| Mix-20% | 3331 | 89.0 / 81.3 | 95.7 / 89.1 | 0.95 | 4.22 |  |
| Mix-40% | 3294 | 92.7 / 86.1 | 97.1 / 93.1 | 0.96 | 4.14 |  |
| Mix-50% | 3202 | 92.3 / 87.3 | 96.2 / 93.4 | 0.96 | 4.66 |  |
| Mix-60% | 2594 | 91.1 / 83.2 | 96.4 / 91.0 | 0.96 | 4.48 |  |

**Table 10.2: Comparison of Dualformer and Mix-$p$ models in the fast mode.**
| Maze 15x15 | Method | 1-Optimal-64 / 3-Optimal-64 | 1-Solved-64 / 3-Solved-64 | SWC | Diversity |
| --- | --- | --- | --- | --- | --- |
| Dualformer (fast) | 91.8 / 87.6 | 97.1 / 94.8 | 0.960 | 9.05 |  |
| Mix-5% | 89.8 / 83.2 | 92.7 / 89.0 | 0.92 | 11.72 |  |
| Mix-10% | 80.3 / 75.6 | 90.5 / 86.9 | 0.88 | 7.01 |  |
| Mix-20% | 78.8 / 77.6 | 87.8 / 86.4 | 0.86 | 1.97 |  |
| Mix-40% | 81.2 / 78.8 | 88.6 / 86.9 | 0.87 | 1.73 |  |
| Mix-50% | 79.0 / 76.4 | 87.7 / 85.8 | 0.86 | 1.92 |  |
| Mix-60% | 81.3 / 78.0 | 89.3 / 87.3 | 0.88 | 1.95 |  |
| Mix-80% | 79.8 / 76.6 | 88.6 / 86.0 | 0.87 | 2.12 |  |
|  | Solution-Only | 72.0 / 68.9 | 82.7 / 80.1 | 0.80 | 1.52 |
| Maze 20x20 | Dualformer (fast) | 90.9 / 84.0 | 97.0 / 94.0 | 0.960 | 17.27 |
| Mix-5% | 86.2 / 74.8 | 94.0 / 87.7 | 0.93 | 15.92 |  |
| Mix-10% | 88.6 / 81.4 | 95.2 / 91.7 | 0.94 | 16.04 |  |
| Mix-20% | 65.2 / 60.5 | 81.7 / 77.1 | 0.79 | 2.72 |  |
| Mix-40% | 63.2 / 59.0 | 80.4 / 77.2 | 0.78 | 2.65 |  |
| Mix-50% | 63.3 / 57.9 | 79.6 / 76.0 | 0.77 | 2.39 |  |
| Mix-60% | 66.8 / 63.6 | 82.4 / 80.1 | 0.80 | 2.50 |  |
| Mix-80% | 56.5 / 51.8 | 74.5 / 70.2 | 0.71 | 2.13 |  |
| Solution-Only | 56.3 / 52.0 | 71.9 / 67.5 | 0.69 | 1.52 |  |
| Maze 25x25 | Dualformer (fast) | 83.9 / 72.9 | 95.5 / 90.6 | 0.940 | 21.23 |
| Mix-5% | 73.4 / 56.9 | 87.0 / 77.3 | 0.85 | 13.37 |  |
| Mix-10% | 80.2 / 68.0 | 90.4 / 83.1 | 0.89 | 16.19 |  |
| Mix-20% | 64.8 / 57.5 | 85.0 / 78.0 | 0.81 | 10.51 |  |
| Mix-40% | 50.7 / 46.3 | 73.8 / 68.9 | 0.69 | 3.60 |  |
| Mix-50% | 56.5 / 50.8 | 77.9 / 71.9 | 0.74 | 3.15 |  |
| Mix-60% | 53.6 / 48.9 | 75.9 / 69.8 | 0.72 | 3.07 |  |
| Mix-80% | 48.2 / 44.2 | 69.6 / 64.4 | 0.66 | 2.88 |  |
| Solution-Only | 39.7 / 34.7 | 60.3 /55.4 | 0.57 | 1.9 |  |
| Maze 30x30 | Dualformer (fast) | 80.0 / 66.0 | 91.8 / 85.7 | 0.906 | 18.23 |
| Mix-5% | 66.3 / 46.1 | 83.6 / 72.0 | 0.82 | 11.85 |  |
| Mix-10% | 76.9 / 65.8 | 90.8 / 84.7 | 0.89 | 19.35 |  |
| Mix-20% | 55.4 / 43.5 | 78.3 / 67.4 | 0.75 | 9.97 |  |
| Mix-40% | 43.4 / 37.3 | 69.7 / 63.2 | 0.66 | 3.79 |  |
| Mix-50% | 43.0 / 38.7 | 68.1 / 62.9 | 0.64 | 3.23 |  |
| Mix-60% | 39.2 / 33.8 | 65.2 / 59.0 | 0.61 | 3.09 |  |
| Mix-80% | 38.0 / 32.3 | 62.2 / 55.5 | 0.58 | 2.75 |  |
| Solution-Only | 30.0 / 26.0 | 54.1 / 47.8 | 0.50 | 1.86 |  |
| Sokoban | Dualformer (fast) | 97.3 / 94.4 | 94.8 / 90.0 | 0.970 | 4.92 |
| Mix-5% | 84.0 / 73.0 | 91.0 / 84.0 | 0.90 | 4.30 |  |
| Mix-10% | 94.0 / 87.0 | 97.0 / 94.0 | 0.97 | 5.58 |  |
| Mix-20% | 94.0 / 89.0 | 97.0 / 95.0 | 0.97 | 4.10 |  |
| Mix-40% | 96.0 / 93.0 | 98.0 / 97.0 | 0.98 | 4.15 |  |
| Mix-50% | 97.0 / 96.0 | 99.0 / 98.0 | 0.99 | 3.38 |  |
| Mix-60% | 98.0 / 97.0 | 99.0 / 99.0 | 0.99 | 3.25 |  |
| solution-only | 86.8 / 83.4 | 92.8 / 90.0 | 0.92 | 1.24 |  |

## 11 Comparison with Different Randomization Strategies

In table˜11.1 below, we show the probabilities of different dropping strategies we use for comparing different randomization strategies.

**Table 11.1: The probabilities of the dropping strategies used for each model.**
|  | Model | Probabilities of Dropping Strategies |
| --- | --- | --- |
| Maze | Dropping Level 1 | $p_{0}=0.5,p_{1}=0.5,p_{2}=p_{3}=p_{4}=0$ |
| Dropping Level 1 + 2 | $p_{0}=0.5,p_{1}=p_{2}=0.25,p_{3}=p_{4}=0$ |  |
| Dropping Level 1 + 2 + 3 | $p_{0}=0.5,p_{1}=p_{2}=p_{3}=1/6,p_{4}=0$ |  |
| Dualformer | $p_{0}=0.45,p_{1}=p_{2}=p_{3}=1/6,p_{4}=0.05$ |  |
| Sokoban | Dropping Level 1 | $p_{0}=0.95,p_{1}=0.05,p_{2}=p_{3}=p_{4}=0$ |
| Dropping Level 1 + 2 | $p_{0}=0.85,p_{1}=0.05,p_{2}=0.1,p_{3}=p_{4}=0$ |  |
| Dropping Level 1 + 2 + 3 | $p_{0}=0.75,p_{1}=0.05,p_{2}=p_{3}=0.1,p_{4}=0$ |  |
| Dualformer | $p_{0}=0.7,p_{1}=0.05,p_{2}=p_{3}=0.1,p_{4}=0.05$ |  |

## 12 Details of the Math Reasoning Experiments

### 12.1 Solution Rewriting

#### 12.1.1 Solution Rewriting Prompt

We use the following prompt template for rewriting the solutions via LLama-3.1-70B-Instruct.

#### 12.1.2 Solution Rewriting Examples

Below, we provide an example of a math question along with four generated solutions that have been rewritten by the Llama-3.1-70B-Instruct model.
It is evident from this example that the original math dataset lacks sufficiently detailed solution steps required to solve the question, such as the formula for a geometric series.

### 12.2 Finetuning and Evaluation Prompts

We use the following prompt for finetuning both the Mistral and the Llama model, following Yu et al. ().

We use the following prompt for the slow mode evalaution.

For fast mode evaluation, we force the generation to directly output the final answer by adding the phrase “The answer is: ”.

### 12.3 Example Outputs

We provide a few example math questions and answers randomly sampled from our trained model, alongside the answers output by the baseline model.
Within each answer box, on the left is the baseline model output (long and lengthy) and on the right is our method output (effective and efficient, colored in blue).

## 13 LLM Generated Solutions for Maze

Readers might wonder whether modern state-of-the-art LLMs can effectively solve our maze problem.
To test this, we randomly selected a $30\times 30$ maze problem and ask the o1-preview model to find the shortest path. O1-preview is the latest reasoning model by OpenAI
which operates in slow mode: it spends more time thinking before they respond.
It turns out that these problems are very challenging for LLMs.
As illustrated in Figure˜13.1, the path suggested by o1-preview incorrectly traverses through the maze walls.
In contrast, Dualformer correctly identifies one optimal path that follows through the maze without any errors.

Figure: (a) o1-preview
Refer to caption: plot/o1/o1.png

Below, we also provide the exact prompt we used for o1 models, and their responses.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - What are the specific failure modes in fast mode?
- Does fast mode performance degrade as maze complexity increases (e.g., with a higher density of wall cells or more structured wall arrangements)?
- How would Dualformer perform on a rectangular maze instead of a square one at test time?
- In the example mazes shown, it appears relatively easy for a human to find the optimal path. How does fast mode perform on more challenging mazes where a human might need more time to find a path?
- For the dropping strategies, how were the initial three sets of probabilities selected?
- The experiment with OpenAI o1 models seems to use them in a zero-shot manner. Have you tried using few-shot approaches and testing on smaller mazes?
- Could you elaborate on how the randomized dropping strategies mimic shortcuts in human cognitive processes?
soundness: 2
strengths: In this work, the authors leverage a purely data-driven approach (randomized trace dropping) to enable controllable dual-mode reasoning in a Transformer model.  This approach is computationally efficient and reduces the need for complex architectural changes. In the pathfinding tasks, Dualformer outperforms baseline models in both fast and slow modes, both in terms of the number of optimal solutions found and the diversity of these solutions.. Additionally, the approach proves effective in improving performance on math problem-solving tasks. The methodology, experiments, and results are clearly presented, making the paper easy to follow and understand.
summary: This paper aims to endow a transformer model with both slow and fast reasoning modes using only a simple data recipe. The proposed model -- Dualformer -- is trained using randomized trace dropping, which enables the model to switch between fast (solution-only) and slow (step-by-step reasoning) processing without any architectural modifications. Dualformer is evaluated on maze and Sokoban pathfinding tasks. The model is evaluated on maze and Sokoban pathfinding tasks and is also applied to fine-tune large language models (LLMs) for solving math problems. Experimental results demonstrate that Dualformer achieves superior performance compared to baseline methods.
weaknesses: It is unclear whether the chosen tasks (mazes, Sokoban) encapsulate fast versus slow thinking as humans would apply them. In particular, simply outputting a solution without explicitly listing the reasoning steps does not necessarily capture fast cognitive processing.

It would be valuable to have more extensive analysis of where and why Dualformer’s fast mode fails. Specifically, examining how task complexity—such as increased wall density or structured obstacle layouts—impacts performance could reveal important limitations.

More investigation into the auto mode, specifically on how the model determines which mode to activate under varying maze difficulties, would add depth. Understanding the criteria or patterns that lead to automatic mode-switching would be very useful.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. **Code Availability**: Will the authors release the code to ensure reproducibility? If not, can they provide more detailed pseudocode or implementation guidelines?

2. **Comparison with SwiftSage**: The authors should include a comparison with SwiftSage and discuss the differences and similarities. How does Dualformer improve upon or differ from SwiftSage?

3. **Complex Scenarios**: Can the authors extend their experiments to more complex tasks such as strategic games (e.g., Werewolf) or combinatorial optimization problems? This would better demonstrate the model's versatility and robustness.

4. **Theoretical Justification**: While the empirical results are promising, a theoretical analysis of why and how the randomization of reasoning traces improves performance would strengthen the paper. Can the authors provide more theoretical insights or hypotheses?
soundness: 2
strengths: The idea of combining fast and slow reasoning modes within a single model is novel and inspired by human cognition, which is a strong conceptual foundation.
summary: The paper proposes **Dualformer**, a new Transformer model designed to integrate both fast and slow reasoning modes, inspired by the dual process theory in human cognition. The model is trained on data with randomized reasoning traces, allowing it to dynamically choose between quick, solution-only outputs (fast mode) and detailed, step-by-step reasoning (slow mode).
weaknesses: 1. **Reproducibility**: The paper does not provide code or sufficient details about the implementation, which raises concerns about reproducibility.

2. **Comparison**: There is no comparison with similar methods such as SwiftSage[1], which should be cited and discussed to contextualize the contributions of this work.

3. **Experimental Scope**: The experimental tasks (maze navigation and Sokoban) are relatively simple and may not fully demonstrate the model's potential in more complex, real-world scenarios. More challenging and diverse benchmarks would strengthen the paper.

### Reference:

- [1] Lin, Bill Yuchen, et al. "Swiftsage: A generative agent with fast and slow thinking for complex interactive tasks." Advances in Neural Information Processing Systems 36 (2024).

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. In section 5.2 (Structured Trace Dropping and Randomization), the MATH dataset is constructed by randomly dropping intermediate reasoning steps from the solution. However, in the examples shown in Answers to Q1 and Q2, the model doesn't skip steps but rather presents solutions more concisely. Does this truly demonstrate learned slow and fast thinking strategies, or is it simply a more compact problem-solving approach? And where does the mismatch between the training demo and inference response come from?
2. Following Limitations 3 and 4, the toy example's limitations seems also stem from the generalizability in approach. In the maze example, the procedure first creates a search path using A*, and the model attempts to mimic this search procedure with random dropout. However, in math problems, the baseline provides a direct path with multiple steps, which is then subjected to random step dropping. This raises questions about where the equivalent of the A* algorithm is in MATH problems and how Structured Trace Dropping is implemented on MATH or other general tasks. The methodology appears to be using different approaches for maze and math problems without clear justification for their equivalence.
soundness: 2
strengths: 1. The paper is well written and easy to follow. 
2. The idea of integrating fast and slow reasoning modes from human cognition theory is novel.
summary: Inspired by slow and fast thinking modes in human cognition theory, this paper introduces Dualformer by training models on reasoning traces with randomized dropping, which allows the model to learn efficient reasoning shortcuts. The performance outperforms baselines on a maze environment and MATH dataset.
weaknesses: 1. The motivation is not super convince to me, why LLM needs to integrate both the fast and slow reasoning modes are not clear, it’s not convinced that human takes benefit from the these two modes then LLM should also follow these two modes given the fact that human brain and LLM are running with two distinct systems.
2. The mathematical foundation of the methodology requires more rigorous discussion. Specifically, it's important to explain why dropping random intermediate steps can lead to better performance. This contrasts with findings in recent reasoning papers, for example [1, 2, 3] where learning longer reasoning steps involving (1) different reasoning strategies; (2) self-refinement; (3) backtrack, etc, led to better performance since it provides the model with more flexible function mapping, which is easier to fit compared to direct query-to-solution mapping. Importantly, a section discussing the foundation of the proposed method is necessary.
- [1] Qin, Yiwei, et al. "O1 Replication Journey: A Strategic Progress Report--Part 1." arXiv preprint arXiv:2410.18982 (2024).
- [2] Kumar, Aviral, et al. "Training language models to self-correct via reinforcement learning." arXiv preprint arXiv:2409.12917 (2024).
- [3] Gandhi, Kanishk, et al. "Stream of Search (SoS): Learning to Search in Language." arXiv preprint arXiv:2404.03683 (2024).
3. A major limitation is that the toy example is too simplistic for understanding the method's effectiveness in real reasoning tasks, though it serves to illustrate the problem scope, and there is one table showing results on MATH. To demonstrate that the model truly learns better solution strategies rather than overfitting on synthetic data, one idea is evaluating in out-of-distribution environments. For example, the model could be finetuned on the maze dataset but evaluated on MATH or GSM8K - problems that could benefit from better thinking/reasoning strategies. 
4. Moreover, in Table 5.6, the differences between the baseline and the method proposed in this paper are marginal. This raises questions about whether the results stem from hyperparameter selection, especially since finetuning in the MATH domain uses the same hyperparameters despite different response lengths. Additionally, the evaluation is conducted with only one seed, and the model uses temperature=0.7 and top-p=0.9 for better diversity. While this approach is valuable for studying whether the model can generate diverse solutions, it introduces uncertainty when evaluating the method's performance.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Why did you not test auto mode on the MATH benchmark?
- How did the training differ between Dualformer and finetuned MATH models?
soundness: 3
strengths: - Considering system 1 & 2 while working with transformers isn’t particularly novel, I thought the way it was included while training the transformer on reasoning traces interesting. Beyond that, the paper provided good motivation for why the auto mode strategy could be more efficient than current models. 
- I found the results on the MATH dataset fascinating and was surprised to see a “more real-world” application of this strategy.
summary: - The paper proposes the question: Can we integrate both the fast and the slow modes into Transformer-based reasoning agents, similar to how humans possess two distinct thinking systems, and let them complement each other?
- Main contributions: A framework for training Dualfromer, a Transformer that solves reasoning and planning tasks inspired by the “System 1 & System 2” thesis by Kahneman in _Thinking, Fast and Slow_. 
- One of the tasks is a maze-like environment where the model is trained to find the shortest path between two cells. The model can either output the solutions and reasoning chain or only the final solution. In “auto mode,” it can decide which of both strategies to engage in. 
- Being able to engage in these different modes is an improvement over current models in the cluster, noticeably Searchformer, which is only able to engage in outputting reasoning chain & solution, potentially requiring lots of inference time. The new model improves here by being trained on randomly sampled reasoning traces and dropping traces with various strategies. 
- They also continue evaluating their new framework by training a Dualformer to solve the Aug-MATH benchmark, where the model outperforms finetuned LLMs.
weaknesses: - While I could understand the performance gains of the new model based on some results being higher in Table 5.1-5.3, I found it very hard to make sense of which models exceed where and why. E.g. I understood that “Auto mode” should, in theory, be less exploratory than “Slow mode”, why is the SWC for 30x30 mazes higher? (0.993>0.989). This is, of course, a nitpick, and I don’t expect the reason here to be super relevant, but I would have, in general, appreciated a better overview of the results. Every score improvement has a blue box, which ultimately doesn’t carry a lot of information value (_if everything is highlighted, nothing is_). It could easily be that I don’t understand a fundamental consideration of this yet for them to be all evaluation indistinction, so I am excited to hear your response. 
- I have a similar, though slightly less confusing, nitpick with the results in the MATH benchmark. I appreciate this inclusion of a finetuned Mistral and Llama, though I would have wished to better understand how good this is compared to other sota models.

Meta-review: This paper focuses on tuning LLMs for reasoning. The reasoning traces are generated using A* search and parts of the search trace are randomly dropped out. This method performs well on maze navigation tasks. On more general reasoning tasks where A* is not applicable, like Hendrycks MATH, the fine tuning traces are generated by rewriting supervised chain-of-thought data. The method also does well here compared to baselines but not as significantly as in the maze cases.

Decision: Accept (Poster)