PaperlayerView exact revision on arXiv

Source-linked AI summary

A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks

Nikunj Saunshi, Sadhika Malladi, Sanjeev Arora

arXiv:2010.03648v2cs.CLcs.AIcs.LGstat.ML

TL;DR

The paper asks why next-word prediction helps downstream text classification, a success for which it says theoretical understanding is limited. It formalizes a sentence-completion connection and shows that cross-entropy-optimal language models yield linearly useful features for a defined class of tasks, while noting limits in the task and model scope.

  • Problem

    The paper asks why next-word prediction produces useful downstream classification features and how improvements in language-model cross-entropy translate into downstream performance.

    Source: §1 · p. 1

  • Method

    The paper reframes classification as sentence completion, formalizes linearly solvable natural tasks, and analyzes language-model features using cross-entropy optimality.

    Source: §1 · p. 2, §7 · p. 11

  • Results

    For natural tasks, an ϵ-optimal language model in cross-entropy yields features with O(√ϵ) classification error.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The analysis motivates the Quad objective, which the paper says provably learns useful features for natural classification tasks.

    Source: §1 · p. 2, §7 · p. 11

    The guarantees are limited to natural tasks with respect to Φ, and extending them to other models and more diverse tasks remains future work.

    Source: §7 · p. 11

Abstract

from arXiv · show

Autoregressive language models, pretrained using large text corpora to do well on next word prediction, have been successful at solving many downstream tasks, even with zero-shot usage. However, there is little theoretical understanding of this success. This paper initiates a mathematical study of this phenomenon for the downstream task of text classification by considering the following questions: (1) What is the intuitive connection between the pretraining task of next word prediction and text classification? (2) How can we mathematically formalize this connection and quantify the benefit of language modeling? For (1), we hypothesize, and verify empirically, that classification tasks of interest can be reformulated as sentence completion tasks, thus making language modeling a meaningful pretraining task. With a mathematical formalization of this hypothesis, we make progress towards (2) and show that language models that are $\epsilon$-optimal in cross-entropy (log-perplexity) learn features that can linearly solve such classification tasks with $\mathcal{O}(\sqrt{\epsilon})$ error, thus demonstrating that doing well on language modeling can be beneficial for downstream tasks. We experimentally verify various assumptions and theoretical findings, and also use insights from the analysis to design a new objective function that performs well on some classification tasks.

1 Introduction

The paper asks why next-word prediction yields useful downstream classification features and studies this question mathematically. It connects classification to sentence completion, formalizes the connection, and situates this analysis alongside prior work on language models and representation learning.

1 Introduction

  • The paper focuses on why cross-entropy training for next-word prediction can produce useful features for downstream text classification.It sets aside the harder question of architectural and algorithmic inductive biases because mathematically precise analysis of deep networks is challenging.Source: §1 · pp. 1–2
  • Classification tasks of interest can be reframed as sentence completion and formalized as natural tasks solvable linearly from conditional word distributions.The paper uses this formulation to analyze when language-model features help with classification.Source: §1 · p. 2
  • An ϵ-optimal language model in cross-entropy learns features that solve natural classification tasks with O(√ϵ) error.The paper also introduces conditional mean features and a theoretically motivated Quad objective, with experiments testing these ideas.Source: §1 · p. 2

1.1 Related work

  • Prior work studied language models as downstream initializations or feature maps and reframed tasks as completion, question answering, or text-to-text prediction.The paper analyzes why autoregressive language-model features work for downstream tasks, distinguishing its focus from masked language models such as BERT.Source: §1.1 · p. 2
  • The authors characterize their work as the first theoretical analysis of language-model features’ efficacy on downstream tasks.They place this contribution alongside prior theory on word and sentence embeddings, contrastive learning, and reconstruction-based representations.Source: §1.1 · p. 3

2 Language modeling and optimal solutions

The section sets notation for contexts, word distributions, and embeddings, then describes cross-entropy language modeling and its optimal solutions. It distinguishes unrestricted distributions from low-dimensional softmax models, whose guarantees concern only an embedding-defined subspace.

2 Language modeling and optimal solutions

  • The language-modeling setup treats contexts as complete or partial sentences and models a distribution over the next word in the vocabulary.A feature map embeds each context, and word embeddings parameterize the softmax next-word distribution.Source: §2 · p. 3, §2.2 · p. 4

