Title: On the Inductive Bias of Stacking Towards Improving Reasoning
Venue: NeurIPS 2024
Accepted: yes
Keywords: stacking

language model

reasoning

inductive bias

efficient training
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

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 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 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.

1 Introduction

With the advent of very large deep learning models, efficient training to 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 to vanilla training [Chen et 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

Shankar Krishnan

Improvementfromstackingatsamelogperplexity

25

20

15

10

5

0

closedbook

openbook

math word

reasoning

QA

QA

problems

primitives

Task groups

3

3

3

2

2

2

1

1

Gradual

stacking

(our method)

are stacked onto itself to 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 on generalization remains a fundamental open question and is largely unexplored.

Modern deep learning models, when trained carefully, have been shown to 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 in various forms including simplicity bias, flatness of learned function, and sparsity. The implicit bias of optimizers, in particular, has been subject to 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 et 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, Dehghani 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 (e.g., greedy layerwise training or gradual stacking) that are gaining popularity.

In this work, we investigate the inductive bias of stacking-based approaches beyond training efficiency. We uncover an intriguing phenomenon pretraining with a variant of stacking is not only efficient, but also has a desirable inductive bias towards improving downstream benchmarks . First, through comprehensive empirical analysis, we discover a novel variant of gradual stacking called MIDAS (MIDdle grAdual Stacking) which copies the middle block of layers of a small network to initialize a larger network (see Figure 1). We demonstrate that MIDAS is 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, we see in Figure 1 that MIDAS has significantly better performance on 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 to 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 at 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 in Liu et al. [2023] for language modeling on synthetic data. However, this is the first evidence of a strong inductive bias for a 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 a model pretrained with MIDAS has much better performance on the reasoning primitives than a model obtained through standard pretraining, as is evident in Figure 1. In light of the above discussion, we state the main contributions of our paper.

We propose a 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 a surprising benefit: beyond enabling efficient training, it also enhances performance on downstream tasks . This improvement is especially notable in tasks that rely on 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 f i and #( f ) to denote the i th layer and the number of layers of the network, respectively. With slight abuse of notation, we use f i,b (where i, b ∈ Z + ) to denote the layers between ( i -1) · b to i · b of a deep network f . In other words, f i,b denotes the i th block of b layers in a deep network f . a 1: k is used to denote a sequence of k scalars { a 1 , . . . , a k } .

Our goal is to learn a function f : X → Y which minimizes the loss E ( x,y ) ∼D glyph[lscript] ( f ( x ) , y ) , for some loss function glyph[lscript] : Y×Y → R + ∪{ 0 } and data distribution D on X ×Y . We are interested in functions of the form f = f L ◦ f L -1 ◦ · · · ◦ f 1 where ◦ and L represent function composition and depth of the network, respectively. We use F L 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 a 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:

where B t is a mini-batch from distribution D and A ( f t -1 , B t , η t ) represents the iterative optimizer update at f t -1 on B t and learning rate η t . 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, it is useful to formally define a stagewise training procedure. In contrast to standard training, k -stage training involves dividing the training process into k 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 L is divisible by k . The following are the key ingredients:

Function class across stages. At stage i , we use function class F d ( i ) where d ( i ) denotes the depth of the network at that stage. When d ( i ) glyph[lessmuch] L , training is more efficient.

Training schedules across stages. As training is divided into k stages, we use T 1 , · · · , T k steps across stages such that ∑ k i =1 T i = T .

0.5

(a) ALBert layer similarity

(c)

MIDAS

block similarity

1.0

1.0

2

m

0.4

1.00

0.29

0.29

1.00

0.22

0.46

0.19

0.36

0.18

0.33

0.17

0.8

1.00

0.28

0.28

1.00

0.26

0.69

0.26

0.68

0.25

0.65

0.21

0.39

0.8

寸

0.3

N

0.22

0.46

1.00

0.61

0.55

0.55

0.6

2N

0.26

0.69

1.00

0.98

0.87

0.41

0.6

6

7

0.2

m

0.19

0.36

0.61

1.00

0.82

0.82

0.4

m

0.26

0.68

0.98

1.00

0.88

0.41

0.4

6

- 0.1

4

0.18

0.33

0.55

0.82

1.00

0.98

- 0.2

4

0.25

0.65

0.87

0.88

1.00

0.42

- 0.2

1

0

1

2

4

6

7

8

6

10

11

-0.0

0.17

0

0.33

1

0.55

2

0.82

E

0.98

4

1.00

5

0.0

0.21

0

0.39

1

0.41

2

0.41

3

0.42

4

1.00

5

0.0

Stage initialization. This is the key component of stagewise training. Given a network f ∈ F d ( i -1) trained in the ( i -1) th stage, let M i ( f ) denote the network initialization for the next stage where M i : F d ( i -1) →F d ( i ) is a 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 of the aforementioned framework. We provide a 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 -stage training setup where model in the previous stage is stacked onto itself to initialize the model in the next stage. In particular, (1) depth d ( i ) = 2 i -1 d (1) grows exponentially, (2) schedule T i is typically T/k or proportional to d ( i ) , and (3) the growth function M i ( f ) = f ◦ f .

Gradual Stacking [Reddi et al., 2023]. In contrast to progressive stacking, gradual stacking linearly increases the model depth by k 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, as follows.

The depth d ( i ) = L · i k grows linearly with the stage.

T i is typically either T/k or allocated proportional or exponential to depth.

M i ( f d ( i -1) ◦ · · · ◦ f 1 ) = f d ( i -1) · · · ◦ f d ( i -1) -( L/k )+1 ◦ f d ( i -1) · · · f 1 . This corresponds to stacking the last L/k layers onto the network to initialize the next stage model.

In the next section, we study a 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 in 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 a deep network and (b) a connection to looped models. Before delving into more technical details, it is important to illustrate these points. We present the case for MIDAS based on 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 b 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 a 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 a connection to looped models - networks with shared parameters between layers.

Observation 2: gradual stacking leads to models resembling looped models. To check the effect of the initialization, we measure the cosine similarity between weights of layers for a 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 [Lan 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 looped models?

Observation 3: looped models exhibit similarity in the middle. In order to study this, we train a 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 is the highest in the middle (see Figure 2a).

The above observations provides a 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 in MIDAS .

where n = #( f ) /b is the number of blocks of b layers in deep network f . Note that operator M ( f, 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.

HistogramofgainoverBaselinefordifferentMidasmodels

0.8

TyDiQAw/ocontext

0.7

TyDiQAwcontext

TyDiQAw/ocontext-mean

0.6

TyDiQAwcontext-mean

Density

0.5

0.4

0.3

0.2

0.1

0.0

-5-4-3-2-101234567

8

AccuracygainoverBaseline

The depth d ( i ) = L · i k grows linearly with the stage, similar to gradual stacking.

T i is typically either proportional to i (linear proportional) or i 2 (square proportional) or exp( i ) (exponential). We will revisit this during our empirical analysis.

We use growth operator M in equation 1 for initializing the next stage, which corresponds to replicating the middle L/k layers to initialize the next stage model.

3.3 Experiments: UL2 Pretraining

In this section, we evaluate MIDAS for standard language model pretraining. We train a 24L decoderonly model with 1.5B parameters using the UL2 objective [Tay et al., 2022] on a 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.1 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α schedule) . For a total training budget of T steps, the schedule PROP -α spends time T i in each stage such that T i ∝ i α for all stages i ∈ [ k ] . Thus T i = i α ∑ k j =1 j α T

PROP-1 schedule has been found to work very well for BERT pretraining [Reddi et al., 2023]. Since UL2 pretraining is a 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 at 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 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 on many standard benchmarks and these are grouped into task categories in Table 1 (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 baseline at 24% speedup, but the average downstream performance for MIDAS (26.8%) is much better than baseline (24.0%). In fact, even MIDAS with PROP-1 schedule which has worse log perplexity is much better on 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 are due to an inductive bias induced by stacking and we dive deeper into this in the next section.

4 Inductive bias of stacking

Results in Table 1 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.

4.1 Downstream performance vs log perplexity

A reasonable expectation from pretraining is that improvements in 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, a 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. It is natural to 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 of skills from it. This is 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 a method by plotting downstream accuracy vs log perplexity isoplots as training proceeds. We use the UL2 1B models that are pretrained with standard (baseline) training and with MIDAS using the PROP-2 schedule (refer to Section 3.3 for more details). In Figure 4, we visualize the downstream vs log perplexity plots for different task groups - closed-book QA, open-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.

13.0

x

Baseline

36

Baseline

30

x

Baseline

X

60

12.5

35

Problems

28

X

Primitives

55

QA

12.0

Book

11.5

X

Book

50

X

lose

11.0

Open

34

x

WordF

26

Reasoning

45

Baseline

10.5

33

Math

24

10.0

40

32

22

9.5

35

9.0

31

2.06

2.04

2.02

2.00

2.06

2.04

2.02

2.00

2.06

2.04

2.02

2.00

2.06

2.04

2.02

2.00

Log perplexity

Log perplexity

Log perplexity

Log perplexity

4.2 Reasoning vs memorization for QA

For a clearer display of the inductive bias, we measure the improvements due to MIDAS on closed book vs 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, we see much larger improvements on open book QA tasks compared to closed book QA tasks, as already evident in Figure 1 and Table 1.

MIDAS is significantly better on Open book QA. To make a direct comparison, we consider TydiQA-GoldP and TydiQA-NoContext 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 on 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 et al., 2021], ASDiv [Miao et al., 2020], AQuA dataset for algebraic word problems, the MAWPS benchmark [Koncel-Kedziorski 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 is because small models have bad arithmetic accuracy. The choice of using calculator or not does not significantly affect the trends of the results. For stacking, we use MIDAS PROP-2 model because it achieves nearly the same perplexity as the baseline model (while being 24% faster), thus, leading to a 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 -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 et al., 2021] through few-shot prompting and fine-tuning. Full results are presented in Table 2. For MIDAS we use the PROP-2 model that has very similar perplexity as the

Copying

(5-shot)

Depth 0

(5-shot)

Depth 1

(5-shot)

Depth 1

(FT)

Depth 2

(5-shot)

Depth 2

(FT)

PSM-calc

(5-shot)

Task

0

20

40

60

80

100

Test accuracy (%)

24.3

86.0

28.3

68.5

19.5

45.0

69.5

14.9

49.7

21.6

43.8

19.0

23.9

62.0

Performance comparison on reasoning primitives

MIDAS

Baseline

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 to 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 or improve after fine-tuning.

Effect of calculator. For LLMs with less than 20B parameters, Wei et 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 to looped models

Given the nature of the growth operator in 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 are 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 1: Induction copying. The 'induction copying' primitive presents a sequence of words, followed by a subsequence selected randomly from within this original sequence, and asks the model to output the next 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 5, 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 of 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 'depth' of the task. We conjecture that this ability is a core function in contextual reasoning, particularly in math. An example of the depth-0 variant is ' u=1; t=0; v=13; y=4; f=22; y= ', and the expected output is 4. An example of the depth-2 variant is ' y=7; f=0; z=3; b=9; x=8; q=y; l=f; m=z; h=x; a=b; n=h; j=m; t=a; i=l; g=q; n= ', and the expected output is 8. Refer to Appendix B for more details.

Primitive 3: Pre-school math (PSM). This tests the model's ability to solve a very simple 'preschool 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-thought) is ' -6+5=-1 '.

