On the Inductive Bias of Stacking Towards Improving Reasoning
Authors: Nikunj Saunshi
Stefani Karp
Shankar Krishnan
Sobhan Miryoosefi
Sashank J. Reddi
Sanjiv Kumar
Venue: NeurIPS 2024 poster (2024)
Keywords: stacking
language model
reasoning
inductive bias
efficient training
Primary area: optimization_for_deep_networks
TL;DR: An intriguing inductive bias of efficient training methods like stacking towards particularly improving tasks that require reasoning, despite having the same pretraining perplexity

Abstract: Given the increasing scale of model sizes, efficient training strategies like gradual stacking have garnered interest. Stacking enables efficient training by gradually growing the depth of a model in stages and using layers from a smaller model in an earlier stage to initialize the next stage. Although efficient for training, the model biases induced by such growing approaches are largely unexplored. In this work, we examine this fundamental aspect of gradual stacking, going beyond its efficiency benefits. We propose a variant of gradual stacking called MIDAS that can speed up language model training by up to 40\%. Furthermore we discover an intriguing phenomenon: MIDAS is not only training-efficient but surprisingly also has an inductive bias towards improving downstream tasks, especially tasks that require reasoning abilities like reading comprehension and math problems, despite having similar or slightly worse perplexity compared to baseline training. To further analyze this inductive bias, we construct {\em reasoning primitives} – simple synthetic tasks that are building blocks for reasoning – and find that a model pretrained with stacking is significantly better than standard pretraining on these primitives, with and without fine-tuning. This provides stronger and more robust evidence for this inductive bias towards reasoning. These findings of training efficiency and inductive bias towards reasoning are verified at 1B, 2B and 8B parameter language models. Finally, we conjecture the underlying reason for this inductive bias by exploring the connection of stacking to looped models and provide strong supporting empirical analysis.

On the Inductive Bias of Stacking Towards Improving Reasoning
Nikunj Saunshi' Google Research nsaunshi wwwpsshh@@ooooe..oooo com
Stefani Karp Google Research stef axekoolll com
Shankar Krishnan Google Research skri Skrisshnan@oogle..o com
Sobhan Miryoosefi Google Research
Sashank J. Reddi Google Research sashank@google.com com
Sanjiy Kumar Google Research sanjivk@google.c com
http:oossfiggoggee....oooo com
Abstract
Given the increasing scale of model sizes, efficient training strategies like gradual stacking [Gong et al., 2019, Reddi et al., 2023] have garnered interest. Stacking enables efficient training by gradually growing the depth of model in stages and using layers from a smaller model in an earlier stage to initialize the next stage. Although efficient for training, the model biases induced by such growing approaches are largely unexplored. In this work, we examine this fundamental aspect of gradual stacking, going beyond its efficiency benefits. We propose a variant of gradual stacking called MIDAS that can speed up language model train- ing by up 40%. Furthermore we discover intriguing phenomenon: MIDAS is not only training-efficient but surprisingly also has an inductive bias towards improving downstream tasks, especially tasks that require reasoning abilities like reading comprehension and math problems, despite having similar or slightly worse perplexity compared to baseline training. To further analyze this inductive bias, we construct reasoning primitives simple synthetic tasks that are building blocks for reasoning and find that a model pretrained with stacking significantly better than standard pretraining on these primitives, with and without fine- tuning. This provides stronger and more robust evidence for this inductive bias towards reason- ing. These findings of training efficiency and inductive bias towards reasoning are verified at IB, 2B and 8B parameter language models. Finally, we conjecture the underlying reason for this inductive bias by exploring the connection of stacking looped models and provide strong supporting empirical analysis.
Introduction
With the advent of very large deep learning models, efficient training reduce the compute and time requirements is becoming increasingly important. Along with efficient optimization procedures, there has been a surge in interest to design efficient training strategies. One practical approach is to use smaller models to initialize larger models. Usually, this results in much faster convergence compared vanilla training [Chen al., 2022, 2016, Gong et al., 2019, Reddi et al., 2023, Wang et al., 2023, Li et al., 2023, Kim et al., 2023, Yao et al., 2024, Wang et al., 2024]. Stacking and growing based approaches have particularly gained traction recently. For instance, gradual stacking [Reddi et al., 2023] is a prominent approach where in each stage the last few layers of the model
"Corresponding author
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Improvement stacking same log perplexity
25
Accuracy gain over baseline
2
15
Gradual
Midas (our method)
stacking
QA
06
pottems
grimEiees
Task groups
Figure I: (a) Pictorial depiction of gradual stacking and MIDAS. (b) Accuracy improvements (in %) for model trained with MIDAS over baseline for various task groups, despite having the same perplexity. For both IB, 2B and 8B models, we see that improvements are mostly positive, and are much larger for tasks that require lot of reasoning.
are stacked onto itself initialize the model's next stage, until the desired depth is reached. This has been shown to significantly speed up BERT pretraining and also has some theoretical justification for the efficiency aspect. While these methods can speed up training, such changes can also induce specific biases into the model. However, the effect of stacking-based approaches generalization remains a fundamental open question and is largely unexplored.
Modern deep learning models, when trained carefully. have been shown exhibit interesting inductive biases, and their success is partially attributed to them. Such biases can arise either from model architecture, optimization techniques, or training strategies, and these biases come various forms including simplicity bias, flatness of learned function, and sparsity. The implicit bias of optimizers, in particular, has been subject extensive research. For instance, the implicit bias of first-order methods like stochastic gradient descent has been studied extensively in overparametrized settings [Gunasekar et al., 2018, Liu al., 2023]. Similarly, the inductive biases of architecture components like self-attention and convolution have also been studied [Edelman et al., 2022, Wang and Wu, 2023]. More recently, there has also been interest in constructs like looped models [Lan et al., 2020, Dehghnni et al., 2018] that share weights across layers. They have been shown to be powerful enough to emulate programmable computers [Giannou et al., 2023] and have the inductive bias to simulate iterative solutions [Yang et al., 2023], thereby yielding models with algorithmic abilities. However, in this vein, very little is known about the implicit biases of newer training strategies (c.g., greedy layerwise training gradual stacking) that gaining popularity.
In this work, we investigate the inductive bias of stacking- approaches beyond training efficiency. We uncover an intriguing phenomenon - pretraining with a variant of stacking is not only efficient, but also has desirable inductive bias towards improving downstream benchmarks. First, through comprehensive empirical analysis, we discover novel variant of gradual stacking called MIDAS (MIDdle grAdual Stacking) which copies the middle block of layers of small network to initialize a larger network (see Figure 1). demonstrate that MIDAS more efficient in training compared to standard training and the previous leading stagewise training approach. However, remarkably, it also yields significantly better performance on many downstream reasoning tasks. For instance, see in Figure 1 that MIDAS has significantly better performance math word problems and reasoning primitives. This performance boost should come as a surprise, since MIDAS uses exactly the same data and fewer training FLOPS compared standard training. In fact, the pretraining perplexity of MIDAS on a validation set matches that of standard baseline training. This strongly suggests that there is some inductive bias for MIDAS play.
In this paper, we formalize and provide strong evidence for such an "inductive bias" MIDAS achieves better downstream evaluations despite performing similarly in terms of pretraining validation perplexity. Thus, the improved quality of MIDAS is not because of better generalization in the pretraining objective, but rather due to its ability to extract more skills and abilities from the pretraining process. This kind of inductive bias phenomenon was first formalized in Saunshi et al. [2022] for contrastive learning and later Liu et [2023] for language modeling on synthetic data. However, this is the first evidence of strong inductive bias for training procedure in real language model

---