2.1 Language modeling using cross-entropy

  • The cross-entropy objective is minimized when the predicted conditional word distribution matches the true distribution on contexts in the language-model corpus.The paper notes that practical optimization, statistical, and expressivity limitations can prevent attaining this optimum.Source: §2.1 · p. 3

2.2 Softmax parametrized language modeling

  • For fixed word embeddings, an optimal softmax model matches the true distribution after projection onto the embedding-defined subspace, not necessarily in the full vocabulary space.The stated recovery guarantee concerns Φp_f*(s), which equals Φp*(·|s) for contexts in the language-model distribution’s support.Source: §2.2 · p. 4

3 Using language models for classification tasks

The paper connects next-word prediction to classification by treating labels as sentence completions, then formalizes which classification tasks can be solved from language-model word probabilities. It illustrates this connection with prompts and indicative completion words, and defines natural tasks to capture the relevant linear-solvability condition.

3 Using language models for classification tasks

  • Classification tasks can be reframed as sentence completion: use next-word probabilities after an input text to predict its label.The paper presents this as an oracle thought experiment and formalizes the reformulation as natural tasks.Source: §3 · p. 4

3.1 Sentence completion reformulation

  • Prompts such as “This movie is ” or “This article is about ” can elicit indicative completion words for sentiment or news categories.The authors report that SST and AG News can be solved using linear functions of probabilities for small subsets of words, with further tasks discussed in Section F.1.Source: §3.1 · p. 5

3.2 Natural classification tasks

  • A task is (τ, B)-natural when a linear classifier with norm bounded by B achieves loss τ using true next-word probabilities as features.For softmax models, the corresponding definition restricts the classifier to the row span of the word-embedding matrix Φ; the converse need not hold for all natural tasks.Source: §3.2 · pp. 5–6

4 Guarantees for language models on natural tasks

The paper derives classification guarantees from cross-entropy optimality for natural tasks, considering arbitrary language models and softmax models with conditional mean features. It also relates those features to the model’s ordinary features under an embedding assumption.

4 Guarantees for language models on natural tasks

  • The guarantees consider arbitrary language models using next-word distributions as features or softmax models using d-dimensional conditional mean features, with models assumed ϵ-optimal in cross-entropy.The softmax setup measures suboptimality relative to the best d-dimensional model using Φ.Source: §4 · p. 6

4.1 Arbitary language models

  • For arbitrary language models, the classification-loss bound improves with smaller cross-entropy error and depends on task naturalness, indicative-word mass, and distribution overlap.The overlap coefficient γ(pT) captures how much of the language-model context distribution covers the downstream distribution; the result is presented as applying regardless of parametrization.Source: §4.1 · pp. 6–7

4.2 Softmax language model with conditional mean features

  • For softmax models, conditional mean features Φpf are guaranteed to perform well on some natural tasks and are empirically verified, but a fixed error term motivates a stronger result.The earlier bound measures suboptimality against the best d-dimensional model and includes fixed error O(Bpϵ*), even when ϵ = 0.Source: §4.2 · pp. 7–8

4.3 Φpf(s) is a linear function of f(s)

  • Under Assumption 4.1, Φpf(s) = Af(s) + b, and the paper reports that pretrained GPT-2 exhibits only a weak version of this linear relationship.The authors note that this relationship may not fully explain why f often performs better than Φpf.Source: §4.3 · p. 8
  • ℓT(f) ≤ τ + O(B√ϵ) under the stated setting, giving a classification guarantee for the language model’s ordinary features.This follows under Lemma 4.3 and Theorem 4.2’s setting.Source: §4.3 · p. 8

5 Extensions

The extensions refine how language-model guarantees transfer across context distributions and introduce Quad, an objective designed to learn useful, interpretable features for natural classification tasks.

5.1 Better handling of distributional shift

  • When prediction errors are random-like and independent of context, the refined transfer coefficient can be approximately 1, avoiding worst-case accumulation across target contexts.The refined coefficient depends on the learned model and can be smaller than the distribution-only coefficient.Source: §5.1 · pp. 8–9