5-shot evaluation results. Figure 5 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 is significantly stronger than the baseline at Depth 0, Copying, PSM-calc, 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 at Depths 1 and 2, we investigate fine-tuning on these tasks as well. We fine-tune on a mixture of 32 depth-1 examples and 32 depth-2 examples (i.e., only 64 examples total), using full-batch gradient descent. Figure 5 reports the validation accuracy on Depth 1 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 B for further fine-tuning and evaluation details.

6 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 to 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.

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. A 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.08112 , 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:2110.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 as programmable computers. In International Conference on 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.15166 , 2023.

Rik Koncel-Kedziorski, Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. Mawps: A math word problem repository. In Proceedings of 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 on Learning Representations , 2020.

Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang, Xuying Meng, Siqi Fan, Peng Han, Jing Li, Li Du, Bowen Qin, et al. Flm-101b: An open llm 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.

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-Yih 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 Kernion, 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.11895 , 2022.

Arkil Patel, Satwik Bhattamishra, and Navin Goyal. Are nlp models really able to solve simple math word problems? In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: 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 Linguistics: 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. In 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, et al. Ul2: 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.09288 , 2023.

Peihao Wang, Rameswar Panda, Lucas Torroba 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 arXiv:2303.00980 , 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 Thirty-seventh Conference on Neural Information Processing Systems , 2023.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought 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:2311.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.

