On the Inductive Bias of Stacking Towards Improving Reasoning

Summary

This paper examines the inductive bias of gradually stacking layers to increase the depth of a smaller model. The proposed stacking variant, MIDAS, enhances training efficiency and discovers a compelling inductive bias that boosts downstream performance, particularly in reasoning tasks.

Strengths

1) This work is well-motivated in its aim to design a strategy for reducing the computational cost of training large models. Like some previous studies, it proposes a straightforward variant of gradual stacking that results in improved training speed and performance on reasoning tasks. 2) They established a link between the surprisingly discovered inductive bias that enhances reasoning performance and the looped transformer, which is specifically designed for such tasks. This work paves the way for further research into understanding the intriguing inductive bias of stacking layers to grow model size, not only for computational gain but also for performance improvement. 3) The paper is well written.

Weaknesses

The benchmarked models of size 1B and 2B are small compared to commonly used models like 7B, 13B etc. It’s hard to conclude yet whether this strategy will scale well. So some scaling study will show whether this work can be really a stronger candidate to replace the baselines.

Questions

1) What is the reason to choose UL2 objective for training with 60% causal LM, 20% prefix LM and 20% span corruption? Do you see similar phenomena by just training with causal LM? Authors are requested to provide some ablation study and discussion about this. 2) From Table 1 we see memorization performance of MIDAS drops compared to baseline. Does the authors have any intuition why this is happening? 3) What is the performance of GRADSTACK for 2B parameters model in Table 1?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

This work shows limited improvement on memorization-based tasks.

Summary

The authors propose MIDAS, an efficient and effective framework for gradually increasing model depth. Their method achieves better performance on some reasoning primitive problems. Also, the authors further provides empirical analysis to support their findings.

Strengths

1. The authors propose MIDAS, a novel variant of gradual stacking, which achieves better training efficiency that baselines. 2. Their experiments show that MIDAS significantly outperforms baselines on certain reasoning primitive problems.

Weaknesses

Can the authors evaluate on more reasoning and common-sense questions for better comparison, such as tasks tested in the Llama paper? Minor: In line 105, “For simplicity, we L is divisible by k” should be “For simplicity, L is divisible by k.”

Questions

Will the authors release the code and dataset to allow the community to further study the interesting phenomenon mentioned in section 5? Can this phenomenon be verified with fewer resources?

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors need to add a separate section for limitations.

Summary

Gradual stacking involves incrementally growing a model by stacking its last few layers to initialize the next stage. A new variant called MIDAS (MIDdle grAdual Stacking) is proposed, which stacks the middle block of layers instead of the last block. This method is found to be more efficient and shows a bias towards improving reasoning tasks. MIDAS demonstrates an inductive bias that enhances performance on downstream tasks requiring reasoning, despite similar or slightly worse perplexity compared to baseline training. This inductive bias was analyzed using reasoning primitives (simple synthetic tasks) which revealed that models pretrained with stacking perform better on these primitives without fine-tuning.

Strengths

1. This paper proposed a novel stacking algorithm for improving reasoning. 2. Experiments show that the algorithm improves performance on four distinct types of reasoning benchmarks. 3. The Deep dive into reasoning section shows very interesting insights on reasoning.

Weaknesses

The authors showed empirical evidence that the proposed algorithm improves reasoning, but reduces memorization. However, the experiments are done on four distinct categories of reasoning benchmarks. I think at least one more reasoning benchmark is needed in each category in order to justify that the improvement difference is indeed due to memorization instead of some dataset artiacts.

Questions

Does it make sense to test on existing synthetic reasoning benchmarks instead of just synthetic reasoning primitives? For example, ProofWriter: https://arxiv.org/abs/2012.13048.

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

Yes, the authors adequately addressed the limitations

Summary

The paper proposes an improvement of the gradual stacking method proposed in Reddi et al. 2023 for efficient training. The improved method relies on an observation that stacking the layers at the end exhibits the similarity between layers at the end and this might be a suboptimal choice but stacking the layers in the middle layers can exhibit the similarity between the layers in the middle layers (similar to looped models). Stacking the middle layers improves the performance on reasoning tasks such as open book QA and math word problem tasks. The paper also shows that even though the validation log perplexity of baseline and gradual stacking methods are similar to the proposed method, the proposed method achieves better performance on the downstream reasoning tasks. At the end, the paper proposes some simple tasks such as induction copying, variable assignment, and pre-school math tasks that might be responsible for improvement on reasoning tasks.

Strengths

- The reasoning behind stacking in the middle rather than at the end is simple and seems effective in experiments. - It is interesting that the paper tried to characterize the improvement of stacking in the middle by characterizing its connection to the looped transformers and showing the benefits in the reasoning-related tasks. The paper also provides extensive experiments to support their claims in this regard. - The paper is well-written and easy to follow overall.

Weaknesses

- In section 5, the paper conjectures that induction copying, variable assignment, and pre-school math tasks are some of the core capabilities of contextual reasoning that can help characterize the inductive bias of the proposed approach. The proposed method MIDAS achieves a significant improvement in these tasks. However, the improvement of MIDAS on downstream tasks (open book QA and math word problems) is marginal compared to the improvements on these tasks. This discrepancy raises the question if there are synthetic tasks for which MIDAS is not performing as well as the baseline which is affecting the performance on downstream tasks. - There are some typos in the paper: - This on line 103? - Line 105?

Questions

See the weakness section.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper discusses some of the limitations in reasonable detail.

Thanks for the rebuttal!

Thank you for the rebuttal! I find the paper interesting, although it is still in its early stages. I keep my score.