5.2 Quad: A new objective function

  • Quad replaces the cross-entropy log-partition function with a quadratic function, yielding an optimal feature map linked to the substitutability matrix’s eigendecomposition.The matrix captures substitutability between words; the optimal feature map has a closed-form solution.Source: §5.2 · p. 9
  • For tasks that are (τ, B)-natural with respect to Φ*, the Quad optimum has task loss at most τ.The theorem gives Φ* = BU_d^⊤ for full-rank B, where U_d contains the top d eigenvectors of the substitutability matrix.Source: §5.2 · p. 10
  • The optimal Quad feature map is the best d-dimensional projection of the substitutability matrix, so substitutable words receive matching embeddings.The paper relates substitutable words to synonyms.Source: §5.2 · p. 10
  • Quad’s results suggest performance comparable to conditional-mean features learned by the cross-entropy objective, consistent with both being linear functions of conditional word probabilities.The comparison is reported in Table 3; no numerical values are supplied here.Source: §5.2 · p. 10

6 Experiments

Experiments test the classification features, sentence-completion reformulation, and a quadratic fit assumption, finding useful performance from fixed GPT-2 features and task-relevant token probabilities.

Tasks using linear function of p*|s

  • Probabilities for roughly 30 task-relevant tokens solve binary and fine-grained SST and AG News classification, and task-specific prompts consistently improve performance, including fine-tuning with small datasets.Even one or two class words yield non-trivial performance.Source: §6 · p. 10

Φpf and f are good features

  • Linear classifiers on fixed pretrained features perform comparably to fine-tuning, while conditional-mean features also work but perform worse than the pretrained features.Adding a task-specific prompt improves performance, and random projections of the token probabilities perform very poorly.Source: §6 · p. 10
  • The fit for the log partition function of features from the pretrained GPT-2 model holds on both its training and unseen data, including WebText.Source: §6 · p. 11

7 Conclusions and future work

The paper explains language-model features for classification through sentence completion and formalizes the connection as natural tasks. It also identifies limits of its analysis and outlines further results and proof material.

7 Conclusions and future work

  • The paper formalizes sentence-completion reformulations as natural tasks, which can be solved linearly using conditional word probabilities, and uses this framework to design Quad.Quad is presented as provably learning useful features for these tasks.Source: §7 · p. 11
  • The analysis does not explain why pretrained features f outperform conditional-mean features Φpf, and its guarantees focus on natural tasks relative to Φ.The authors also identify guarantees for other models and tasks, and theory for fine-tuning, as future directions.Source: §7 · p. 11

A Overview

  • The appendices provide proofs, further experiments, and additional discussion of natural tasks, word embeddings, and the assumptions behind the results.They include experiments on more tasks, larger-corpus Quad training, and the O(√ϵ) dependence.Source: A Overview · p. 16

B Better handling of distributional shift

The distribution-shift extension defines transferability coefficients that account for learned-model errors, then states strengthened guarantees for natural tasks.

B Better handling of distributional shift

  • The refined transferability coefficients incorporate prediction errors and their covariance across contexts, rather than depending only on the language-model and target context distributions.The coefficients can exploit the learned model’s inductive bias and be smaller than the distribution-only coefficient.Source: B Better handling of distributional… · pp. 16–17
  • The strengthened theorems give classification guarantees for tasks that are (τ, B)-natural with respect to the relevant feature map or embeddings.The supplied theorem statements introduce ϵ-optimal language models and natural-task conditions, but do not include the complete bound expressions.Source: B Better handling of distributional… · pp. 16–17
  • If prediction errors are random-like and independent of context, the refined coefficient can be approximately 1, preventing errors from accumulating on target contexts.This bypasses the worst-case transfer captured by γ(pT).Source: B Better handling of distributional… · p. 17

C Quad: A new objective function

The Quad objective is introduced because cross-entropy does not give a simple closed form for optimal word embeddings; Quad makes the optimal feature map linear in next-word probabilities and gives the optimal embedding matrix a closed-form characterization. The section also frames its discussions of natural tasks as intuition rather than formal argument.