A Experimental Details

A.1 Pretraining details

Model architecture. We use a decoder-only model and train it using the UL2 objective [Tay et al., 2022] with 60% causal LM, 20% prefix LM 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 to the 1B model, except it uses 48 layers instead of 24. The 8B model uses 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 to 0.86 epochs of C4, 9 epochs of Wikipedia, 0.58 epochs of Arxiv and 0.44 epochs of Github.

Training details. For the 1B and 2B models, we use a cosine learning schedule with a peak learning rate of 0 . 01 that decays to 0 . 001 in the end, and use a batch size of 512. For the 8B model we use a peak learning rate of 0 . 001 and decay it to 0 . 0001 , and use a batch size of 1024. Peak learning rate was tuned to 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.

B Details for contextual reasoning primitives

In this section, we provide further details corresponding to Section 5.

All evaluations in Section 5 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, we present the inputs in 0-shot form here vs. their 5-shot form. In 5-shot form, which is how we conduct the 5-shot evaluations, each example is separated by two consecutive newline characters.

Fit a linear map from Y, to Y+1 -Yj.

0.8

r2score

cosine

0.7

level of linearity

0.6

0.5

0.4

0.3

0

2

4

6

8

10

layer index

Fit a linear map from Y, to Y+1 -Yj.

6°0

