Auto-Regressive Next-Token Predictors are Universal Learners

Listed in decreasing order of significance.

TinyStories experiment is anecdotal, compares to wrong model?

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.

Relatedly, it seems that 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.

Linear Decoders should be defined more slowly

• 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.
• A figure depicting the construction would be welcome: I had to work it out with pen and paper on the margin.
• "Under some margin conditions and using a convex surrogate loss function, this class is in fact learnable using SGD." This sentence must be supported by either a citation or a proof. As it is currently phrased it is too vague and not well-enough argued for to be used as a true statement (the current level of rigour is more appropriate for a tangential side-note).
• Observe that this class is learnable in polynomial time: Say PAC-learnable, becuase this paper uses multiple notions of learnability.

Length complexity is not well-defined

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.

On that note, it would be illuminating (and add to the cohesion of the paper) if the authors spell-out the length complexity of the construction from Theorem 7.

Missing citations

• 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 Transfomers").
• 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.

Other minor writing comments

• Page 4 above Definition 2: "...for all sub-sequences $z_{<t}$" should be prefixes (sub-sequences is correct but less accurate).
• Definition 2: Should say "for every $\epsilon, \delta \in (0,1)".
• 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.
• 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.
• Theorem 3: It would be nice to write "then \mathcal{H} = \mathcal{H}_1 \times \dots \times \mathcal{H}_T". Ideally, theorem statmenets should be as self-contained as possible as they are often used for quick reference.
• 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.
• Theorem 5: If the previous comment is accepted, this statement should be updated.
• Theorem 6: If my understanding of the proof is correct (see Summary of this review, above), 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.
• 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.
• 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.
• 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$.
• Section 2.3.1, above Theorem 11: It would be much better to avoid using $\approx$ notation in favor of notation with more well-defined meaning, such as $O(\cdot)$.
• Section 2.3.1, above Theorem 11: Should be "and a sample complexity of $\approx k \log n$" (not "the sample complexity").

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

(2) More analysis would be useful on how length complexity scales with problem complexity for different hypothesis classes.

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

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

(5) 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.

(6) 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.

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

• Multiplication experiment: unlike the TinyStories experiment, the model is not linear -- is this important? Why not use the same model for both experiments?

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

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