training. While our real-world benchmarks already provide strong evidence, in order to better isolate the contributing factors, we construct simple synthetic tasks that are building blocks for reasoning, called reasoning primitives. We find that model pretrained with MIDAS has much better performance on the reasoning primitives than a model obtained through standard pretraining, as evident in Figure 1. In light of the above discussion, we state the main contributions of our paper.
We propose novel variant of gradual stacking, called MIDAS, that achieves better training efficiency than gradual stacking.
Our investigation of the inductive bias in gradual stacking approaches, particularly with MIDAS, reveals surprising benefit: beyond enabling efficient training, it also enhances performance on downstream tasks. This improvement is especially notable in tasks that rely context and reasoning abilities.
We provide strong evidence of the aforementioned phenomenon on several datasets that have previously been used to demonstrate reasoning capabilities.
We construct simple synthetic tasks that are building blocks for reasoning and demonstrate that MIDAS performs significantly better than baseline training on these tasks. These datasets may be of independent interest to the LLM reasoning community.
Finally, we conjecture the reason behind improved reasoning capabilities of MIDAS by presenting connections between gradual stacking and looped models and provide strong empirical evidence to support it.
2
Problem Setup
In this section, we first present the problem setup and background material needed for this paper. Before we discuss the problem setting, we set up the following notation for the rest of the paper.
Notation. For a deep network f. we use fo and #(f) to denote the 7th layer and the number of layers of the network, respectively. With slight abuse of notation, we use fis (where i, E Z to denote the layers between (i - 1) to i b of a deep network f. In other words, Si,, denotes the 7th block of b layers in deep network f. ai is used denote sequence of k: scalars {ar,... {a1,...,,0 a..).
Our goal is to learn a function ʃ X' + y which minimizes the loss E(a.)) {f(x),, y). for some loss function e. y x + R ((0 and data distribution D on X' x V. We are interested in functions of the form f = fi-1 o.... fi where o and L represent function composition and depth of the network, respectively. We use Fz to denote the function class consisting of functions of this form. Given samples from the distribution D, we typically use an iterative stochastic optimizer (e.g., SGD) to learn function that minimizes the loss. We note that the optimization procedure is inconsequential to the arguments in the paper. For standard training, each iteration is of the form:
f" -== --= + A(5"-1, Br, 11),,
(Standard Training)
where B, is a mini-batch from distribution D and A(f B., ne) represents the iterative optimizer update at ft-1 on B, and learning rate ?)- The computation cost and memory requirement for training typically increases linearly with the depth, making even simple algorithms, like SGD, slow for very large models. Throughout this paper, we use T to denote the total number of training iterations.
2.1 k-stage training
Since we primarily focus on stagewise training approaches, is useful to formally define a stagewise training procedure. In contrast to standard training, k-stage training involves dividing the training process into stages, and at each stage, using the the model from the previous stage to initialize the model in the current stage. For simplicity, we assume I. is divisible by k. The following are the key ingredients:
1. Function class across stages. At stage depth of the network at that stage. When d(i) < L, training is more efficient.
function class Fa(i) where d(i) denotes the
2. Training schedules across stages. As training is divided into k stages, we use Tı, ... Tr steps across stages such that Σ
Σí=1 T. = =.
3

---