r2 score

0.8

cosine

level of linearity

0.7

0.6

0.5

0.4

0.3

0.2

0

5

10

15

20

25

30

35

layer index

For each dataset below, the inputs are separated from the targets by the '|' character (this is not a token in the input), and the targets are colored in red.

Figure 5 uses the following evaluation datasets, in the following order:

Copying (random-letter words)

Variable assignment depth 0 (code)

Variable assignment depth 1 (code)

Variable assignment depth 1 (code)

Variable assignment depth 2 (code)

Variable assignment depth 2 (code)

Pre-school math (PSM)

Copying (random-letter words):

pum nyj gdq ocu rzk jbw mlz eny kyx uni rzk jbw mlz eny kyx ___. ->|uni

Copying (real words):

eat fit ban sea vet zit pea cat van tea sea vet zit pea cat ___. ->|van

Variable assignment depth 0 (basic):

Variable assignment depth 1 (basic):

Variable assignment depth 2 (basic):

Variable assignment depth 0 (math):

Variable assignment depth 1 (math):

Variable assignment depth 2 (math):

Variable assignment depth 0 (code):

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 l=1 y=3 a=6 Question: What is the value of k? Answer:

|17 Variable assignment depth 1 (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 l=7 c=13 y=f o=c r=e u=k n=l Question: What is the value of o? Answer: |13 Variable assignment depth 2 (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 v=4 s=17 y=21 q=j l=s

B.2 Fine-tuning details

For fine-tuning, we use the 'code' variant of the variable assignment task, depths 1 and 2, in 0-shot form (i.e., no 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 tune on each, and report our results as an average across the 3 runs. Table 7 reports the standard deviations as well.

Regarding hyperparameters, we continue to use AdaFactor [Shazeer and Stern, 2018] with the same hyperparameters as in the pretraining phase, with the exception of learning rate and batch size. We use

a 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 ∈ { 200 , . . . , 300 } , chosen to be significantly after training has converged to 100% accuracy (we do not observe overfitting in this range as training continues), we 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 6 includes 5-shot evaluation results for all contextual reasoning primitives. Rows 1, 9, 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 we 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 a memorization task. We evaluate arithmetic for MIDAS and baseline training, and we see that arithmetic is 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.e., '-6+5' vs. '-1', we can successfully disentangle reasoning and memorization in our evaluation. This is equivalent to having access to a calculator, so we call it 'PSM with calculator' or 'PSM-calc' in Figure 5.

Fine tuning. Table 7 presents the fine-tuning results from Figure 5 along with corresponding standard deviations (across the 3 trials).

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 a good summary of the paper's contributions. 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 in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in 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 as long as 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 no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 important 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

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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 be 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 a dataset and/or model, the authors should describe the steps taken to 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 is a specific model and empirical evaluation, it may 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 is 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 to 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as 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 our work and prior work appear, we highlight them (Section 3 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/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Due to the computational demands of exploring a 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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 calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It 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 a 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 of range (e.g. negative error rates).

If error bars are 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: We do not report these details in the submission but can include them in a 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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs 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 NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed the NeurIPS Code of Ethics. To the best of our knowledge, this work conforms.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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/analytical paper not tied to any specific application or deployment.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 a 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 to 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: No data or models have been released with this paper.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a 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 are 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 is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was 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 to 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS 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.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1) 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1) 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?\n\nLIMITATIONS:\nThis work shows limited improvement on memorization-based tasks.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

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

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

WEAKNESSES:
Can the authors evaluate on more reasoning and common-sense questions for better comparison, such as tasks tested in the Llama paper?

Minor:
In line 105, “For simplicity, we L is divisible by k” should be “For simplicity, L is divisible by k.”

QUESTIONS:
Will the authors release the code and dataset to allow the community to further study the interesting phenomenon mentioned in section 5?

Can this phenomenon be verified with fewer resources?

LIMITATIONS:
The authors need to add a separate section for limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

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

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

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

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

LIMITATIONS:
Yes, the authors adequately addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\nSee the weakness section.\n\nLIMITATIONS:\nThe paper discusses some of the limitations in reasonable detail.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-21"}