C Quad: A new objective function

  • Quad addresses cross-entropy’s lack of a simple closed-form optimal embedding solution, while its optimal feature map is linear in next-word probabilities and its optimal embedding matrix has a closed-form solution.The authors also argue that the resulting features can solve some natural tasks.Source: C Quad: A new objective function · p. 17
  • Quad replaces cross-entropy’s log-partition function with a quadratic function, and its optimal solution depends on the substitutability matrix’s eigendecomposition.The matrix captures substitutability between words with identical conditional probabilities across contexts.Source: C Quad: A new objective function · p. 17
  • Theorem 5.2 characterizes the optimal embedding matrix as Φ*=BU_d^T, and a fixed Φ yields optimal features f*(s)=(ΦΦ^T)^-1/2Φp*_{.|s}; natural tasks then have loss at most τ.B is specified as a full-rank d-by-d matrix in the supplied theorem passage.Source: C Quad: A new objective function · p. 18

D More on natural tasks

  • The section’s discussions of natural tasks are intended to provide intuition for definitions and results, and may not be formal or precise.Source: D More on natural tasks · p. 18

D.1 Sentence completion reformulation ≡natural task

This section connects sentence-completion reformulations to natural-task bounds, discusses qualitative conditions on word embeddings, and outlines proof strategies for arbitrary and softmax language models. It also identifies the softmax variant of Pinsker’s inequality as a proof technique.

D.1 Sentence completion reformulation ≡natural task

  • Sentence-completion tasks yield linear predictors over p*|s with hinge loss at most 4 Bayes-Error(T), and are (τ,B)-natural with τ=O(Bayes-Error(T)) and B=O(α(Windicative)^−1).The resulting bound is smaller when the task has low Bayes error and indicative completions have large, frequent-word mass.Source: §D.1 · pp. 18, 20

D.2 Nice properties of word embeddings Φ

  • If Φ captures top semantic directions associated with synonym structure and the task separates coarse meaning differences, projection residuals can be small, yielding small τ′ and B′; these claims are qualitative.The authors leave formalizing desirable properties of Φ to future work.Source: §D.2 · pp. 20–22

D.3 Proofs for Section D.1

  • The proof analyzes predictor gpT from Equation (14): since gpT(s) ≤ 1, its hinge loss is ℓ(gpT(s), y) = (1 − ygpT(s))+ = 1 − ygpT(s) for every s, y; Proposition D.2 sets B = α(Windicative)−1.Source: §D.3 · pp. 22–23

E.1 Proof sketch

  • The following subsections sketch arguments for Theorems 4.1 and 4.2, with full proofs of strengthened versions presented afterward.Source: §E.1 · p. 23

E.1.1 Proof sketch for arbitrary language models: Theorem 4.1

  • For arbitrary language models, the proof bounds classification error by linking loss to prediction-error covariance, transferring covariance from pT to pL, and using Pinsker’s inequality to connect it to cross-entropy.The covariance transfer step uses the transferability coefficient; the final step averages over pL.Source: §E.1.1 · p. 24

E.1.2 Proof sketch for softmax language models: Theorem 4.2

  • For softmax models, the cross-entropy bound uses modified Pinsker under a classifier-in-span-Φ condition and an evaluated model that computes softmax using Φ, comparing against the best softmax model.The bound therefore incurs suboptimality relative to the best softmax model rather than the absolute best language model.Source: §E.1.2 · p. 24
  • The derivation uses the softmax variant of Pinsker’s inequality in Lemma E.4 before presenting detailed proofs.Source: §E.1.2 · p. 24

E.2 Proofs for arbitrary language models

These proof sections establish downstream guarantees by relating classification loss to language-model error and applying the natural-task conditions. They also connect softmax features to unrestricted features under an assumption, and derive the Quad solution from the substitutability matrix’s leading eigenspace.

E.2 Proofs for arbitrary language models

  • For arbitrary language models, the proof bounds downstream performance on natural tasks using covariance of prediction errors, then relates that covariance to language-model suboptimality.The proof combines these two steps to establish the theorem’s guarantee for ϵ-optimal language models.Source: §E.2 · pp. 25–26

E.3 Proofs for softmax language models

  • For fixed embeddings, the softmax-language-model proof establishes a corresponding natural-task guarantee for features from an ϵ-optimal d-dimensional model.The supplied theorem passage states the setting and guarantee but does not include the bound’s full expression.Source: §E.3 · pp. 26–27

