Can Language Models Learn to Skip Steps?

Summary

The paper explores the ability of language models to skip steps in their reasoning processes. The authors introduce a controlled framework to stimulate step-skipping behavior by iteratively refining models to generate shorter and accurate reasoning paths. The study demonstrates that models can develop this ability under guidance, leading to increased efficiency without sacrificing accuracy. The paper presents empirical results showing enhanced generalization capabilities in out-of-domain scenarios after fine-tuning on expanded datasets that include both complete and skipped reasoning sequences.

Strengths

- The empirical results are robust in three domains, showing benefits in efficiency of the proposed method. - The paper is clearly written and well-organized, making it easy to follow the authors' methodology and findings.

Weaknesses

- Only one backbone model is considered. Experiments across model families and model sizes should be considered to show the generalization ability of the proposed methods. - The OOD test is actually the harder in-domain test. I am curious about the across domain effect of the proposed method For example, how's the effect of training on "Analog of Algebra" and test on "Multi-digit Addition", given that the skip ability should be a general ability across different domains? - In methodology: "We begin with a training dataset D0, which contains detailed full-step reasoning answers to the questions." -> How the full-step reasoning data created?

Questions

see weakness

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Limitations

yes

Summary

This paper proposes an iterative training method that helps sequence models learn to skip steps. The method starts from a training set with full-length solutions or mixed with some skipped-length solutions. At each stage a model learns these solutions with the instruction “Solve it in n steps” and is prompted to generate shorter answers. Correct shorter answers are added to the training set. The effect of this approach is tested using LlaMa-7B on three tasks, including algebraic evaluations, multi-digit addition, and a direction inference task.

Strengths

The proposed skip reasoning pipeline is interesting and was evaluated against a diverse set of tasks with different levels of OOD generalization tests. The authors conducted detailed analyses to understand the effect of the training pipeline, e.g., Figure 5 with the multi-digit addition is very informative. The overall presentation is very clear and easy to follow.

Weaknesses

My main concern is the generalizability of this method. As shown in the paper, the model largely benefited from the warm start setup that includes some skipped problems in the first training set, and has trouble generalizing to problems requiring more steps. One interesting generalization test would be to train on a mixture of all three tasks, but withhold adding skipped step instances for one task, and see if the model can generalize skipping steps on the withheld task. The shorter answers generated during the iterative process also don't seem quite “generated by the model itself”, as filtering out correct answers would require oracle knowledge. Is it assumed that correct answers are any exact subset of the full-length solution?

Questions

The accuracy metric measures final answer accuracy, are intermediate steps correct? I'm unsure if it makes sense to read too much into the average step metric when accuracy is low. In figure 4, what does the accuracy look like for only problems where the model skipped steps? What do you think make multi-digit addition and directional inference more difficult than the algebra task? Especially that accuracy and average step for OOD problems are still pretty bad for multi-digit addition even with the warm start and a few iterations in. What is the change in the ratio between D_init and D_skipped over the iteration process? What’s the range of i (in n-i) of the added skip-step instances in D_skipped under warm start? Line 221 has incorrect figure reference.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Noted in the paper.

Summary

This paper proposes to teach LLMs to deliberately skip steps when doing complex tasks involving multi-step reasoning. The authors use self-generated inference paths with fewer steps to fine-tune the models, which is similar to self-distillation. The authors conduct experiments on a few controlled tasks show that the proposed approach can effectively reduces the reasoning steps while maintain performance.

Strengths

1. The idea of teaching LLMs to skip steps following the human reasoning process is intuitive and makes sense. 2. The proposed method is overall technically sound and well described. 3. The paper is in general well-written and easy to follow. 4. Experimental results confirm the effectiveness of the proposed approach, at least on these "artificial" tasks.

Weaknesses

1. The experiments are not solid because the tasks considered in the experiments are very artificial and not representative for real-world reasoning tasks. The paper could be made much stronger by conducting tasks/datasets such as GSM8K/MATH or coding tasks, instead of simple reasoning tasks. Without the empirical study on realistic tasks, it is hard to confirm the contribution and usefulness of the proposed metric.

Questions

N/A

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes

