Paper: Auto-Regressive Next-Token Predictors are Universal Learners (ICLR-2024)

Reviewer comment (Evaluation):
I am not sure what are the actual results being reported with the TinyStories experiment: What I found is a footnote on a 1.2 difference in perplexity between the linear predictor and GPT-2 Small---but I'm not sure what to make of this quantity. And there is a statement that the linear predictor "often does produce coherent text". But how often? And how do you measure coherence? While I lack the expertise to suggest a concrete experiment, I expect a higher level of rigour for a paper at ICLR. Perhaps the authors could refer to the TinyStories paper and reproduce some of the experiments there with their model.

Author response:
TinyStories experiment: following your remarks, we will conduct and report additional thorough and rigorous experiments on the TinyStories dataset in the final version of the paper. We will use some of the experimental settings in the original TinyStories paper to provide more quantitative results on the TinyStories dataset. We will also compare our linear model to the smaller transformer models studied in the TinyStories paper. However, note that we do not claim that linear models are better than transformers for language modeling, and certainly they can be much less efficient in terms of parameters. The primary reason for studying linear auto-regressive models is because we can achieve a more complete theoretical understanding of their optimization and expressive power, something that we are not able to do for Transformers. Hence, the TinyStories experiment is mainly used for demonstrating that linear models can generate reasonable text and give non-trivial performance, a result that we found surprising.

Reviewer comment (Experiments):
In the TinyStories paper there is a 28M-parameter model that achieves performance comparable to GPT2-XL. Therefore, the comparison of the linear predictor to GPT2-Small seems inappropriate---much better performance can be attained with 5x less parameters. Perhaps the authors should compare the linear predictor to that model instead. Discovering that a linear predictor requires significantly more parameters than a (good) transformer model to achieve comparable performance would undermine the main message of the paper (that the power of LLMs can be attributed to AR rather than architecture). Please correct me if you disagree that comparing to the TinyStories transformer is a more fair comparison towards this end.

Author response:
TinyStories experiment: following your remarks, we will conduct and report additional thorough and rigorous experiments on the TinyStories dataset in the final version of the paper. We will use some of the experimental settings in the original TinyStories paper to provide more quantitative results on the TinyStories dataset. We will also compare our linear model to the smaller transformer models studied in the TinyStories paper. However, note that we do not claim that linear models are better than transformers for language modeling, and certainly they can be much less efficient in terms of parameters. The primary reason for studying linear auto-regressive models is because we can achieve a more complete theoretical understanding of their optimization and expressive power, something that we are not able to do for Transformers. Hence, the TinyStories experiment is mainly used for demonstrating that linear models can generate reasonable text and give non-trivial performance, a result that we found surprising.

Reviewer comment (Writing):
Given the significance of linear ARs (linear decoders) throughout the paper, I suggest defining them more slowly and in their own Definition environment---rather than presenting them as an "Example". A Definition environment is more easy to refer back to, and encourages more formal writing.

Author response:
Definition of linear decoders: we have updated the paper to reflect your comments, giving a more structured introduction of linear decoders.

Reviewer comment (Theory):
As length complexity is currently defined in Definition 9, a class $\mathcal{H}$ has infinitely many length complexities for computing a given $\mathcal{F}$. This is because if $T$ is a length complexity of $\mathcal{H}$ for computing $\mathcal{F}$, then so is any $T' \geq T$. Instead, you should define the length complexity of $\mathcal{H}$ computing $\mathcal{F}$ to be the minimal $T$ for which the conditions stated in the definition hold.

Author response:
Length complexity definition: we updated the paper to clarify the definition of length complexity.

Reviewer comment (Novelty):
Missing citation — "Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective" by Feng et al. (2023): This work gives a complexity-theoretic explanation to the success of chain-of-thought, especially in mathematical reasoning. The techniques seem entirely different, and it is focused on transformer architectures (rather than focusing on the autoregressive aspect of LLMs, as in the paper currently under review). Still, it clearly fits in three of the four topics covered by the related work (all save for "Beyond Transformers").

Author response:
Missing citations: we added the citations you suggested to the updated manuscript.

Reviewer comment (Novelty):
Missing citation — "Fast Learning Requires Good Memory: A Time-Space Lower Bound for Parity Learning" by Raz (2016): This prominent result (FOCS best paper, JACM) should probably be mentioned in the first paragraph of Section 2.3.1 that discusses hardness results for parities.