E.4 Proofs for Section 4.3

  • Under Assumption 4.1, Φp_f(s)=Af(s)+b, and the resulting feature relationship yields ℓ_T(f)≤τ+O(B√ϵ) for natural tasks.The proof passage also describes Assumption 4.1 as approximately true when word embeddings are Gaussian-like.Source: §E.4 · pp. 27–28

E.5 Proofs for Section C

  • The Quad optimum uses the top d eigenvectors of the substitutability matrix, and with fixed Φ its optimal feature map is a normalized projection of p*_{.|s}.The proof assumes distinct eigenvalues for simplicity of presentation.Source: §E.5 · pp. 28–30

E.6 Proofs for supporting lemmas

The supporting proofs connect classification loss to covariance bounds, use distribution-shift arguments to transfer bounds, and establish inequalities used in the main results. The experiment details specify the model and the smaller SST-2 version used in the experiments.

E.6.2 Handling distribution shift

  • The distribution-shift proof transfers a bound from the task distribution p_T to the language-model distribution p_L.The supplied proof passages identify this transfer as a step in completing the lemma.Source: §E.6 · p. 31, §E.6.2 · pp. 34–35

E.6 Proofs for supporting lemmas

  • The supporting lemmas use Pinsker-type inequalities and smoothness arguments to relate language-model suboptimality to bounds used in the main proofs.The softmax variant of Pinsker’s inequality is established using convexity of the KL-based loss.Source: §E.6 · pp. 31–33

E.6.1 Classification loss to covariance of error

  • For classification loss, the proof bounds loss using covariance of prediction error through Lipschitzness of the loss and Jensen’s inequality.Source: §E.6.1 · p. 34

F Experiment Details

  • Experiments use the 117M-parameter small GPT-2 model, and the SST-2 version has 6,920 training and 1,821 test examples.The authors say this SST-2 version is smaller than the GLUE version and better fits their sentence-completion hypothesis.Source: F Experiment Details · p. 35

F.1 Solving downstream tasks using f and Φpf

The experiments evaluate fixed language-model features and conditional mean features on downstream classification, then test the Quad objective and a quadratic approximation of the log-partition function. Quad is comparable to cross-entropy features in the reported cases, while the conditional mean features are comparable to the original features across tasks.

F.1 Solving downstream tasks using f and Φpf

  • Conditional mean features had comparable performance to GPT-2 features across sentiment and non-sentiment tasks.For sentiment tasks, prompts almost always boosted performance, and two-dimensional positive/negative class-word features were more sample-efficient than standard 768-dimensional features.Source: §F.1 · p. 36

F.2 Finetuning Experiments

  • Tables 3 and 4 compare Quad features with cross-entropy features on IMDb and Amazon, with a prompt condition also reported.The IMDb experiment fixes the cross-entropy output embeddings and initializes Quad with the SVD of pretrained embeddings; the Amazon caption notes validation loss was still decreasing at measurement.Source: Table 3 · p. 38, Table 4 · p. 38, §F.3 · p. 38

F.3 Testing Quad objective

  • Quad achieved performance comparable to cross-entropy in both tested cases, but did not outperform it on SST.Adding a prompt consistently improved performance for both objectives; the authors leave the performance gap and broader evaluation for future work.Source: §F.3 · p. 39

F.4 Learning the quadratic approximation of the log-partition function

  • A quadratic approximation of the log-partition function is fit to support the assumed quadratic form and the resulting linear relation between features.The fit is learned from a regression loss, with a symmetric positive-semidefinite matrix parameterized as A = UU^T.Source: §F.4 · pp. 39–40

F.5 Experimentally Checking Theorem 4.2

The experiment tests whether Theorem 4.2’s square-root relationship appears in practice by tracking language-model cross-entropy and conditional-mean-feature loss across training checkpoints. The observed points exhibit a square-root trend.

F.5 Experimentally Checking Theorem 4.2

  • In practice, the experiment found some square-root trend by plotting validation-set model cross-entropy against conditional-mean-feature performance on SST-2 across training checkpoints, although Theorem 4.2 provides only an upper bound.Source: §F.5 · pp. 40–41
Loading 2010.03648v2…