Summary

The paper proposes a method for training an LLM to solve reasoning problems using fewer verbalized reasoning steps than it is naturally encouraged to by a fixed training dataset. The resulting model is shown to maintian or improve performance on in-distribution data and OOD data testing extrapolation w.r.t. length or compositionality, while using fewer reasoning steps at inference time. Analysis shows that performance gains are concentrated around problems requiring an intermediate number of reasoning steps, rather than very few reasoning steps. Experiments are conducted with Llama-2-7b on three synthetic datasets. The method itself works by using warm-start data with mixed length reasoning demonstrations, followed by bootstrapped training data created by controlling model generations with control codes (instructions) combined with filtering model generations for correctness to create new gold data.

Strengths

- Very important: The idea of shortening reasoning steps, particularly to mimic human reasoning that is variable in its verbalized length, is a very interesting and practical direction. - Very important: Results are positive and promising for model generalization at an increased level of efficiency. Particularly interesting are results suggesting that model performance can increase on difficult OOD data by virtue of skipping some reasoning steps. Initially, CoT was found to improve OOD generalization, but it seems that this iteration on CoT could improve OOD performance even more in some situations. - Important: The paper is overall straightforward to read and understand, with only a few exceptions. - Of some importance: The connection to easy-to-hard generalization was interesting to me. That this method could improve OOD performance, specifically length/compositional generalization, was very interesting.

Weaknesses

- Important: What are the instructions at inference time? Do you require a ground-truth number of reasoning steps to run the model at inference time? If so, this important detail is missing from the paper and could make the method difficult to use in practice for problems if it is not know how difficult they are in advance. Would the method be robust to misspecified instructions at inference time? - Important: I find it a little confusing to reconcile the results of Sec. 5.1 with Sec. 5.2. Sec. 5.1 makes it look like using fewer steps greatly hurts model performance, while Sec. 5.2 makes it seem like using fewer steps does not hurt performance (specifically the Warm start rows, relative to Cold start baselines). - Of some importance: The data is a little artifical. There are existing reasoning and compositional reasoning benchmarks that could be appropriate for this work (though they could require stronger models), including SCAN (https://arxiv.org/abs/1711.00350), GSM8k, StrategyQA, and MATH datasets. However, this is not a major weakness as using clean, controlled datasets is advantageous for studying these kinds of phenomena and they enable automatic construction of warm start data.

Questions

- Why keep the cold start data in the training data if the bootstrapped data is good or better? Do you have ablations that suggest what mixture of the data is best? - Suggested experiment: if you could have two models that are similar except for one being better at long-context reasoning, it would be interesting to see how your method affects each model. The reason for this is that compressing reasoning length could be beneficial by virtue of reducing the context length, rather than some other inherent benefit like allowing the model to spend more computation on harder steps. Such an experiment would help disambiguate if the improvement comes from shortening context length or from using fewer steps. - Note L.68-69 is heavily disputed by follow work on ToM, e.g. https://arxiv.org/pdf/2310.19619 - Just so you’re aware, some highly related work has appeared contemporaneously: (1) https://arxiv.org/pdf/2405.14838, (2) https://arxiv.org/pdf/2407.06023 - L.34: use an em-dash rather than single dash here - L.221: Fig7(a) should read Fig4(a)

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Discussion was adequate, but it could be worth mentioning that experiments were conducted with only one model and three synthetic datasets.

Respose to rebuttal

>We only require the step number as input when generating the skipped data for the training set during the iteration phase. At inference time… Thanks for the clarification! This highlights to me that there may not be strong control over how much step skipping the model does. Rather, the model learns to skip some portion of the time based on the training data. I think this is fine. I think the strength of this paper is interesting results around step skipping and model generalization, and not a new method for controlling how many steps a model takes when solving a problem. >We apologize for the confusion regarding the results in Sections 5.1 and 5.2… Thanks this makes sense! >We are actively working on additional experiments using these datasets Great, please do upload these if they finish in time. --- Based on the above discussion, I plan on keeping my score at 7. If the results were especially interesting on additional, harder datasets, I could increase it, but I think the core contribution of the paper is already solid based on the three datasets used.