Author response:
Missing citations: we added the citations you suggested to the updated manuscript.

Reviewer comment (Writing):
Page 4 above Definition 2: "...for all **sub-sequences** $z_{<t}$" should be **prefixes** (sub-sequences is correct but less accurate).

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Theory):
Definition 2: Should say "for every $\epsilon, \delta \in (0,1)".

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Theory):
Definition 2: Should say "If there exists $m \colon (0,1)^2 \to \mathbb{N}$" such that... Otherwise, the order of quantifiers / role of $m$ is unclear.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Definition 2: Should say "returns w.p. $\geq 1-\delta$ a function $h \in \mathcal{H}$". Also, the first use of w.p. should be explicitly defined ("with probability (w.p.)") so that the paper is accessible to a broad audience.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Theorem 3: It would be nice to write "then $\mathcal{H} = \mathcal{H}_1 \times \dots \times \mathcal{H}_T$". Ideally, theorem statements should be as self-contained as possible as they are often used for quick reference.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Definition 4: I would say that "$h$ computes $f$ w.r.t $\mathcal{D}$". That is, explicitly include $\mathcal{D}$ in the definition, since it is crucial. Likewise for the definition of _approximates_.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Theory):
Theorem 5: If the previous comment is accepted, this statement should be updated.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Theorem 6: If my understanding of the proof is correct, it is worth adding a sentence explaining what the dataset is to the body of the paper—rather than just saying the main fact the proof relies on, explain how it is used.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Theorem 7 and Corollary 8: Use the notation $\mathcal{H}^{\mathrm{lin}}$ here when referring to linear AR functions/models. Otherwise, the notation surprisingly re-appears on the next page.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Section 2.3.1: This is a matter of taste, but I would find it more informative to use $\mathcal{P}_n$ to denote the class of parities on $n$ bits. This is because $\mathcal{F}$ was previously used to refer to a generic class of functions.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Writing):
Section 2.3.1, second paragraph after Theorem 10: when defining $\mathcal{F}_{n,k}$, it should be $A \in \binom{[n]}{\leq k}$. Note the square brackets around $n$.

Author response:
Other minor writing comments: we addressed all the comments in the updated paper. Thank you again for these detailed and valuable suggestions.

Reviewer comment (Reproducibility):
Question — Do the authors intend to release the code used for their experiments so that the reader may reproduce them?

Author response:
We also plan to publish the code for reproducing all the experiments along with the final paper.

Reviewer comment (Theory):
The theoretical results rely on very strong assumptions about availability of chain-of-thought training data, which may be unrealistic.