22
110
330
620
82
625
0.2
E21
33
-3
130
E31
1:30
6:2
6.0
1:00
-11
ALBert layer similarity
(b) GRADSTACK block similarity
(c) MIDAS block similarity
Figure 2: (a) For an ALBert model trained with weight sharing across all layers, we measure the functional similarity between layers by looking at the top 1% activated neurons in each MLP layer and measure the intersection--ver-niiin having the same parameters, a natural functional similarity structure emerges around the middle. (b) For a UL2 model trained with GRADSTACK, we measure the cosine similarity between every pair of layer blocks for the first feedforward layer weights. (c) The same similarity measured for MIDAS. The cosine similarities for stacking based models suggests a strong connection to looped models, and MIDAS has closer similarity structure to ALBert style looped models than GRADSTACK.
(IoU) metric for each pair of layers. Despite all layers
3. Stage initialization. This is the key component of stagewise training. Given network f Fa(i-1) trained in the (i 1)th stage, let M,(f) denote the network initialization for the next stage where Mi : Fa(i1) Fa(s) is growth operator.
Almost all the recent stagewise training procedures are different instantiations of this framework, using different training schedules and stage initializations. We will revisit some prominent instantiations of the framework in the next section.
2.2 Progressive & Gradual Stacking
Progressive and gradual stacking are two special instantiations the aforementioned framework. We provide brief description of these approaches since they are important for our discussion.
Progressive Stacking [Gong et al., 2019]. This is a simple instance of k-sstge training setup where model in the previous stage stacked onto itself initialize the model in the next stage. In particular, (1) depth d(i) 2 d(11 grows exponentially, (2) schedule T; is typically T/k k or proportional to d(i), and (3) the growth function Mi(J) = fo
Gradual Stacking [Reddi al., 2023]. In contrast progressive stacking, gradual stacking linearly increases the model depth by in each stage. It only stacks the last L/k layers of model from the previous stage to initialize the model in the next stage, follows.
1. The depth d(i)= = L-i
x grows linearly with the stage.
2. T. is typically either T/k or allocated proportional or exponential to depth.
3. Mi(fa(i-) ... o fi) = fa(i-1) ... fa (41//1))(1//)+1 fa(i-))
f1. This corresponds to stacking the last L/k layers onto the network to initialize the next stage model.
In the next section, we study novel variant of gradual stacking that enables faster training and exhibits an interesting inductive bias, which we examine carefully.
3 Algorithm: MIDAS
We present the MIDAS algorithm this section. We first discuss the motivation behind this variant of gradual stacking and then formally define the algorithm.
3.1 Motivation
The motivation for MIDAS touches upon two crucial aspects: (a) the role of different layers in deep network and (b) a connection looped models. Before delving into more technical details, important to illustrate these points. We present the case for MIDAS based three observations.

---

Observation 1: gradual stacking breaks the natural role of layers. Recall that gradual stacking initializes a larger model by duplicating and stacking the last block of from the smaller model. Thus in the newly initialized model, the second-last block of b layers will be the same as the last b layers of the smaller model (see Figure 1). Intuitively, this is undesirable since the last few layers have been shown to play a different role compared to other layers for Transformer models [Belrose et al., 2023]. We further validate this in Figure 6. Thus, duplicating the last few layers can break the natural role of layers at the initialization, making it suboptimal choice. However, it is plausible that the similarity structure across layers is broken after continued training and the initialization is inconsequential. The next observation shows that this is not true, and establishes connection to looped models - networks with shared parameters between layers.
gradual stacking leads models resembling looped models. To check the effect of the initialization, we measure the cosine similarity between weights of layers for model pretrained with gradual stacking. In Figure 2b, we observe that indeed the layers continue to have very high cosine similarity at the end of training, thus establishing a connection between stacking and looped models like ALBert [L.n et al., 2020] and Universal Transformers [Dehghani et al., 2018), Unsurprisingly, the similarity structure for gradual stacking is lopsided towards the end of the model, which raises the question: Is this similarity structure natural for for looped models?
Observation 2:
Observation 3: looped models exhibit similarity in the middle. In order to study this, we train prototypical looped model, ALBert, where all layers share the same parameters. Surprisingly. despite parameters being shared, a natural similarity structure emerges between layers: yet again the first and last layers tend to be functionally dissimilar to other layers, whereas the functional similarity between layers the highest the middle (see Figure 2a).
The above observations provides strong motivation for stacking in the middle rather than at the end, thus inspiring our MIDAS algorithm.
3.2
MIDAS algorithm
First we define the following mapping operator that is useful for stage initialization MIDAS. M) = fr. o ... fin/21.6 o f1/22/.0-... f1,b,
(1)
Replication
where n = / is the number of blocks of b layers in deep network f. Note that operator [ b) expands the size of the network by size b. Based on this operator, MIDAS can again be described as a simple instantiation of the k-stage training framework, as seen below. For completeness, the pseudocode for MIDAS in listed in Algorithm 1.
Histogram of gain over Baseline line for different Midas models
Algorithm 1 MIDAS
TyDrQ wio suntect
TyDiQA w contect
TyDIQA wa context mean
Require: Schedule Ti-k Mhrr optimizer update A (see Section 2), data distribution D.
TydiQA
Demoity
Initialize fl.0 E Fr,kk-
for = 1 + h: do
for 1. = 1 - T, do
02-
Sample batch B, from D. fait pa.t-1 + A(fs.1 B., 9()
01
00
end for
Accuracy gain over Baseline
Initializer for next stage:
Figure 3: Histogram of accuracy improvements for models trained with MIDAS over baseline. The data points are MIDAS 1B models listed in Table 1. The figure shows that MIDAS-based models have much higher improvement in the contextual version of TyDiQA compared to the non-contextual version.
got10 = M L.kk)
(see Equation 1) end for return jk.
5

---

1. The depth d(i) L- grows linearly with the stage, similar to gradual stacking.
2. T: is typically either proportional to i (linear proportional) or i² (square proportional) or exp(i) (exponential). We will revisit this during our empirical analysis.
3. We use growth operator M in equation 1 for initializing the next stage, which corresponds to replicating the middle Lyk layers to initialize the next stage model.
3.3 Experiments: UL22 Pretraining
In this section, we evaluate MIDAS for standard language model pretraining. We train 24L decoder- only model with 1.5B parameters using the UL2 objective [Tay et al., 2022] on mixture of C4, Wikipedia, Arxiv and Github. The observations also hold for GPT-style autoregressive language modeling. To enable fair comparison, we cached the pretraining dataset and so all methods are trained for the same number 500B tokens in the same order, using the same batch size (refer to Appendix A. for more details on the training setup). We pretrain models with three methods: (a) standard training (Baseline), (b) gradual stacking (GRADSTACK) and (c) our proposed method MIDAS. The goal is to compare them with respect to validation loss and downstream performance on several diverse benchmarks. Motivated by the proportional schedules from prior work, we try the following generalized proportional schedules for gradual stacking and MIDAS.
Definition 3.1 (PROP-a schedule). For total training budget of T steps, the schedule PROP-Q spends time Ti in each stage such that T3 xx i for all stages E [k] Thus T:
= Σ -
PROP-1 schedule has been found work very well for BERT pretraining [Reddi al., 2023]. Since UL2 pretraining is harder task, we also explore less aggressive schedules like PROP-2 and PROP-3 that spend more time on larger models.
Efficiency and perplexity findings. We summarize the main results in Table 1, for various stacking methods and schedules. Firstly, we note that for all schedules, MIDAS has significantly better validation log perplexity than GRADSTACK the same speedup level. This suggests that stacking in the middle is a lot more effective for optimization than stacking at the end of the model. With the PROP-2 schedule, MIDAS is 24% faster and nearly matches the baseline's 's log perplexity. Additionally, we observe that the findings are robust to the choice of block size for stacking.
Downstream benchmark evaluations. While perplexity can serve as a decent proxy for model quality, there is growing evidence that it is not the best measure [Liang et al., 2023]. Downstream benchmark evaluations serve as a more holistic measure for quality and are out-of-distribution evaluations of skills. To this effect, we evaluate MIDAS many standard benchmarks and these are grouped into task categories in Table (refer to Appendix A.2 for more detailed evaluations on individual tasks). The accuracy for task category is an average over representative tasks from that group. For instance, for closed book QA task, we consider an average accuracy on TriviaQA, TydiQA (no context), NaturalQuestions and WebQuestions.
Surprisingly, we find that downstream improvements for MIDAS are significantly larger than the improvements in perplexity. In particular, MIDAS with PROP-2 schedule has very similar perplexity to bascline at 24% speedup, but the average downstream performance for MIDAS (26.8%) is much better than bascline (24.0%). In fact, even MIDAS with PROP-1 schedule which has worse log perplexity is much better downstream evaluations. Similar trends of better downstream evals holds for the 2B parameter model. The improvements are particularly large for open book QA and math word problems, both of which are tasks that require reasoning abilities whereas memorization tasks like closed book QA do not improve. We conjecture that these downstream improvements due to inductive bias induced by stacking and dive deeper into this in the next section.
4 Inductive bias of stacking
Results in Table demonstrate that MIDAS not only yields training speedups, but also improves downstream evaluations when trained on the same number of tokens as standard training. This suggests that stacking can extract more skills out of the same data. Here, we take a closer look at these improvements in downstream evaluations through the lens of an inductive bias of stacking.

---

Table 1: Downstream evaluations for UL2 pretrained models with 1B, 2B and 8B parameters. Comparisons include standard training (Baseline), gradual stacking (GRADSTACK) from [Reddi al., 2023] and our proposed method MIDAS. The downstream evaluations are averaged over tasks within 3 task groups. See Appendix A for precise tasks included in each task group. For each cateory and model size, we highlight the top model is bolded and the second best model underlined. Firstly, MIDAS is much better than GRADSTACK, thus justifying stacking in the middle. Secondly, MIDAS can match the log perplexity of baseline training while being roughly 24% faster. Furthermore, even the schedule with 40% speedup has much better downstream evaluations compared to baseline, even though it has worse log perplexity. The improvements are particularly large for task groups that require reasoning (open book QA, math word problems).
Loss (U)
Closed Book QA (1)
Open Book QA (1)
Math Word Problems (1)
All Tasks
d(i// (olak
Schedule
Speedup
Average (1)
(1
(a
Parameters
Baseline
24
1x
1.966
13.2
33.3
23.5
24.0
GRADSTACK
PROP-1
1.39%
2.045
10.3
31.4
23.5
22.6
MIDAS
PROP-1
1.39k
2.028
11.6
34,5
30.3
26.7
MIDAS
PROP-1
1.41x
2,032
10.6
36.1
27.0
25.6
GRADSTACK
PROP-2
1.24x
2.024
11.0
31.6
17.3
20.4
MIDAS
PROP-2
1,24x
2.009
11.7
36.3
29.0
MIDAS
PROP-2
1,26x
2.012
11.9
37.3
29.8
MIDAS
PROP-3
1,16x
1999
12.5
34.8
33,3
28.3
2B Parameters
Baseline
1x
1,966
15.2
39.1
27.1
28.0
MIDAS
PROP-1
1,39x
1.947
14.0
38.9
32.0
29.5
GRADSTACK
PROP-2
1.24%
1.945
14.2
37.0
24.5
25.9
MIDAS
PROP-2
1,24x
1.929
15.7
40.2
38.2
32.9
Parameters
Bascline
72
1x
18411
21.1
39.6
34.9
32.8
MIDAS
PROP-2
1 1.26x
1.844
21.8
40.0
43.1
36.4
4.1 Downstream performance vs log perplexity
reasonable expectation from pretraining is that improvements the pretraining objective would correlate with improvements in model quality and downstream performance. This notion of transfer has even been theoretically formalized for language modeling in Saunshi et al. [2020], Arora and Goyal [2023]. Thus, based on this, natural explanation for the downstream improvements of stacking would be that it generalizes better on the pretraining objective. However, as we see in Table 1, downstream performance of MIDAS is better despite having similar or worse validation perplexity - hence this is not simply the case of better generalization to unseen pretraining data. is natural ask: If not perplexity, what explains this downstream phenomenon?
Since pretraining objective is just a proxy objective for model quality, it is plausible that different training strategies and model architectures can extract different levels skills from it. This because there are multiple ways of doing well on the pretraining tasks, and some training strategies can be biased to pick one solution over another one. This behavior has been formalized as the inductive bias in pretraining by recent work [Saunshi et al., 2022, Liu et al., 2023] at the same level of validation pretraining loss, different optimization algorithms could have vastly different downstream performance. We hypothesize that a similar phenomenon is at play when it comes to stacking.
Isoplots. Inspired by this phenomenon of different downstream performance at the same perplexity, we visualize the inductive bias of method by plotting downstream accuracy vs log perplexity isoplots as training proceeds. We use the UL2 IB models that are pretrained with standard (baseline) training and with MIDAS using the PROP-2 schedule (refer to Section 3.3 for more details). Figure we visualize the downstream vs log perplexity plots for different task groups closed- QA, opon-book QA and math word problems. We observe a very interesting trend MIDAS and baseline training can have different isoplot behaviors and the divergence is different for different tasks.

---

Beseline
Baseline Midees
Baseline Mides
Mides
36
20
Cooed Book OA
12.0
Moth Word Problems
22
Reasoming Primitiver
22.3
Beseline Midas
23.0
22.3
24
13.0
32-
9.
31
2.2 perplesity
perplexity
perplexity
perplexity
Figure 4: Downstream evalulation vs validation log perplexity isoplots as training proceeds for baseline and MIDAS models trained on the same data (stacking is 24% faster here). On the y-axis we track the performance on various task groups closed book QA, open book QA, math word problems and our reasoning primitives from Section On the x-axis the log perplexity is presented in the reverse order, thus downstream performance for both methods improves log perplexity gets lower. For closed book QA (memorization) tasks MIDAS has very similar trends to baseline. For open book QA tasks and math word problems, MIDAS has much better downstream performance at an equivalent log perplexity. This showcases the inductive bias MIDAS towards better overall quality and better reasoning abilities.
4.2 Reasoning vs memorization for QA
For clearer display of the inductive bias, we measure the improvements due to MIDAS on closed book open book QA tasks. It is reasonable to assume that closed book QA tasks require strong memorization abilities whereas open book QA tasks require some reasoning abilities to infer answers from the context that is provided. On average, see much larger improvements on open book QA tasks compared to closed book QA tasks, as already evident Figure and Table 1.
MIDAS is significantly better on Open book QA. To make direct comparison, consider TydiQA-Goldd GoldP and TydiQA-NoCondee tasks the datasets are identical and the only difference is whether or not additional context is provided (the answer for the contextual version is guaranteed to be inferred from the given context). In Figure 3, we see that the improvements by various MIDAS based models on the contextual version of TydiQA are much higher than those on the non-contextual version. This provides direct evidence of the bias of MIDAS towards improving tasks that require reasoning. Furthermore, we find that the memorization performance of stacking improves as the schedule spends more time the larger model.
4.3 Reasoning in math tasks
To test reasoning abilities, we evaluate the language models on various math word problem datasets like SVAMP [Patel al., 2021], ASDiv [Miao al., 2020], AQuA dataset for algebraic word problems, the MAWPS benchmark Kedziol-K et al., 2016]. We report 5-shot evaluation for the pretrained model on these tasks. Following Wei et al. [2022], we use an external calculator to do the arithmetic and evaluate the models on their ability to compute the correct expression for the answer. This because small models have bad arithmetic accuracy. The choice of using caleulator or use MIDAS PROP-2 model
not does not significantly affect the trends of the results. For stacking, because achieves nearly the same perplexity as the baseline model (while being 24% faster), thus, leading to fair comparison based on the previous notion of inductive bias.
MIDAS is significantly better on Math/Reasoning tasks. Detailed results can be found in Table 5. For most math tasks, we observe that the MIDAS-bas based pretrained model is significantly better than the baseline model, especially for the MAWPs benchmark. This provides further evidence of better math and reasoning capabilities of MIDAS.
GSM8K fine-tuning. We also evaluate the 2B and 8B models on harder math problems from the GSM8k dataset [Cobbe al., 2021] through few-shot prompting and fine-tuning. Full results are presented in Table 2. For MIDAS the PROP-2 model that has very similar perplexity as the

---

Table 2: Evaluation on math tasks, including math word problems from Table and harder task GSM8k. For GSM8k we report accuracy with 8-shot prompts and with finetuning. We also report accuracy on all tasks after using an external calculator to fix arithmetic errors; this corresponds to w/ calc. Overall the use of calculator improves the accuracy for all models on all tasks. The benefit of MIDAS over baseline even higher with calculator.
Model
Pretraining
Math WPs (5-shot)
GSM8k (8-shot)
GSM8K (Finctune)
Loss (U)
cale.
cale.
W/o cale.
calc.
W/o cale.
cale.
2B Parameters
Baseline
1.926
15.4
27.1
3.0
MIDAS
1.929
22.5
38.3
3.0
4.1
10.4
14.5
8B Parameters
Baseline
1.841
27.3
34,9
4.5
6.6
MIDAS
1.844
32.9
43.1
5.5
7.4
15.2
18.7
Performance comparison reasoning primitives
100
Test accuracy (%%
Baseline
80-
69.5
40.7
45.0
24.3
21.6
14.5
Copying (5-shoti
Dopth (5 shot)
Death
Depth S-shht)
Depth (FT)
PSM shat)
(5 (5-shat)
Task
Figure 5: Accuracy improvements for model trained with MIDAS over baseline for representative reasoning primitives, despite having the same perplexity. We see clear improvements for MIDAS on almost all the primitives, both with 5-shot evaluation and after fine-tuning (FT) for the depth 1 and 2 primitive.
baseline model. We find that MIDAS has much higher accuracy after fine-tuning, thus suggesting that the benefits of the inductive bias continue after fine-tuning and are not just restricted few-shot evaluations. In particular, on the test set, the accuracy metric increased from 5.3% (for the baseline model) to 10.4% (for MIDAS) for the 2B model (these numbers were produced by computing the average score over three runs with different random seeds). Similarly the GSM8K accuracy of the 8B model improves from 12.3% to 15.2%. This suggests that MIDAS not only improves the performance on harder math tasks, but also that the gains remain improve after fine-tuning.
Effect of calculator. For LLMs with less than 20B parameters, Wei al. [2022] found that models often solve math problems correctly but make arithmetic errors, leading to low accuracy. Wei et al. [2022] remedied this by computing all arithmetic expressions using a Python program as an external calculator. In Table 2 we find that this improves the accuracy for our models too. Interestingly, the gap between MIDAS and baseline gets even larger with calculator use in almost all comparisons. We believe this is because arithmetic abilities are closer to memorization for smaller models [Razeghi et al., 2022] and calculator use makes the problem closer to reasoning, since now the model only has to infer the right expression. We believe this interplay between reasoning and memorization for math problems deserves further investigation.
4.4 Connection looped models
Given the nature of the growth operator each stage, we hypothesize that stacking based models are close to looped models. The layer duplication that happens at every stage ensures that blocks of layers start from a common initialization. We measure the similarity between different blocks of layers by measuring cosine similarities between the parameter vectors (see Figure 2). Since looped models have been conjectured to solve algorithmic problems [Giannou et al., 2023] by finding iterative solutions [Yang et al., 2023], we conjecture that the better reasoning abilities of MIDAS due to this connection to looped models We believe exploring this further is a very fruitful direction.

---

5
Deep dive into reasoning improvements
To further investigate the nature of this inductive bias, we construct various simple synthetic tasks to help tease apart the model's capabilities. We conjecture that these simple tasks capture core basic capabilities needed for contextual reasoning, and we therefore call these tasks "contextual reasoning primitives". They are: induction copying, variable assignment, and pre-school math (PSM), discussed further below. Overall, across various few-shot evaluations and fine-tuning, we see significant performance gaps between MIDAS and baseline training, suggesting that we have successfully isolated some of the basic capabilities at which MIDAS excels relative to baseline training. We refer the reader to Appendix B for more results and the exact input format.
Primitive I: Induction copying. The "induction copying" primitive presents a sequence of words, followed by subsequence selected randomly from within this original sequence, and asks the model to output the nex! word in the sequence. A simplified example is: "pum nyj gdq ocu rzk jbw mlz eny kyx uni rzk jbw mlz eny kyx", and the expected output is "uni". This primitive is inspired by the "induction head" mechanism introduced in Olsson et al. [2022], which is posited to be the basic mechanism for in-context learning more generally. In Figure task "Copying", we present results for 3-letter words of random letters, separated by spaces, with a sequence length of 10 and a subsequence length 5.
Primitive 2: Variable assignment. The "variable assignment" primitive tests the model's ability to associate a value with a variable name and apply this ability compositionally. which we test by varying the "ddpth" of the task. We conjecture that this ability is a core function in contextual reasoning, particularly in math. An example of the depth-O) variant is "u=1; t=0; v=13; y=4; f=22; y=", and the expected output is 4. An example of the 2epth-2 variant is "y=7; f=0; 2=3; b=9; x=8; q=y; 1=f; m=z; h=x; a=b; n=h; j=m; t=a; i=1; g=q; n=", and the expected output 8. Refer Appendix B for more details.
Primitive 3: Pre-school math (PSM). This tests the model's ability to solve a very simple "pre- school math" problem by correctly associating multiple values and variables simultaneously and applying this association to a particular task. An example is "z=6; b=5; i=-z+b; i=", and the expected answer (with chain-of of-tthughht is "-6+5=-1".
5-shot evaluation results. Figure presents the results for representative tasks, with more results in Appendix B. Overall, we see that MIDAS outperforms baseline training across all tasks. In particular, we see that MIDAS significantly stronger than the baseline at Depth O Copying. PSM-cale, and Depth 1, in decreasing order of magnitude of the performance gap. Depth-2 is much harder and is at random guessing (20%) for both models.
Fine-tuning results. Due to the difficulty of the variable assignment task Depths and 2, investigate fine-tuning on these tasks as well. We fine-tune on a mixture of 32 depth-1 examples and 32 d2pth-2 examples (i.e., only 64 examples total), using full-batch gradient descent. Figure 5 reports the validation accuracy on Depth I and Depth 2 after fine-tuning on this mixture (tasks "Depth 1 (FT)" and "Depth 2 (FT)"). Overall, we see that fine-tuning with just 64 examples significantly improves performance, resulting in MIDAS outperforming the baseline by a gap of over 20% validation accuracy at both depths. See Appendix for further fine-tuning and evaluation details.
Conclusions and future work
In this work we propose a novel stacking method that outperforms previous stacking methods and speeds up language model pretraining by 25-40%. In the process, we uncover a very intriguing inductive bias of stacking its ability improve downstream reasoning tasks. Through extensive empirical analysis, the paper makes a strong case for the presence and significance of this inductive bias. We believe this deserves further attention and exploration since understanding this inductive bias could unlock new approaches to improving model quality, reasoning in particular. The reasoning primitives start to provide more insights by isolating the reasoning improvements and we hope that the dataset is useful for future research on improving reasoning. Finally, understanding the dichotomy between memorization and reasoning, and how this affects the performance on various tasks, is an interesting direction to pursue.
10

---

Acknowledgments. We thank Srinadh Bhojanapalli and Vaishnavh Nagarajan for discussions on the role of layers and memory vs contextual tasks, respectively, in the early stages of the project. We also thank Satyen Kale for valuable feedback throughout the project.
References
Sanjeev Arora and Anirudh Goyal. theory for emergence of complex skills in language models. arXiv preprint arXiv:2307. 15936 2023.
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens. arXiv preprint arXiv:2303.cm 00003/08/211 12, 2023.
Cheng Chen, Yichun Yin, Lifeng Shang. Xin Jiang, Yujia Qin, Fengyu Wang, Zhi Wang, Xiao Chen, Zhiyuan Liu, and Qun Liu. bert2BERT: Towards reusable pretrained language models. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, 2022.
Tianqi Chen, Ian Goodfellow, and Jonathon Shlens. Net2net: Accelerating learning via knowledge transfer. International Conference on Learning Representations, 2016.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 14168, 2021.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2018.
Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, 2022.
Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers programmable computers. In International Conference Machine Learning, 2023.
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of bert by progressively stacking. In International conference on machine learning, pages 2337-2346. PMLR, 2019.
Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 10-15 Jul 2018
Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, et al. Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling. arXiv preprint arXiv:2312.. 25 15..16666 2023.
Rik Koncel- Kedzior-kiii Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. Mawps: A math word problem repository. In Proceedings the 2016 conference of the north american chapter of the association for computational linguistics: human language technologies, 2016.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference Learning Representations, 2020.
Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang, Xuying Meng, Siqi Fan, Peng Han, Jing Li, Du, Bowen Qin, et al. blm-101b: An open Il and how to train it with $100 k budget. arXiv preprint arXiv:2309.03852. 2023.
Percy Liang, Rishi Bommasani, Tony Lee, Dimitris Tsipras, Dilara Soylu, Michihiro Yasunaga, Yian Zhang, Deepak Narayanan, Yuhuai Wu, Ananya Kumar, et al. Holistic evaluation of language models. Transactions on Machine Learning Research, 2023.
11

---

Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre- training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning. PMLR, 2023.
Shen- Yun MiaO, Chao-Chun Liang, and Keh- Su. A diverse corpus for evaluating and developing english math word problem solvers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 975-984, 2020.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield- Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kermion, Liane Lovitt. Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, arXiv preprint arXiv:2209. / /895, 2022.
Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nlp models really able solve simple math word problems? In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies. Association for Computational Linguistics, 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 2020.
Yasaman Razeghi, Robert L Logan IV, Matt Gardner, and Sameer Singh. Impact of pretraining term frequencies on few-shot numerical reasoning. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Findings of the Association for Computational Linguisties: EMNLP 2022.
Association for Computational Linguistics, 2022.
Sashank Reddi, Sobhan Miryoosefi, Stefani Karp, Shankar Krishnan, Satyen Kale, Seungyeon Kim, and Sanjiv Kumar. Efficient training of language models using few-shot learning. In Proceedings of the 40th International Conference on Machine Learning, 2023.
Nikunj Saunshi, Sadhika Malladi, and Sanjeev Arora. A mathematical exploration of why language models help solve downstream tasks. International Conference on Learning Representations, 2020.
Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. In Proceedings of the 39th International Conference on Machine Learning, 2022.
Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596 4604. PMLR, 2018.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Dara Bahri, Tal Schuster, Steven Zheng, al. U12: Unifying language learning paradigms. In The Eleventh International Conference on Learning Representations, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09220 099889288.......88 2023.
Peihao Wang, Rameswar Panda, Lucas Torrob Hennigen, Philip Greengard, Leonid Karlinsky, Rogerio Feris, David Daniel Cox, Zhangyang Wang. and Yoon Kim. Learning to grow pretrained models for efficient transformer training. arXiv preprint 200300000 2023.
Yite Wang, Jiahao Su, Hanlin Lu, Cong Xie, Tianyi Liu, Jianbo Yuan, Haibin Lin, Ruoyu Sun, and Hongxia Yang. LEMON: Lossless model expansion. In The Twelfth International Conference on Learning Representations, 2024.
Zihao Wang and Lei Wu. Theoretical analysis of the inductive biases in deep convolutional networks. In Thirth- Conference Neural Information Processing Systems, 2023.
12

---

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, al. Ohain-of-thougee prompting elicits reasoning in large language models. Advances in neural information processing systems, 2022.
Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoulos, Looped transformers are better at learning learning algorithms. arXiv preprint arXiv:23/1.12424, 2023.
Yiqun Yao, Zheng Zhang, Jing Li, and Yequan Wang. Masked structural growth for 2x faster language model pre-training In The Twelfth International Conference on Learning Representations, 2024.
13

---

A Experimental Details
A.1 Pretraining details
Model architecture. We use decoder-only model and train it using the U1.2 objective [TTy et al., 2022] with 60% causal L.M, 20% prefix LMM and 20% span corruption. The 1B model uses 24 layers, model dimension of 2048, hidden dimension of 5120 and 32 attention heads. The 2B model is very similar the IB model, except uses 48 layers instead of 24. The 8B model 72 layers, model dimension of 2048, hidden dimension of 16384 and 16 attention heads.
Dataset. We use a mixture of C4 (57%) [Raffel et al., 2020], Wikipedia (17%), Github (17%), Arxiv (9%); the proportions are motivated by the dataset used for Llama pretraining [Touvron et al., 2023]. All models are trained for 512B tokens that are precached so that all model see exactly the same data in the same order. This corresponds 0.86 epochs of C4, 9 epochs of Wikipedia, 0.58 epochs of Arxiv and 0.44 epochs of Github.
Training details. For the IB and 2B models, we use a cosine learning schedule with a peak learning rate 0.01 that decays 0.001 in the end, and use batch size of 512. For the 8B model we use peak learning rate of 0.001 and decay it to 0.0001, and use a batch size of 1024. Peak learning rate was tuned be optimal for baseline training. All experiments use the AdaFactor optimizer [Shazeer and Stern, 2018] and sequence length of 1280.
A.2 Additional downstream evaluations
In this section we share further experimental details related to the results summarized in the Table 1.
Trivia QA
TyDi QA (w/o Context)
Natural Questions
Web Questions
d(i//
Schedule
Speedup
Parameters
Bascline
24
1x
28.1
12.0
4.5
8.1
GRADSTACK
PROP-1
1.39x
22.4
10.1
3.0
5.8
MIDAS
PROP-1 I
1.39x
25.0
11.7
3.7
5,9
MIDAS
PROP-1
1.41x
22.9
9.6
3.5
6.5
GRADSTACK
PROP-2
1,24x
22.9
4.0
5.9
MIDAS
PROP-2
1.24x
26.4
10.4
3.7
MIDAS
PROP-2
1.26x
25.5
10.9
3.8
7.4
MIDAS
PROP-3
1.16x
26.9
12.0
4.5
6.8
2B Parameters
Baseline
48
1x
33.6
12.8
5.9
8.7
MIDAS
PROP-]
1.39x
31.1
11.7
5.6
7.8
GRADSTACK
PROP-2
1.24x
32.0
12.5
5.8
6.7
MIDAS
PROP-2
1.24x
34.6
13.0
6.3
8.9
SB Parameters
Baseline
72
1x
47.0
15.2
9.6
12.9
MIDAS
PROP-2
1 1.26x
47.9
17.0
9.2
13.1
Table 3: Closed Book QA
14

---

TyDi QA (w/ Context)
SquadV2
DROP
QuAC
CoQA
Schedule
Specdup
1B Parameters
Baseline
31.4
41.1
22.9
18.8
$2.6
GRADSTACK
PROP-1
1.39X
34.3
36,9
MIDAS
PROP-1
1.39x
39.1
24.3
18.7
MIDAS
PROP-1
1.41x
37.0
44.9
25.0
18.4
55.1
GRADSTACK
PROP-2
1.24x
30.0
41.0
22.1
17.2
47.8
MIDAS
4
PROP-2
1.24x
35.5
46.6
24.4
19.7
55,4
MIDAS
PROP-2
1,26x
38.2
46.3
24.8
19.9
57.3
MIDAS
PROP-3
1.16x
33.6
40.2
24.7
19.5
35.9
2B Parameters
Baseline
48
1x
42.5
49.6
25.1
20.6
57.8
MIDAS
PROP-1
1.39x
37.7
48.9
26.1
20.1
61.8
GRADSTACK
PROP-2
1.24x
38.0
47.9
23.6
19.0
56.7
MIDAS
PROP-2
1.24x
41.8
48.0
27.9
20.7
8B Parameters
Baseline
51.8
25.9
19.6
MIDAS
PROP-2
1.26x
27.0
20.5
64.8
Table 4: Open Book QA
ASDiv
MAWPS Add/Sub
MAWPS Multi-Arith
MAWPS Single-Eq
MAWPS Single-Op
SVAMP
d(i)/i
Schedule
Speedup
IB Parameters
Baseline
24
1x
21.7
39.0
1.7
30.5
34.2
13,9
BRADSTACK
4
PROP-1
1.39x
19.1
38.8
2.0
31.1
35.2
13.1
MIDAS
PROP-1
1.39x
27.7
45.1
2.8
40.2
16.9
MIDAS
3
PROP-1
1.41x
25.8
45.1
2.5
33.1
40.7
14.8
GRADSTACK
PROP-2
1.24x
15.2
29.1
1.0
24.6
26.3
7.6
MIDAS
4
PROP-2
1.24x
26.3
51.9
3.3
39.4
40.0
13.0
MIDAS
3
PROP-2
1.26x
28.6
39.0
3.0
41.1
50.4
16.8
MIDAS
PROP-3
1.16x
28.9
55.7
1.5
41.1
$0.9
21.8
2B Parameters
Baseline
48
1x
27.9
41.5
3.2
37.4
36.5
16.4
MIDAS
8
PROP-1
1.39x
29.0
56.2
1.0
41.9
45.9
18.1
GRADSTACK
8
PROP-2
1.24x
22.7
43.0
3.2
30.5
33.1
14.3
MIDAS
PROP-2
1.24x
34.7
58.2
7.3
50.0
$7.5
21.8
8B Parameters
Baseline
72
35.0
44.6
3.7
46.0
22.9
MIDAS
PROP-2
1.26x
39.3
60.8
5.2
54.9
66.0
32.2
Table 5: Math World Problems
B Details for contextual reasoning primitives
In this section, we provide further details corresponding to Section 5.
All evaluations in Section were performed on the 1B-parameter models. For MIDAS, we use the variant with block size 4 and the PROP-2 schedule.
B.1 Exact input format
Expanding on Section 5, here we provide the format of the inputs and target outputs. The only caveat is that, for simplicity of presentation, present the inputs in 0-shot form here vs. their 5-shot form. In 5-shot form, which how we conduct the 5-shot evaluations, each example separated by two consecutive newline characters.
15

---

a linear map from Yo Yius-V
a linear map from Y, to Yous -Y
score
0.9
score
osne
0.-
osine
lnvertty
06-
05
04-
0.3
03
02
Figure 6: Measure of linearity for different layers in pretrained BERT-Bas and BERT-Large models. For each layer i, we fit linear map A, between inputs Y, and the output of the Transformer block (without the residual connection), Y1+1 Y.. We then measure the r2 score and cosine similarity for the learned linear fit. The first and last few layers demonstrate a much higher level linearity compared to the rest of the layers.
For each dataset below, the inputs are separated from the targets by the "I" character (this is not token the input), and the targets are colored in red.
Figure uses the following evaluation datasets, in the following order:
1. Copying (random-letter words)
2. Variable assignment depth (code)
3. Variable assignment depth (code)
4. Variable assignment depth (code)
5. Variable assignment depth (code)
6. Variable assignment depth (code)
7. Pre-school math (PSM)
Copying (random-letter words):
Fill in blank:
pum nyj gdq rzk jbw mlz eny kyx uni rzk jbw mlz eny kyx ->/uni
Copying (real words):
Fill blank:
eat fit ban sea vet zit pea cat van tea sea vet zit pea cat ~>/van
Variable assignment depth (basic):
16

---

Fill in blank:
o=14 s=4
u=8
m=10
q=12 == m=---. ->110
Variable assignment depth (basic):
Fill in blank:
g= b=24
v=3
s=23
h=20
k=b
a=s
n=y
f=g
d=h
a=_--.
->123
Variable assignment depth (basic):
Fill blank:
w=24
1=12
d=16
e=5
j=9
g-j
y=e
r=1
k=d
hew
v=g
i=r
c=h
t=k
p=y
_=_...
->124
Variable assignment depth (math):
17

---

The following is a set of simple mathematical equations.
n=22
r=16
w=13
v=6
k=10
What is the numerical value of n?
Answer: 122
Variable assignment depth (math):
The following is a set of simple mathematical equations.
h=20
w=9
c=22
j=11
v=5
g=c
k=w
a=j
s=h
o=v
What is the numerical value of s?
Answer: 120
Variable assignment depth (math):
The following a set of simple mathematical equations.
g
v=24
k=15
p=6
c=10
t=p
s=g
a=c
y=v
n=
1=g
j=t
m=y
i=a
What is the numerical value of j?
Answer: 16
Variable assignment depth (code):
18

---

a-6 The following is a very short Python program. Use the program to resolve the value of the variable in the question.
Program:
q=12
k=17
1=1
y=3
Question: What is the value of k?
Answer:
|17
Variable assignment depth (code):
The following is a very short Python program. Use the program to resolve the value of the variable in the question.
Program:
k=11
f=21
e=10
1=7
c=13
y=f
0Fc
r=e
u=k
n=1
Question:
What is the value of o?
Answer:
|13
Variable assignment depth (code):
The following is a very short Python program. Use the program to resolve the value of the variable in the question.
Program:
t=13
j=14
y=4
s=17
y=21
qFj
1=s
19

---

e
h
x=v
b=x
f=e
n
a=h
i=l
Question:
What is the value of i?
Answer:
|17
Pre-school math (PSM):
Fill in blank:
k=1 j=8 1=-k+j 1=_-.. ~>1-1+8=7
Arithmetic:
-3+2=-1
-6+1=-5
+9-7=2
-6-4=-10
-6-1=-7
+1+9=110
B.2 Fine-tuning details
For fine-tuning, we use the "code" variant of the variable assignment task, depths and 2, in O-shot form (i.e., in-context examples). Due to the randomness of the data generation process and the rather small size of each dataset (64 examples), we randomly generate 3 different 64-example fine-tuning datasets (consisting of 32 depth-1 examples and 32 depth-2 examples), fine and report our results as an average across the runs. Table 7 reports the standard deviations as well.
each,
Regarding hyperparameters, we continue to use AdaFactor [Shazeer and Stern, 2018] with the same hyperparameters in the pretraining phase, with the exception of learning rate and batch size. We use
20

---

constant learning rate of 0.001, which was chosen to match the final learning rate of the pretraining phase. We use full-batch training with our 64-example datasets. We then evaluate performance separately on depth 1 and depth 2.
For every step i E {200,.... 300), chosen to be significantly after training has converged to 100% accuracy (we do not observe overfitting in this range training continues), evaluate performance on a 1000-example holdout set. For smoothing purposes, we average over steps 200 through 300 and report the final averaged performance.
B.3 Full 5-shot and fine-tuning results
5-shot. Table includes 5-shot evaluation results contextual reasoning primitives. Rows 10, 11, and 14 are the rows which appear in Figure 5.
When performance is better than random guessing, MIDAS consistently outperforms the baseline in rows 1-11.
For pre-school math (rows 12-14), the value report in Figure 5 is "with calculator". This is because the pre-school math task actually combines two capabilities: reasoning and arithmetic. Arithmetic can be thought of as memorization task. We evaluate arithmetic for MIDAS and baseline training, and we see that arithmetic quite poor for both models (7.8% and 9.6%, respectively, in Table 6). However, by evaluating PSM with chain-of-thought and only assessing the accuracy of the reasoning chain itself, i..., "-6+5" vs. "-1", we can successfully disentangle reasoning and memorization in our call it "PSM with calculator"
evaluation. This is equivalent to having access calculator, "PSM-cale" in Figure 5.
Task Copying (random-letter words)
MIDAS (%)
Baseline (%)
Random guessing( %)
24.3
14.9
10
Copying (real words)
17.8
10.3
10
Variable assignment depth 0 (basic)
35.6
32.1
20
Variable assignment depth (basic)
20.6
21.9
20
Variable assignment depth (basic)
18.9
17.7
20
Variable assignment depth 0 (math)
92.8
50.1
20
Variable assignment depth (math)
26.5
19.2
20
Variable assignment depth 2 (math)
20.4
18.8
20
Variable assignment depth 0 (code)
86.0
49.7
20
Variable assignment depth 1 (code)
28.3
21.6
20
Variable assignment depth 2 (code)
19.5
19
20
Pre-school math (PSM), no calculator
7.8
9.6
n/a
Arithmetic-one accuracy
9.7
10.3
n/a
Pre-school math (PSM), with calculator
69.5
62
n/a
Table 6: 5-shot results for all variants of the contextual reasoning primitives. This is an expanded set
compared to Figure
Fine tuning. Table presents the fine-tuning results from Figure 5 along with corresponding standard deviations (across the trials).
Task
MIDAS (%)
Baseline (%)
Random guessing( %)
Variable assignment depth 1 (code)
68.54 + 7.69
43.75 + 5.54
20
Variable assignment depth 2 (code)
44.97 + 7.26
23.88 + 1.56
20
Table 7: Fine-tuning results corresponding to Figure 5's 2 fine-tuning tasks. Additionally, this table reports the standard deviation across the runs with std dev.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The Abstract and Introduction provide good summary of the paper's contri- butions. They express measured excitement about the results without overpromising.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: It has been discussed throughout the paper. In particular, MIDAS provides limited improvement on memorization- based tasks which has been discussed in Section 3.3 and Section 4.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
22

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification:
No theoretical result is provided in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Section 3 and Appendix A are sufficient for reproducing the main results discussed in this paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
23

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: We do not provide access to the data and code, but the data and models come from prior works; where differences between work and prior work appear, we highlight them (Section and Section A).
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips) cc/
conoonoyollylloollo for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to rum to reproduce the results. See the NeurIPS code and data submission guidelines (https:
// //nips.coo ccpps.ccppppp wttccc//pubu. wcc./pulcc/eaeeccccccecceeeeeele wwt.cc/gggieeeeeeee htt////odbbbissseeee wwabbmisssooooo.l...oo..o.oooooe missPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experimental details have been provided in Appendix A.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due to the computational demands of exploring reasonably large set of experiments, it was not feasible to perform multiple runs per setting.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
24

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative emror rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification:
We do not report these details in the submission but can include them final version.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / http//nnurrrp.ppe wwt.rriss.ccppuuuuee cccppbliic/thieeeeeeeeeeeeeeeeee htte//hhiissiieieeeeeeeeeeeeeeee csGuidelines?
Answer: [Yes]
Justification: We have reviewed the NeurlPS Code of Ethics. To the best of our knowledge, this work conforms.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
This a foundational/analyyical paper not tied to any specific application or deployment.
Guidelines:
25

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No data or models have been released with this paper.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [No]
Justification: The creators of data and models cited, but we did not include licenses for all assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
26

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
27

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "inlbL2VBiN", "reviewer": "Reviewer_UVBV", "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.", "questions": "See the weakness section.", "limitations": "The paper discusses some of the limitations in reasonable detail.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The reasoning behind stacking in the middle rather than at the end is simple and seems effective in experiments. \n- 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. \n- 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.  \n\n- There are some typos in the paper:\n   - This on line 103? \n   - Line 105?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback and insightful questions.\n\n**Q1**: *the improvement of MIDAS on downstream tasks (open book QA and math word problems) is marginal compared to the improvements on these tasks [reasoning primitives].*\n\nA: While the improvement on these benchmarks is lower compared to improvement on primitives, the absolute magnitude of improvement is still quite high. For instance, for the 1B model, MIDAS with Prop-2 schedule improves open book QA from 33.3 -> 36.3 (+9% relative) and math word problems from 23.5 -> 29.0 (+23% relative). The reviewer is correct in observing that improvements on reasoning primitives from 35.2 -> 58.4 (+65% relative) are even larger. Our hypothesis is the following: solving each benchmark dataset requires some combination of different skills like memorization, reasoning, etc. with varying proportions. Based on the analysis from Section 4.2, we observed that MIDAS helps a lot more on tasks that require more reasoning (open book QA) than the memorization version of the same QA tasks. Reasoning primitives were designed to specifically isolate reasoning skills and minimize the need for memorization. Thus, we believe that MIDAS improves the most on primitives. Tasks like open book QA do require some reasoning from context, however, they can also benefit a lot from memorization of facts and world knowledge where MIDAS does not improve over baseline. (For instance the model may choose to ignore the context and answer from memory, since after all it has decent closed book QA accuracy). This, we believe, can dampen the magnitude of improvements. Math word problems intuitively require more reasoning than memorization compared to open book QA, and thus improvements are larger. We believe that understanding this interplay between reasoning and memorization is a very interesting and important future direction.\n\n\n**Q2**: *... raises the question if there are synthetic tasks for which MIDAS is not performing as well as the baseline*\n\nA: That is a good question. Since “memorization” is the slice where MIDAS has the least improvement (slight regression in some cases), we would need to construct a synthetic dataset that tests for memorization abilities. Even on closed book QA benchmarks, MIDAS is almost neutral in various cases, so finding a task where MIDAS is substantially worse would be an interesting direction. We could not come up with a very good synthetic task to isolate this effect, and if the reviewer has some suggestions, we are happy to try those.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZyrfdLjhKf", "reviewer": "Reviewer_1oxX", "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.", "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.", "limitations": "Yes, the authors adequately addressed the limitations", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. This paper proposed a novel stacking algorithm for improving reasoning. \n2. Experiments show that the algorithm improves performance on four distinct types of reasoning benchmarks. \n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback and suggestions.\n\n**Q1**: *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 artifacts.*\n\nA: Based on the reviewer’s suggestion, we evaluated MIDAS and baseline models on another category of tasks called “story completion”, which includes tasks like LAMBADA, StoryCloze and Hellaswag. These measure the model's ability to correctly complete a story or premise. These tasks are also closer to reasoning, since the completion needs to be figured out from the given premise. The results for the 1B UL2 MIDAS models are presented below. Please refer to the response to Reviewer s9u7 for more evaluations on the story completion category.\n\n| Model | Lambada | StoryCloze | HellaSwag |\n|--|--|--|--|\n| 1B Baseline | 16.1 | 73.1 | 58.7 |\n| 1B MIDAS-Prop2 (1.33x) | 18.4 | 74.6 | 58.9 |\n| 1B MIDAS-Prop3 (1.25x) | 25.5 | 74.9 | 60.4 |\n\n**Q2**: *Does it make sense to test on existing synthetic reasoning benchmarks instead of just synthetic reasoning primitives? For example, ProofWriter*\n\nA: Thank you for the suggestion. One of the motivations behind new reasoning primitives was to evaluate these models on very simple and basic tasks that can isolate the benefits of MIDAS, which we could not do successfully with existing benchmarks. That said, it also makes sense to test on more complex synthetic tasks like ProofWriter. Unfortunately, we were not able to set up this evaluation task in time for the response since we prioritized some other experiments, but we will try to include this in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EDEW6adkWd", "reviewer": "Reviewer_48NW", "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.", "questions": "Will the authors release the code and dataset to allow the community to further study the interesting phenomenon mentioned in section 5?\n\nCan this phenomenon be verified with fewer resources?", "limitations": "The authors need to add a separate section for limitations.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. The authors propose MIDAS, a novel variant of gradual stacking, which achieves better training efficiency that baselines.\n2. 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?\n\nMinor:\nIn line 105, “For simplicity, we L is divisible by k” should be “For simplicity, L is divisible by k.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and suggestions. We tried to address them below by running new experiments.\n\n**Q1**: *Can the authors evaluate on more reasoning and common-sense questions for better comparison, such as tasks tested in the Llama paper?*\n\nA: Based on the reviewer’s suggestion, we run evaluations on common sense reasoning benchmarks PiQA, HellaSwag, Winogrande, ARC-E, ARC-C from the Llama paper. \n\n| Model | PiQA | HellaSwag | Winogrande | ARC-E | ARC-C | Average (5) |\n|--|--|--|--|--|--|--|\n| 1B Baseline | 75.4 |\t58.7 |\t59.8 |\t63.2 | \t31.7 | 57.8 |\n| 1B MIDAS Prop-2 (1.33x) | 74.0 | 58.9 | 58.9 | 63.1 | 31.4 | 57.3 | \n| 1B MIDAS Prop-3 (1.25x) | 74.0 | 60.4 | 59.0 | 62.5 | 32.7 | 57.7 |\n|--|--|--|--|--|--|--|\n| 2B Baseline | 75.6 | 65.0 | 62.3 | 66.9 | 35.2 | 61.0 |\n| 2B MIDAS Prop-2 (1.24x) | 76.0 | 65.5 | 62.8 | 67.0 | 34.5 | 61.1 |\n\nWe did not include BoolQ because the evals are very noisy and close to “trivial” accuracy. Furthermore the trends are highly sensitive to which metric is used (accuracy vs auc-pr).\n\nOverall, we find the MIDAS is roughly neutral related to baseline at ~25% speedup. We do not observe a strong inductive bias for these tasks. Our hypothesis is that common-sense questions require a significant component of “memorization” of world knowledge of what is and what is not plausible in the physical world. \n\n---\n\n**Q2**: *Will the authors release the code and dataset to allow the community to further study the interesting phenomenon mentioned in section 5?*\n\nA: We will release the reasoning primitives dataset and code to aid future research on this topic.\n\n---\n\n**Q3**: *Can this phenomenon be verified with fewer resources?*\n\nA: The observation that MIDAS (middle stacking) is better than gradual stacking also holds at the 120M parameter scale, with decoder-only models like GPT-2 small and also with masked LM with BERT-Base. We did not report these results in the paper since the focus was on larger scale models, but we can include these results in the revision. The inductive bias towards reasoning and non-trivial performance on reasoning primitives may require closer to a billion parameters to see non-trivial few-shot performance.\n\nWe will fix the typos pointed out by the reviewer and will include more discussion about limitations, particularly regarding the understanding of the inductive bias.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZvGNUwLc4R", "reviewer": "Reviewer_s9u7", "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.", "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.\n2) From Table 1 we see memorization performance of MIDAS drops compared to baseline. Does the authors have any intuition why this is happening? \n3) What is the performance of GRADSTACK for 2B parameters model in Table 1?", "limitations": "This work shows limited improvement on memorization-based tasks.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n2) 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.\n3) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and many suggestions. We tried to address them below by running new experiments.\n\n**Q1**: *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.*\n\nA: Based on the suggestion, we pretrained an 8B model on 100B tokens (instead of 400B tokens in the interest of time) for baseline and MIDAS. We try three learning rates to observe the effect of hyperparameters and report the results below. We include the evaluations corresponding to the categories in Table 1. We also include a new category for story completion tasks (average of LAMBADA, Story Cloze, Hellaswag) that also require reasoning from context, inspired by other reviewer suggestions. The average column now refers to an average of 18 tasks instead of 15. Overall the trends are very similar to 1B and 2B models: at 1.24x speedup MIDAS has much better task average. Interestingly, MIDAS seems a lot more robust to the choice of learning rate. We will include these 8B experiments and a longer one in the revision.\n\n| Model | Loss (val) | Closed book QA | Open book QA | Math Word Problems | Story Completion | Task Average (18) |\n|--|--|--|--|--|--|--|\n| Baseline (LR=1e2) | 1.962 | 13.7 | 35.1 | 27.7 | 49.7 | 30.3 |\n| MIDAS (1.26x) (LR=1e2) | 1.917 | 16.1 | 39.1 | 36.0 | 54.2 | 35.5 |\n|--|--|--|--|--|--|--|--|\n| Baseline (LR=3e3) | 1.911 | 16.5 | 37.2 | 30.8 | 53.2 | 33.1 |\n| MIDAS (1.26x) (LR=3e3) | 1.905 | 18.2 | 39.2 | 30.9 | 54.7 | 34.3 |\n|--|--|--|--|--|--|--|--|\n| Baseline (LR=1e3) | 1.898 | 17.9 | 38.6 | 27.6 | 51.8 | 32.5 |\n| MIDAS (1.26x) (LR=1e3) | 1.909 | 17.3 | 36.3 | 33.0 | 55.8 | 34.2 |\n\n\n---\n\n**Q2**: *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?*\n\nA: The observations that MIDAS is better than gradual stacking, and the inductive bias of MIDAS towards improved reasoning, also hold for GPT-style causal language modeling. We started with the UL2 objective and stuck with it, partly because it also uses causal LM for 60% of the data. Based on the reviewer's suggestion, we ran causal LM training for the 1B model and reported the results below. We include the evaluations corresponding to the categories in Table 1 for Causal LM training of the 1B model and also the new categories described in the response to **Q1**. The trends are very similar to the UL2 models.\n\n| Model | Loss (val) | Closed book QA | Open book QA | Math Word Problems | Story Completion | Task Average (18) | Primitives |\n|--|--|--|--|--|--|--|--|\n| 1B LM Baseline | 2.43 | 12.3 | 33.9 | 30.3 | 42.8 | 30.4 | 51.6 |\n| 1B LM GradStack (1.33x) | 2.43 | 10.3 | 32.1 | 22.3 | 42.3 | 26.5 | 39.9 |\n| 1B LM MIDAS (1.33x) | 2.44 | 11.4 | 36.6 | 29.1 | 50.2 | 31.3 | 58.7 |\n\n\n---\n\n**Q3**: *memorization performance of MIDAS drops compared to baseline. Does the authors have any intuition why this is happening?*\n\nA: This is an interesting open question and we attempted an initial analysis in Section 4.2, where we notice a clear inductive bias of MIDAS towards improving open book QA problems (which is closer to reasoning) more than closed book versions of the same questions (which is closer to memorization). Our speculative hypothesis is that due to the connection to looped models, MIDAS has a slightly lower “effective” number of parameters, which is known to correlate with memorization abilities [1]. This seems to be more than compensated for by the improved reasoning abilities of MIDAS. We believe this phenomenon deserves further exploration and understanding.\n\n\n---\n\n**Q4**: *What is the performance of GRADSTACK for 2B parameters model in Table 1?*\n\nA: Below we report results for GradStack with Prop-2 schedule for the 2B model, including the new task categories.\n\n| Model | Loss (val) | Closed book QA | Open book QA | Math Word Problems | Story Completion | Task Average (18) | Primitives |\n|--|--|--|--|--|--|--|--|\n| 2B Baseline | 1.926 | 15.2 | 39.1 | 27.1 | 54.4 | 32.4 | 54.4 |\n| 2B GradStack (1.24x) | 1.945 | 14.2 | 37.0 | 24.5 | 51.5 | 30.2 | 64.2 |\n| 2B MIDAS (1.24x) | 1.929 | 15.7 | 40.2 | 38.3 | 53.6 | 36.3 | 78.3 |\n\n\n\n\n[1] Allen-Zhu, Li. Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws. 2024", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for constructive feedback and their appreciation for the results in the paper. We have responded to reviewer questions independently. Following reviewer suggestions, we ran and reported the following new experiments:

- Run GradStack for the 2B UL2 model that was missing in Table 1. MIDAS continues to be much better than GradStack (see Q4 of reviewer s9u7)
- Pretraining with causal LM objective on 1B model. Trends are very similar to UL2 pretraining objective (see Q2 of reviewer s9u7)
- 8B pretraining with 100B tokens to verify that the results scale with model size. MIDAS @ 1.26x speedup has better downstream evals than baseline training, and is more robust to choice of learning rate. (see Q1 of reviewer s9u7)
- Add new evaluations on another category of tasks, “story completion”, that includes Lambada, StoryCloze and HellaSwag. The inductive bias of MIDAS towards reasoning shows up here too. (see Q1 of reviewer 1oxX)
- More evaluations on common sense tasks from Llama paper (PiQA, HellaSwag, WinoGrande, ARC-E, ARC-C). MIDAS @ 1.25x speedup is neutral on average compared to baseline training. (see Q1 of reviewer 48NW)


We hope this addresses any remaining concerns by the reviewers, and we are happy to engage in more discussions.

Meta-review: A paper that contributes to the field. 

The reviewers and myself were positive about the paper, recognizing its novel approach and its potential impact on improving training efficiency and reasoning capabilities in language models. The proposed MiDAS technique is seen as a significant improvement over previous stacking methods, particularly in enhancing performance on reasoning-related tasks, which is a critical area of research. The empirical results are solid, and the connection made between stacking and looped models is particularly insightful, paving the way for further research in this area. The paper is also well-written and easy to follow, which was noted by multiple reviewers.

Overall, strenghts:
Novelty and Effectiveness: The idea of stacking middle layers instead of the last few layers is both novel and effective, as demonstrated by the experiments. Empirical Validation: Extensive experiments validate the proposed method, showing significant improvements in reasoning tasks. Well-Written: The paper is clearly presented, making the complex concepts easier to understand.

The following are suggestions for minor expansions and clarifications:
The paper could benefit from additional evaluations on more diverse reasoning benchmarks to further substantiate the claim that MIDAS enhances reasoning over memorization. Consider adding synthetic reasoning benchmarks like ProofWriter to provide a broader perspective. Moreover, the drop in performance on memorization tasks, compared to baseline models, was noted. A more detailed analysis and possible mitigation strategies could strengthen the paper. Understanding and explaining this trade-off more deeply would be valuable. 

In summary, this paper presents a well-supported and innovative approach to improving reasoning in language models and it will be a good contribution to the field.
Decision: Accept (poster)