Author response:
Indeed our theoretical results rely on availability of chain-of-thought training data. We note that many tasks do have CoT data for training (or fine-tuning), and that training on data with CoT demonstrations was found to greatly improve performance. For example, the GSM8K dataset contains CoT answers, and recent works suggest that augmenting the training data with different forms of CoT dramatically improves performance (e.g. https://arxiv.org/pdf/2309.12284.pdf). We do agree however that our theoretical results rely on availability of “unrealistic” CoT data, and we acknowledge this in the paper as well. As is often the case in many theoretical works, we study an idealized setting that allows us to give concrete analysis of specific simple models, studying their capabilities and limitations. We believe that our main theoretical result in this setting, showing that simple models such as linear autoregressive predictors can learn to compute any (computable) function with CoT, is surprising and novel.

Reviewer comment (Theory):
More analysis would be useful on how length complexity scales with problem complexity for different hypothesis classes.

Author response:
Our work provides a novel theoretical framework for studying the interaction between problem complexity and length complexity, and we derive some preliminary analysis of length complexity for specific problems. Since our main contribution in this work is in introducing such a theoretical framework, we leave a further investigation to future work.

Reviewer comment (Experiments):
Additional validation on more complex architectures like Transformers would strengthen the conclusions about training scheme vs architecture.

Author response:
Following your suggestion, we will add additional experiments comparing non-autoregressive to autoregressive Transformers in the final version of the paper.

Reviewer comment (Theory):
The proposed linear models are not exactly equivalent to the classical linear models analyzed. Some parameters are shared across time steps.

Author response:
Indeed, the linear model trained in our experiments is in fact a linear network, which is somewhat different from the models we analyze. However, we note that the overall model in the experiment still computes a linear function of the inputs. In other words, the models used in the experiments are still linear auto-regressive models (as defined in the theory section), but they are trained with a linear-network parameterization which enforces some weight sharing.

Reviewer comment (Experiments):
Non autoregressive models are not discussed, tested and compared. It will be to know if given large amount of training data, non-autoregressive models can perform on par with autoregressive models.

Author response:
Our theoretical results imply that, at least for linear models, non-autoregressive models are significantly weaker than autoregressive models. This is true regardless of the amount of training data used, since non-autoregressive linear models cannot express rich functions that autoregressive linear models can. We note that very recently, a similar theoretical result was shown for the Transformer architecture (https://arxiv.org/pdf/2310.07923.pdf).

Reviewer comment (Experiments):
On MULTIPLICATION experiments, models are trained and tested on 4-digit numbers. The model is likely to memorize the patterns instead of generalizations, especially this paper use special tokenization for digits and signs.

Author response:
Generalization on the multiplication experiment: note that we split the data to train and evaluation pairs, and report the performance when multiplying pairs of numbers that were never seen during training. Therefore, the model generalizes to unseen pairs, applying the same algorithm learned during training.

Reviewer comment (Evaluation):
The TinyStories experiment (Section 3.1) lacks quantitative evaluation. It is unclear if the examples shown are representative.

Author response:
TinyStories experiments: We will conduct and report additional thorough and rigorous experiments on the TinyStories dataset in the final version of the paper, and add additional prompts and completions. We do note that the examples shown in the paper are representative and are not “cherry-picked” examples.

Reviewer comment (Experiments):
Multiplication experiment: unlike the TinyStories experiment, the model is not linear -- is this important? Why not use the same model for both experiments?

Author response:
Multiplication experiment: Indeed, we use an MLP instead of a linear network for the multiplication experiment, as we found them to give more consistent performance. We do note that these models are still extremely simple compared to the standard transformer architecture.

Reviewer comment (Theory):
There appears to be a potential mismatch with realistic chain-of-thought prompting in that the learnability theory developed here assumes that the tasks are available together with their full sequences of intermediate steps in the training set, whereas in reality prompts containing chain-of-thought demonstrations may be quite unnatural, and are quite likely not to have appeared in that form in the training set. In this sense, the learnability theory, if understood as applying to unsupervised LLM training, does not explain why LMs would be able to deal with unnatural prompt formats.

Author response:
Availability of chain-of-thought during training: We agree that in some cases, chain-of-thought reasoning for specific tasks does not appear in the training set. However, note that many tasks do have CoT data for training (or fine-tuning), and that training on data with CoT demonstrations was found to greatly improve performance. For example, the GSM8K dataset contains CoT answers, and recent works suggest that augmenting the training data with different forms of CoT dramatically improves performance (e.g. https://arxiv.org/pdf/2309.12284.pdf). Our work focuses on such settings, where CoT is available as part of the training data.

Reviewer comment (Experiments):
In the multiplication experiment, the calculations are decomposed into "more intermediate steps than in Liu&Low". Furthermore, they involve padding to get all strings to have the same length, apparently unlike Liu&Low. Both of these steps may impact the comparison with the results from Liu&Low. Furthermore, unlike the MLP, GPT-3.5 is (judging by Figure 2) not prompted to output intermediate steps. Thus, it is unclear what one can learn from the comparison between the models in Figure 2 (right), and the statement "outperforms GPT-4" in the introduction may not be fully supported.

Author response:
More intermediate steps in the multiplication experiment: indeed, our model is trained with longer chain-of-thought data compared to the paper of Lie&Low, and is also compared to GPT-3.5/4 which do not use CoT in their answer. On the other hand, in our work we use an extremely simple architecture (shallow MLP) compared to the gigantic transformers in competing methods. Our goal in our experiment is not to convince readers that MLPs or linear networks are better than transformers - they clearly are not. Our only goal is to emphasize that longer CoT data can make such simple models competitive with much larger and more complex networks. In other words, instead of adding complexity to the model, we are adding complexity to the CoT training data.

Reviewer comment (Reproducibility):
Section 3.1: what is the objective function -- is it cross-entropy? Is softmax applied on the linear output?

Author response:
Section 3.1: yes, we train with standard cross-entropy loss, with softmax over the outputs. We clarify this in an updated version of the manuscript.