A Mathematical Exploration of Why Language Models Help Solve Downstream Tasks
Authors: Nikunj Saunshi
Sadhika Malladi
Sanjeev Arora
Venue: ICLR 2021 Poster (2021)
Keywords: language models
theory
representation learning
self-supervised learning
unsupervised learning
transfer learning
natural language processing

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

Published as a conference paper at ICLR 2021
A MATHEMATICAL EXPLORATION OF WHY LAN- GUAGE MODELS HELP SOLVE DOWNSTREAM TASKS
Nikunj Saunshi, Sadhika Malladi & Sanjeev Arora
Princeton University
(nsaunshi, small ladi, arora)ecs http:/css..iinnnttoo...oo
ABSTRACT
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 mathematical study of this phenomenon for the downstream task of text classification by considering the following questions: (I) What is the intuitive connection between the pretraining task of next word prediction and text classification? (2) How
mathematically formalize this connection and quantify the benefit of language modeling? For (1), we hypoth- esize, and verify empirically, that classification tasks of interest can be reformu- lated as sentence completion tasks, thus making language modeling a meaningful pretraining task. With mathematical formalization of this hypothesis, we make progress towards (2) and show that language models that are e-optimal in cross- entropy (log-perplexity) learn features that can linearly solve such classification tasks with ((E)) error, thus demonstrating that doing well on language modeling can be beneficial for downstream tasks. We experimentally verify various assump- tions and theoretical findings, and also use insights from the analysis to design a new objective function that performs well on some classification tasks.
INTRODUCTION
The construction of increasingly powerful language models has revolutionized natural language processing (NLP). Using gigantic text corpora and a cross-entropy objective, language models are trained to predict distribution over the next word to follow a given context (piece of text). Pretrained language models are useful for many downstream NLP tasks, either as initializations (Ramachandran
et al. 2017 Howard & Ruder 2018) or source of contextual word embeddings (McCann et al.)
2017 Peters et al.) 2018). Recent models (Radford et al. 2019 Brown et al., 2020) have even bypassed the need for careful fine-tuning and have demonstrated strong performance downstream tasks without fine-tuning. This work aims to understand this incredible success of language models.
Since next word prediction is powerful test of language understanding, at an intuitive level it is believable that doing well on language modeling can help with many diverse NLP tasks. At the same time, it is quite intriguing how improvements in the test perplexity of language models translate to better downstream performance. Attempting to understand this phenomenon naturally raises the following questions: (a) why should training the next-word prediction task, with the cross-entropy objective, result in useful features for downstream tasks? (b) what role do induccive biases of the model architecture and training algorithms play this empirical success? Given the nascency of deep learning theory, it is very challenging to say anything mathematically precise about (b) for deep networks. Given these difficulties, this paper focusses on the mathematical study of (a) by exploring if and how quantitative improvements on downstream NLP tasks can be mathematically guaranteed for language models that do well on the cross-entropy objective. As a first cut analysis, restrict attention to text classification tasks and the striking observation that they can be solved fairly well with linear classifiers on top of fixed language models features, i.e. without finetuning treat models as black boxes, just first-order optimality conditions of the
(Table D. Although cross- cnossyyytropy objective reveal interesting properties of learned features, leading to an understanding of their success classification tasks. Insights from the analysis help us construct simple objective

---

Published as a conference paper at ICLR 2021
(Quad), that provably learns useful features for classification tasks, as also verified empirically, We summarize our contributions along with an overview of the paper below.
In Section we set up notation and formally describe language modeling and the ubiquitous low- dimensional softmax paramctrization, along with a description of the cross-entropy objective and properties of its optimal solutions. We then describe the observation, in Section [311 that text clas- sification tasks of interest can be reformulated as sentence completion tasks. Amenability to such a reformulation is mathematically formalized (Section 3.2) as the classification task being natural task: tasks that can be solved linearly using conditional distribution over words following an input text. Section
4 presents our main results, theorems 4.1 and 4.2 that use the above formalization to mathematically quantify the utility of language model features on natural tasks: e-optimal language model (in cross-entropy) will do ((V))well on such tasks. Theorem 4.2 shows a stronger result for low-dimensional softmax models by leveraging new tool, conditional mean features (Defini- tion (. which we show (Section (6 to be effective in practice. The usefulness of the language model features themselves is demonstrated by arguing weak linear relationship between them and conditional mean features. In Section 5.2 we present a new mathematically motivated objective (Quad) that has formal guarantees. Experiments in Section 6 verify the sentence completion refor- mulation idea and the good performance of conditional mean features on standard benchmarks.
1.1 RELATED WORK
Text embedding methods: Prior to language models, large text corpora like Wikipedia (Merity) ee al. (016)) were used to learn low-dimensional embeddings for words (Mikolov et al.] 201 a Pennington al., 2014) and subsequently for sentences (Kiros et al., 2015; Arora al. 2017
Pagliardini et al. 2018 Logeswaran & Lee 2018) for downstream task usage. These methods were
inspired by the distributional hypothesis (Firth 1957 Harris 1954), which posits that meaning of
text is determined in part by the surrounding context. Recent methods like BERT (Devlin et al.)
2018) and variants (Lan et al. 2019 Yang al., 2019 Liu et al. 2019) learn models from auxiliary tasks, such as sentence completion, and are among the top performers on downstream tasks. In this work we consider autoregressive models and make distinction from masked language models like BERT; Table shows that language model and BERT features have comparable performances.
Language models for downstream tasks: We are interested in language models (Chen & Good-
man 1999). especially those that use neural networks to compute low-dimensional features for con-
texts and parametrize the next word distribution using softmax (Xu & Rudnicky. 2000 Bengio
et al. 2003). Language models have shown to be useful for downstream tasks as initializations (Ra-
machandran et al. 2017 Howard & Ruder 2018) or as learned feature maps (Radford et al.) [017]
McCann et al.) 20117; Peters et al. 2018). The idea of phrasing classification tasks as sentence com-
pletion problems to use language models is motivated by recent works (Radford et al.) 2019 Puri & Catanzaro 2019 Schick& & Schütze 2020) that show that many downstream tasks can be solved by next word prediction for an appropriately conditioned language model. This idea also shares similarities with work that phrase suite of downstream tasks as gusserng tasks (McCann eet al., 2018) or rext-to-text tasks (Raffel et al., 2019) and symbolic reasoning as fill-in-the-blant tasks Talmor et al. 2019). Our work exploits this prevalent idea of task rephrasing to theoretically analyze why language models succeed on downstream tasks.
Relevant theory: Since the success of early word embedding algorithms like word2vec (Mikolov et al. 2013a) and GloVe (Pennington et al. 2014), there have been attempts to understand them theoretically. Levy & Goldberg (2014) argue that word2vec algorithm implicitly factorizes the PMI matrix. Noise Contrastive Estimation (NCE) theory is used to understand word embeddings Dyerr 2014) and to show parameter recovery for negative sampling based conditional models (Ma) & Collins 2018). A latent variable model (Arora et al. (0166) is used to explain and unify various word embedding algorithms. Theoretical justification is provided for sentence embedding meth- ods either by using a latent variable model (Arora et al. 2017) or through the lens of compressed sensing (Arora et al. 2018). Also relevant is recent work on theory for contrastive learning (Arora et al. |219 Tosh et a 2020b al Wang & Isola 2020) and reconstruction-basse methods (Lee al.) 2020). which analyze the utility of self-supervised representations learned for downstream tasks. Our work is the first to analyze the efficacy of language model features on downstream tasks.
2

---

Published as a conference paper at ICLR 2021
2 LANGUAGE MODELING AND OPTIMAL SOLUTIONS
We use S to denote the discrete set of all contexts, i.. complete or partial sentences (prefixes), W to denote the vocabulary of words, with V [W] being the vocabulary size. For a discrete set A, let A denote the set of distributions on A. We use p. PL E As to denote probability distributions over S, and P.|s: P1 E Aw to denote conditional distributions, where (ij(ww) is the predicted probability of word w following context s and (.)((w) denotes the true conditional probability. Boldface P.|s. Pijs RY denote vectors of probabilities for P.js. Pi,s Aw- For RV, v(w) indexes the coordinate for w E P.I. P.I=(()) (w) is the probability of w according to p.|s- We use du E Rd denote a d-dimensional embedding for word w: word embeddings are stacked into the columns & E RdXV, We use f S - I2 for a feature map from contexts to d-dimensional embeddings, e.g. f(s) can be the output of a Transformer model for input context & E S. For embeddings {0.)ses with 0g E TR (any D), we use {0.] to denote g S > ReD such that g(s) 0..
2.1 LANGUAGE MODELING USING CROSS-ENTRONS ENTROPY
Language model aims to learn the true distribution of a text corpus and popular approach to do so is through next word prediction. Given context (c.g., a sentence
S). it predicts a distribution P,ls over the word to follow, c.g. for the context "The food was ", the model could place high prob- abilities words "delicious", "expensive". "bland", etc. We use pr. denote the true distribution over the context set S in the language modeling corpus. A standard approach is to minimize the ex- pected cross-entropy loss between the true distribution pijs and the model prediction P.ja- We define the cross-entropy loss for a language model with output vector of probabilities {pess as
Exem (Pp.s))
E supr. W
[- (o
a-pr.
To understand what language models learn, we look at the optimal solution of the cross-entropy objective. While one cannot practically hope to learn the optimal solution due to optimization, statistical and expressivity limitations, the optimal solution at least tells us the best that language modeling can hope to do. A well-known property of cross-entropy objective is that its optimal solution is pi..- which can be proved by noting that Exent. (P. = (PT P(s)) +
Proposition 2.1 (Cross-entropy recovers P.j.)- The unique minimizer of too (p..e) is p.j. -P.. for every E support(pL).
2.2 SOFTMAX PARAMETRIZED LANGUAGE MODELING
Unlike traditional language models like n-gram models, neural language models parametrize the conditional distribution P.|a as softmax computed using low dimensional embeddings. For an embedding 0 E Rd,, the softmax distribution over W using word embeddings & E RdxV is o (w) e &w Zo, where Zo = Ew'ew e du" is the partition function. While po, depends on D. we will use po instead whenever & is clear from context. Just like Pi we can interpret
Pe E RV as a vector of probabilities for the distribution pa-
We now describe the abstraction for softmax models that is applicable most neural models. A language model first embeds a context & into f() E IRd using a feature map S + Rd that is parametrized by an architecture of choice (e.g. Transformer (Vaswani al. 2017)). The output conditional distribution is set to be the softmax distribution induced by the context embedding f(s)
and word embeddings 4, i.e. p.|a = Pf(x)- The cross-entropy in its familiar form is presented below
Exem (
E E S-PL
[- log(P)())(00)) = E SPPL
[-f(s) dhw) + log log(( (()) w>pij
E
(2)
wrppjj
We rewrite it as Exeml (f, 4)
E. [excat, (f(s), 4) where exemt
(0. (0. Ewent, (pe. +)) is the cross-
arpL
entropy loss for context s that uses embedding A. Analogous to Proposition 2.1 we would like to know the optimal d-dimensional feature map f and the induced conditional distribution pf
A finite minimizer may always exist. This handled in Section 4 that deals with e-optimal solutions.
3

---

Published as a conference paper at ICLR 2021
Proposition 2.2 (Softmax models recover P. a subspace). Fix fixed D, if f arg a (f +) exists, then (p)- = pp for every & E suppor((pL).
Unlike Proposition 2.1 (s) E RV is only guaranteed to be equal to P.a E RV on the d- dimensional subspace spanned by rows of & E Rdxv_ We may not learn Pie exactly when < V. but this result at least guarantees learning P. on linear subspace determined by word embeddings ②. This forms the basis for our main results later and is proved by using the first- order optimality condition, i.e. Volsent (f)
O, Ys E S. The gradient of cross-entropy is Valsem. = -... + VaZajza = - 1 + Ppp. Setting it to completes the proof. We use the properties of optimal solutions to understand why language models help with classification tasks.
3 USING LANGUAGE MODELS FOR CLASSIFICATION TASKS
Sections 2.1 and 2.2 suggest that language models aim to learn P. or a low-dimensional projection 'P' Thus to understand why language models help with downstream tasks, natural starting point is to understand how access to P. can help with downstream tasks. In thought experiment, we use oracle access to P. for any and demonstrate that sentence classification task can be solved by reformulating it as a sentence completion problem and using P. to get completions to predict the label. This sentence completion reformulation is mathematically formalized as natural tasks.
3.1 SENTENCE COMPLETION REFORMULATION
For exposition, we consider the sentence classification task of sentiment analysis, where the inputs movie reviews (subset of S) and labels belongs to {=1), denoting positive and negative reviews.
Classification task sentence completion: Can we predict the label for a movie review s by using Pi,.? One way is to use P.j. to compare probabilities of ":)" and ":(" following a movie review and to predict sentiment based on which is higher. This seems like reasonable strategy, since ":)" is likelier than ":(" follow positive movie review. One issue, however, that P.J. will place much higher probability on words that start sentences, like "The"", rather than discriminative words useful for the task. To allow a larger set of grammatically correct completions, we can append prompt like "This movic is " at the end of all movie reviews and query probabilities of indicative adjectives like good, bad, interesting, boring etc. that better indicators of sentiment. This approach of adding a prompt can also work for other classification tasks. For the AG news dataset (Zhang et al.) 2015) containing news articles from categories (world, science/tech., sports, business), prompt like "This article is about can help solve the task. The theoretical and practical relevance of prompts discussed in Theorem 4.T, and Section 6 respectively. We note that the choice of prompts and completion words is less important than the underlying idea of sentence completion reformulation and its formalization.
Solving tasks using a linear function of p: The above process is actually a sub-case of using a linear classifier on top of P/j. E VRV. For sentiment analysis, w4 = ":)" and w = ":(", then the sign of P.j. (w+) P. (w_) can predict the sentiment. This strategy can be expressed as v ''],,. where the linear classifier v E R has v(w+) 1, v(w-) = -1 and vww) = for w' E WW w- Similarly with the prompt, we can assign positive weights in to adjectives like "good" and negative weights adjectives like "boring". Strength of sentiment in different ad- jectives (e.g., "good" vs "amazing") can be captured through different weights. This equivalence between sentence completion reformulation and linear elassifier on P. further explored Sec-
tion D.I Other tasks can be similarly solved with a different set of words for each class. We verify experimentally that SST and AG news tasks be solved by a linear function of probabilities of just a small subset of words in Section 6 and for many other classification tasks in Section F.I thus lending credibility to the sentence completion view.
4

---

Published as a conference paper at ICLR 2021
3.2 NATURAL CLASSIFICATION TASKS
We now translate the above sentence completion reformulation into a reasonable mathematical char- acterization for classification tasks of interest. Firstly we formally define text classification tasks and the standard metric for performance of linear classification on fixed features. A binary classification piece of text from S and the label y is in {+1}. Given feature map g S - RD (arbitrary D), T is solved by fitting a linear classifier R on top of g(s) and the metric classification loss is
task T characterized by a distribution pT over S x {+1}, where the input
(g(g, v) [(.))>>T
[e(v T g(s),))] (G()) = inf (g,
(3)
vER
where E is I-Lipschitz surrogate the 0-1 loss, like the hinge loss {(.,)) (1- vū)+ or the logistic loss f( y) = log(1 - - For given embeddings {Os}eEs, the classification loss is written as e ((0.),, v) (. [e( 0..
We now formalize classification tasks amenable to sentence completion reformulation, from Sec- tion 3.1, as (T, -natuurall tasks, i.e. tasks that achieve a small classification loss of 7 by using a linear classifier with loe -nom bounded by B on top of features Pi E RV
min Definition 3.1. A classification task T is (T, -)nnturall if VER" 11
e (T(p)))))) ST.
While we motivated this formalization of linear classification over P. in Section 3.1 we provide a mathematical justification in Section D.I along with interpretations for and B that relate them to the Bayes optimal predictor and probability mass of indicative words respectively. Low dimensional softmax models, however, only learn P. in the subspace of 4, per Proposition 2.2 Thus also interested in subset of tasks that this subspace solve.
Definition 3.2. Task T (T, B) nnatural w.r.t E RdxV if
min VEnon-spand (1)l|l
e (r(1) ST.
Note that every (r, Bn-nttural task w.r.t. T is trivially (r, B)-ntuual,, though the converse may not hold. However it can be argued that if ① has some "nice properties", then (r. Bna-tturall tasks of interest will roughly also be (T, -ntttrrl w.r.t. $b Capturing the synonym structure of words can be such a nice property, as discussed in Section D.2 A better understanding of these properties of word embeddings can potentially enable better performance of language models on downstream tasks. In fact, Section 5.2 describes a carefully designed objective that can learn word embeddings with desirable properties like synonyms having similar embeddings. In the subsequent sections, we use the above formalization to show guarantees for language models on natural tasks.
4 GUARANTEES FOR LANGUAGE MODELS ON NATURAL TASKS
We now show guarantees for features from language models on natural tasks in two cases: 1) for an arbitrary language model {p. where use V-dimensional features P.|x E RRV for downstream tasks and 2) for softmax language model (f, o) where we use new d-dimensional features @p E Rd. Since cannot practically hope to learn the optimal solutions described in propositions 2.1
and and 2.2 we only assume that the language models are c-optimal in cross-entropy. We first define P to
be the minimum achievable cross-entropy and (en((()) ( be the minimum achievable cross-entropy
by a d-dimensional softmax language model using Q; clearly Ex I Exxen( (4).
Cx =
Esem ((
E
inf Exem, (0, b))
(4)
xent
smpL
BERA
We first present the results for arbitrary language models with a proof sketch that describes the main ideas, following which we present our main results for softmax language models.
4.1 ARBITARY LANGUAGE MODELS
We show guarantees for a language model that is c-optimal, i.e. Exem esem ((p.))) -tem I c, on (r, B)- natural tasks. An important consideration is that the language model distribution Pr, of contexts is
EEtending to k-way tasks is straightforward.
fo makes sense since ||"",_la & || ooo is dual norm of || I|1-
5

---

Published as a conference paper at ICLR 2021
often diverse superset of the downstream distribution py (defined in Section 2.2 thus requiring us to show how guarantees of p. = P.j. on average over the distribution $ ~ PL transfer to guarantees on a subset pr. In the worst case, all of the c error in cross-entropy by {p.,.) is incurred sentences from the subset pr, leading to pessimistic bound In practice, however, the erross might be more evenly distributed across PL. thus bypassing this worst case bound. As a first step, present the worst case bound here; stronger guarantees are in Section [5.1 The worst-case coefficient (PT), defined below, captures that PT is ((P)) )-fracion of PL.
sentences,
( - sup(> (0, PL (s)
(5)
We now present our results that applies to any language model, regardless of the parametrization (e.g., n-gram models, softmax models). The result suggests that small test cross-entropy (hence test perplexity) is desirable to guarantee good classification performance, thus formalizing the intuition that better language models will be more useful for downstream tasks.
Theorem 4.1. Let {p.ss} be a language model that is e-optimal, i.e. lsene((P.1s)) Eient S e, for
some 0. For a classification task T that is (T, B)-nttural, have
eT (1p.)t) ST V2BPe (T(P)))
This upper bounds classification loss on task T for V-dimensional features (p.1.) } from an e-optimal language model. We discuss factors that lead small upper bound and corresponding intuitions.
is small: learned language model has smaller cross-entropy pog-peele)))ty)
т is small: task can be solved well through a sentence completion reformulation with a set of
indicative words completions, in Section 3.1 and has small Bayes error (cf. Section D.I)
B is small: set of indicative words has high probability mass in P. (cf. Section D.I. This could
potentially explain the superior performance when prompts are added (Section 6
(PT) is large: pT is closer PL: note that 7(PT) S 1 with equality if and only if pT = PL
Thus the bound captures meaningful intuitions about good performance of language models on downstream tasks. We provide a detailed proof sketch in Section E. and strengthened version of this (Theorem I.I is presented in Section E.6 Proving this result requires connecting the clas- sification loss with language modeling cross-entropy loss and dealing with distribution mismatch; we present a rough outline to do so below. Since T is (r, Baanural,,, let v* be the classifier with ||o Imx S and ((P))))) 0²) S r. The result follows from the following inequalities:
ey ((p.)).0") - E (Tp u0)
E [( SupT
T (p.1. - -P.
Lipschitzness + Jensen's
E [(0 (p-10 -p. S ((P)) SupT
E
[10" (p.1. - p1
Transfer pT to Pr.
ampe
VU RY, (v (p.). - -.ip))))))}) < 2 (p lxeet,s (P.).))
Pinsker's inequality
The first and third inequalities (Lemma E.8 and Lemma E.3) connect the classification loss to the cross-entropy loss in language modeling, while the second inequality deals with distribution mis- match between Pr. and pr. We now present stronger result for softmax models.
4.2 SOFTMAX LANGUAGE MODEL WITH CONDITIONAL MEAN FEATURES
We now consider a softmax language model with feature map f that satisfies (xemi [f, 1b E (T) S suboptimality measured w.r.t. the best d-dimensional model, unlike Theorem 4.1 Note that Theorem 4.1 can be invoked here to give a bound of & (T(P)(()) < T + O(B natural tasks, where cT. = ft ( f is the suboptimality of the best d-dimensional model. The fixed error of O(B (ct.)) (even when c = 0), however, is undesirable. We improve on this by proving stronger result specifically for softmax models. Inspired by Proposition 2.2, our guarantees are for features ( E Idd called conditional mean features.
on (T, B)-
Definition 4.1 (Conditional Mean Features). For a feature map S Rdd and + E RdXVV
define conditional mean features hpp S - TRd, where Tpp((8) (p where Pf(s) E RRV
Forr instance pr is 0.001 fraction of pz. {rl1a} could have 1000 error on pT and 0 error on of PL-

---

Published as a conference paper at ICLR 2021
We now present the result for softmax language models that has similar implications Theorem 4.1 but with menoeneet subtle differences.
Theorem 4.2. For fixed D let f be features from c-optimal d-dimensional softmax language model, i.e. Exem( 1) - g (D) ()= S c. For a classification task T that is (T, B) B-natturall w.r.t. D,
e (r(p)) + (2B2e
This result guarantees good performance of conditional mean features bpy on some natural tasks, thereby suggesting novel way to extract features for downstream tasks. We empirically verify the good performance of p on classifications tasks (Section 6) and also find a Ol ((ve)-like behavior (Section The proof (Section E.3 is similar to that of Theoreme 4.1, the main difference being the use of the following inequality, proved using a softmax variant of Pinsker's inequality (Lemma E.4
Yv E row-span (D), (v (Pr(s) -
-p)))) Pi.))²² - 2| (ex.ntts(ppr
inf (PF-
f (*)ERR
The more general result (Theorem 5.1) replaces 7(P)) with a more refined coefficient (Section 5.1 While guarantees only for natural tasks w.r.t. $ Section [.aiiinnDD]discusseee why this might be enough for tasks of interest if word embeddings & satisfy nice properties.
4.3
(15(8) IS A LINEAR FUNCTION OF f())
Theorem 4.2 shows that bp is useful for linear classification. However, using feature map f directly is more standard and performs better in practice (Section 6). Here we argue that there is a linear relation between ff and h if word embeddings O satisfy a certain Gaussian-like property, which we show implies that tasks solvable linearly with Φpy are also solvable linearly using f.
Assumption 4.1. There exists symmetric Positive semidefinite matrix A E IR RRdxd vector b E Rd and a constant e E such that log( Zo) 2 9 TTAA A6 + b + effor any E RRd.
If word embeddings were distributed as Gaussians, i.c. V columns of & are sampled from N(u, )) independently. it not hard show (Lemma E. 1) that log( Zg) = _ 1 E9 + 8 H + log(V). While some papers (Arora et al.] 2016 Mu & Viswanath 2018) have noted that word embeddings are fairly random-like in the bulk to argue that the log partition function constant for ||0]|2 1, our quadratic assumption is bit stronger. However, empirically we find the fit to be very good, as evident in Figure Under the above assumption, we can show a linear relation between f and hpy- Lemma 4.3. Under Assumption 4.1 feature map ʃ satisfies ( Af(s) + b, Corollary 4.1. Under same setting as Lemma and
and Theorem 4.2 (()) S + O(B())
This shows that f itself is good for natural classification tasks. However, in practice, the linearity between f and bp, only weakly holds on features from pretrained GPT-2 (Radford et al. 2018)
(s) A5 bll"
The fractional residual norm of the best linear fit, i.e. r
measured for
[
different distributions (r = 0 is perfect fit) are 0.28 for SST, 0.39 for AG News, and 0.18 for IMDb contexts. This non-trivial linear relationship, although surprising, might not completely explain the success of f, which usually performs better than Opy; we leave exploring this to future work.
5
EXTENSIONS
5.1 BETTER HANDLING OF DISTRIBUTIONAL SHIFT
The bounds in the previous section use the coefficient 7(PT) to transfer guarantees from PL PT and we define a more refined notion of transferability here. The coefficient 7(PT) ) is independent of the learned model and assumes worst case distribution errors. For the refined coefficient, we first
define the error made in predicted probabilities by a softmax language model f as A (prr))) =
Pf(a) Piss- For any distribution E As, we define uncentered covariance of a function g S
D as Ep(g) Egpp [g(s)g(s)]] The refined transferability coefficient is then defined as
70s ppr) = (1 1pppe (AA()))))) -tΣ, (pPr))))) Epe (9 (Pprc)))) |1)
7

---

Published as a conference paper at ICLR 2021
Thue tog partition versus learned quadrani tunction
Wel
302
50
202
-50
1200
100
100
-50
5
Tue pertition
Thue log partitin
Tre kog partition
Figure I: Learned quadratic function v/s log partition function on various datasets for features com- puted from pre-trained GPT-2 to verify Assumption 4.1 We also plot the y x line for reference.
We state the refined result for softmax language models; detailed results are deferred to Section B
Theorem 5.1 (Simplified). In the same setting as Theorem 4.2 ( S T
2. T(Pribpy)
It is easy show that (ppT; 2 7(PT), so this is indeed stronger bound. The coefficient ((P:; ppy) measures how average error on f on pt can propagate to PT. This can potentially be much smaller than 7(Pr) due to some inductive biases of f. For instance, if errors made by p. independently of s, then Ep,. (P))))))))) z Ep( (AA (prcn)) R Ex--o[nnT] making r(p; bpy) = 1. Independence prevents accumulation of language modeling error contexts from pT. bypassing the worst case transfer of ((PT).
the model are random-like, i.e. (prrc)
5.2 QUAD: A NEW OBJECTIVE FUNCTION
In Definition 3.2 we discuss how low dimensional softmax language models learn a linear projection of . only solving tasks that lie in the row span of word embeddings D. Although & defines tasks that language model features
solve, the standard cross-entropy objective does not lend simple closed form expression for optimal $. This motivates the construction of our Quad objective, that has two nice properties: (1) the optimal feature map f" is a linear function of p., and thus solve some natural tasks, and (2) the optimal Φ" has an intuitively meaningful closed- form solution.
1 1 [())]?
(Juad(J $)
E So-pt
(-f(s) sw)
(6)
The Quad objective is very similar the cross-entropy objective from Equation 2 with the log partition function replaced by quadratic function, inspired in part by Assumption 4.1 We can derive the optimal solution &² that depends on the eigen-decomposition of a substitutability matrix.
Definition 5.1.
The substitutability matrix is defined to be Ω = E [Pio RVVV If
srpp
Ω² = U is the eigendecomposition, then UA E VVxx xd is matrix of top d eigenvectors of
The matrix Ω captures substitutability between pairs of words. Words w and w' are substitutable if they have identical conditional probabilities for every context E S and thus can replace occur- rences of each other while still providing meaningful completions. By definition, these words satisfy [w] 22" [w']. Such pairs of words were called "free variants" in the work on distributional se- mantics (Harris 1954), and capture the notion of synonyms; more Section D.22
Theorem 5.2. Let f" x arg min j. (quaa([; Q). Then xx = BUJ, for full rank B E Rdxdd Also, for a classification task TT that is (T, B)-naturall w.rl. th", we have (
Thus f excels on natural tasks w.r.t. 4², which in turn, is the best d-dimensional projection of Ω Thus words w. w' W that are synonyms (hence substitutable) will satisfy oū = Di fulfilling the desired property for word embeddings discussed in Definition 3.2
We train using the Quad objective and compare its performance to a similarly trained GPT-2 lan- guage model. The results in Table 3 suggest that Quad performs comparably Φp, from the cross-entropy objective, which fits our theory since both are linear functions of P. Section F3 has more details and experiments. The goal of testing Quad is to demonstrate that theoretical msights can aid the design of provably effective algorithms. Refer Section C for more details on Quad.
8

---

Published as a conference paper at ICLR 2021
Table Accuracy (%) on k-way linear classification using fixed GPT-2 features. Good performance of features fs), conditional mean features Φpy(s) and meaningful subset of I 30 (and I 2k) coordinates of P)(*) verify the sentence completion reformulation and main results. The numbers right below the features denote dimensionality of the features. An asterisk indicates that we added a task-specific prompt. Other baselines are fine-tuning (FT, Section [2) and random projection of
(f()) (rand. proj.). Sentence version of SST (train/test: 6.9K/1 (K) is used.
[(s)
ps)(s) 768
Pf(s) (subset)
Pf()) (class words)
Pf(e) (rand. proj.)
Task
k
FT
768
I 30
S
768
SST
2
87.5
83.3
82.6
78.7
67.5
91.4
SST#
2
89.4
87.3
85.4
79.1
76.4
SST fine
5
49.2
43.5
44.0
39.2
23.1
50.2
SST fine
5
49.4
48.6
47.6
40.3
28.8
53.5
AG
4
90.7
84.6
83.8
75,4
58.5
94.5
AG
91.1
88.2
86.1
75.1
63.7
94.4
EXPERIMENTS
We use experiments to verify (1) linear classification on fixed language model features does compa- rably to fine-tuning the features, (2) sentence completion reformulation (Section (.1),, i.c. tasks can be solved using probabilities for indicative words, (3) conditional mean features are effective.
Tasks using linear function of p We validate our claims from Section 3 that classification tasks be solved by linear functions of Pij.- Since P, is never available, instead use the output features f()) and probabilities p.ls := (s) from a small pretrained GPT-2 model (Radford et al. 2019). Table T demonstrates that on binary and fine-grained Stanford Sentiment Treebank (SST) (Socher et al. 2013) and AG News (Zhang et al., 2015) tasks, probabilities Pf(s) of just 30 or so task-relevant tokens (see Section
solve the tasks. Even just one/two token per class ("class words") yields non-trivial performance. Furthermore, we validate the sentence completion reformulation in Section 3.1 by using the probabilities Ps(a) after adding task specific prompt and consistently observing improved performance, including for fine-tuning (FT) with small datasets.
Φpy and f are good features: We first that linear classification over fixed features (() from the pretrained model performs comparably to the FT baseline. We further validate Theorem 4.2 by verifying that the conditional mean features bp (s) also linearly solve downstream tasks fairly well. This performance is comparable to, but always worse than f(s), as seen in columns 3 and 4 of Table We again find that adding prompt improves performance. Note that random projection of (s()) same dimensions as (s) has very poor performance. Section E.S has results for wider range of classification tasks. Evidence for Assumption 4.1 is provided by learning a quadratic function to
fit the log partition function of features from pretrained GPT-2 model (see Section F.4p. Figure demonstrates that the fit holds for its training and unseen data (e.g.. WebText (Radford et al. 2019
7 CONCLUSIONS AND FUTURE WORK
We provide intuitive and mathematical explanations for the success of language model features classification tasks by reformulating them as sentence completion problems. This reformulation formalized as natural tasks: those that can be solved linearly using the conditional probability distribution . Insights from our analysis help design the Quad objective that provably learns good features for these natural tasks. We hope our analysis will inspire other mathematical insights into language models. While Section 4.3 argues linearity between conditional mean features Ppy and f, it is insufficient to explain the observed superiority of over hp,. We leave exploring this limitation of our analysis to future work. Guarantees for softmax models are for natural tasks w.r.t.
4, thus knowing the optimal d-dimensional word embeddings x for (f 4) is also important. Other meaningful directions include providing guarantees for other successful models like BERT (Devlin et al. 2018) and more diverse downstream tasks. Although
would like to show stronger guarantees by exploiting model and algorithmic inductive biases, as well as study the setting of fine- tuning language model features, lack of a good theory of deep learning the current bottleneck.
Acknowledgments: Sanjeev Arora, Sadhika Malladi and Nikunj Saunshi are supported by NSF, ONR, Simons Foundation, Amazon Research, DARPA and SRC.

---

Published as a conference paper at ICLR 2021
REFERENCES
Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. A latent variable model approach to pmi-based word embeddings. Transactions of the Association for Computational Lingnistics, 2016.
Sanjeev Arora, Yingyu Liang, and Tengyu Ma. A simple but tough-to-beat baseline for sentence embeddings. In Proceedings of the International Conference on Learning Representattonss 2017.
Sanjeev Arora, Mikhail Khodak, Nikunj Saunshi, and Kiran Vodrahalli. A compressed sensing view of unsupervised text embeddings, bag-of-n-grand and LSTMs. In Proceedings of the Interna- tional Conference on Learning Representations, 2018.
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. In Proceedings of the 36th International Conference on Machine Learning, 2019.
Sören Auer, Christian Bizer, Georgi Kobilarov, Jens Lehmann, Richard Cyganiak, and Zachary Ives. Dbpedia: A nucleus for web of open data. In Proceedings of the 6th International The Semantic Web and 2nd Asian Conference on Asian Semantic Web Conference, 2007.
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin. A neural probabilistic language model. Journal of machine learning research, 2003.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 2020.
Stanley F Chen and Joshua Goodman. An empirical study of smoothing techniques for language modeling. Computer Speech & Language, 13, 1999.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint wrXiv:1110004005, 2018.
Chris Dyer. Notes on noise contrastive estimation and negative sampling. arXiv preprint 1rX1v:1410.8251, 2014.
John R Firth. A synopsis of linguistic theory, 1930-195 1955. Studies in linguistic analysis, 1957.
Zellig Harris. Distributional structure. Word, 1954.
Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification.
arXiv preprint arXiv: 18880006466 2018
Minqing Hu and Bing Liu. Mining and summarizing customer reviews. In Proceedings of the Tenth ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, 2004.
Mikhail Khodak, Nikunj Saunshi, Yingyu Liang, Tengyu Ma, Brandon Stewart, and Sanjeev Arora. A la carte embedding: Cheap but effective induction of semantic feature vectors. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguisties (Volume 1: Long Papers), 2018
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.69880 69800, 2014.
Ryan Kiros, Yukun Zhu, Russ R Salakhutdinov, Richard Zemel, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Skip-thought vectors. In Advances in neural information processing systems, 2015.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sori- cut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: /9942 2019.
Jason D. Lee, Qi Lei, Nikunj Saunshi, and Jiacheng Zhuo. Predicting what you already know helps: provable self-supervised learning. arXiv preprint arXiv:2008. 01064, 2020.
10

---

Published as a conference paper at ICLR 2021
Omer Levy and Yoav Goldberg. Neural word embedding implicit matrix factorization. In Ad- vances neural information processing systems, 2014.
Xin Li and Dan Roth. Learning question classifiers. In Proceedings of the 19th international conference on Computational linguistics- Volume 1, 2002.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, arXiv preprint orxiv: 1692, 2019.
Lajanugen Logeswaran and Honglak Lee. An efficient framework for learning sentence representa- tions. In Proceedings of the International Conference on Learning Representations, 2018.
Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for conditional models: Consistency and statistical efficiency. In Proceedings of the Conference Empirical Methods in Natural Language Processing, 2018.
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the ACL: Human Language Technologies, 2011.
Julian J. McAuley, Rahul Pandey, and Jure Leskovec. Inferring networks of substitutable and com- plementary products. CoRR, 2015.
Bryan McCann, James Bradbury, Caiming Xiong, and Richard Socher. Learned in translation: Contextualized word vectors. In Advances in Neural Information Processing Systems, 2017.
Bryan McCann, Nitish Shirish Keskar, Caiming Xiong, and Richard Socher. The natural language decathlon: Multitask learning as question answering. arXiv preprint 0rXiv:/806.087.0 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 1eXXvv1609.07880 2016.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word represen- tations in vector space. arXiv preprint arXiv: 2013a.
Tomas Mikolov, Ilya Sutskever, Kai Chen, Greg Corrado, and Jeff Dean. Distributed represen- tations of words and phrases and their compositionality. In Advances in neural information pro- cessing systems, 2013b.
Jiaqi Mu and Pramod Viswanath. All-but-the-top: Simple and effective postprocessing for word representations. In Proceedings of the International Conference on Learning Representations, 2018.
Matteo Pagliardini, Prakhar Gupta, and Martin Jaggi. Unsupervised learning of sentence embed- dings using compositional n-gram features. Proccedings of the North American Chapter of the ACL: Human Language Technologies, 2018
Jeffrey Pennington, Richard Socher, and Christopher D Manning. Glove: Global vectors for word representation. In Proceedings of the 2014 conference on empirical methods in natural language processing (EMNLP), 2014.
Matthew E Peters, Mark Neumann, Mohit lyyer, Matt Gardner, Christopher Clark, Kenton Lee, and Luke Zettlemoyer. Deep contextualized word representations, arXiv preprint 2018.
Raul Puri and Bryan Catanzaro. Zero-shot text classification with generative language models. arXiv prepring 1912. 2019.
Alec Radford, Rafal Jozefowicz, and Ilya Sutskever. Learning to generate reviews and discovering sentiment. arXiv preprint a704v 2017.

---

Published as a conference paper at ICLR 2021
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language
understanding by generative pre-training 2018. URL /ttps://SS-s--wett2
am.zzonaw.ccom wwt
Larguagguurderrstanddnnnpppprp
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl Blog. 2019.
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. arXiv preprint arXiv: 10683, 2019.
Prajit Ramachandran, Peter Liu, and Quoe Le. Unsupervised pretraining for sequence to sequence learning. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, 2017.
Timo Schick and Hinrich Schütze. It's not just size that matters: Small language models are also few-shot learners. arXiv preprint #xXi:2009.00///8 2020.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over sentiment treebank. In Proceedings of the 2013 conference on empirical methods in natural language pro- cessing, 2013.
Alon Talmor, Yanai Elazar, Yoav Goldberg, and Jonathan Berant. olmpics-on what language model pre-training captures. arXiv preprint arXiv: 2019.
Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive learning, multi-view redun- dancy, and linear models. arXiv preprint arXiv:2008. 10150, 2020a.
Christopher Tosh, Akshay Krishnamurthy, and Daniel Hsu. Contrastive estimation reveals topic posterior information to linear models. arXiv preprint arXiv:2003.02234 2020b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1804.08800.07440 2018
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity the hypersphere. arXiv preprint arXiv:2005. /0242, 2020.
Theresa Wilson and Janyce Wiebe. Annotating opinions in the world press. In Proceedings of the Fourth SIGdial Workshop of Discourse and Dialogue, 2003.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierrie Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface's trans- formers: State-of-the-an natural language processing. arXiv preprint 191003791, 2019.
Wei Xu and Alex Rudnicky. Can artificial neural networks learn language models? In Sixth inter- national conference on spoken language processing. 2000.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in neural information processing systems, 2019.
Xiang Zhang. Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. In Advances in Neural Information Processing Systems 28, 2015.
12

---

Published as a conference paper at ICLR 2021
A
OVERVIEW
Section
B
is a more detailed version Section 5. and Section detailed version of Section 5.2
Section D.1
has a discussion about why natural tasks are reasonable formalization for the sentence
completion reformulation and also interpretations for T and in the definition natural tasks. discusses desirable properties of word embeddings O like capturing synonym structure in words. Section contains proofs for all results, including proof sketches for the main results in
Section D.2
Section E. Lemma E,4 is the softmax variant of Pinsker's inequality that prove and use for
main results.
Section DectionFccontiinn contains many experimental findings that consolidate many of our theoretical results. Section provides the information about subsets of words used for results in Table [ and also additional experiments test the performance of pretrained language model embeddings f more downstream tasks and also verifying that conditional mean embeddings Φp, do well on these tasks. In Section F3 we present additional results for Quad objective trained on larger corpus and tested on SST. Section|E./provides additional details on how A, b and from Assumption 4.1 are learned and also further verification of the assumption on more datasets. Finally, Section E55 experimentally verifies the ( dependence from Theorem 4.2
B BETTER HANDLING OF DISTRIBUTIONAL SHIFT
While the bounds above used 7(P)) to transfer from the distribution pr to pT. we define a more refined notion of transferability here. While 7(P) only depends pL and pT, the more fined notions depend also on the learned language model, thus potentially exploiting some in- ductive biases. We first define the notion of error made in the predicted probabilities by any predictor a. as A (8) = P.s= Pij.- Thus for any softmax language model f we have A (PrO)) (s) Pf(x) P.j. For any distribution p E As, we define the covariance of func- [g(8)g(8) We define 3 coefficients for the results to follow tion g S + IR Ep(g) sup
Definition B.1. For any distribution p E As, we define the following
70: (p)))) = (1) Epp (A E EpL (A (p: (p.1.)) = (1/Epe (D
(A(P
112
(7)
"E ( (p)
(8)
70; 4ps) = (p: (P)())))
(9)
We notice that (p(A(r.))) = E Ppp)(pp - (p(DA(p)))
srup
Φ E A (D. a) TT We are now ready to state the most general results.
Theorem B.1 (Strengthened Theorem 4.1). Ler {p.ia} be language model that is c-optimal, i.e.
Evou (p.1a) f < for some > O. For a classification lask T that is (T, B)-natural, we have
2B2
l (Ip ST+
7(P: (p)
For classification task T that is (r, B)-natural w.r.f. T we have
2
e ((p eT ((Dp. ST+
No (pr: (ppp))
Theorem 5.1 (Strengthened Theorem 4.2 For a fixed & let be features from an e-optimal d- dimensional softmax language model, i.e. Exemlf (f 1) -Ee
- (Φ) S where ( (④) is defined in Equation For a classification task T that is (T, -)nntrrla w.r.t. D we have
2B
èT (1P)a)) s ey (r(p)) ST+
7PT: 2py)
Thiis is not exactly the covariance since the mean is not subtracted, all results hold even for the usual covariance.
13

---

Published as a conference paper at ICLR 2021
Discussions: is not hard to show that the coefficients satisfy 74 (Pr: (p.js)) 2 7(P:: (p.1s)) 2 ((P)) and y(p::(pp))) 2 ((T)) ). thus showing that these results are strictly stronger than the ones from the previous section. The transferability coefficient a measure of how guarantees on Pt using a language model can be transferred to another distribution of contexts and it only depends on the distribution of contexts and not the labels. Unlike ((P)), the coefficients in Definition B.1 depend on the learned models, either {p. or (P)()),, and can be potentially much smaller due to the inductive bias of the learned models. For instance, if errors made by the model are random-like i.e. ((p. .) .3))) p. independently of s, then Ep, (p..)) Ep(A(p.i.)) Egy= J. making y(p (x}) 1. Independence prevents language modcling error from accumulating on contexts from pT. bypassing the worst case transfer of (PTT).
C QUAD: A NEW OBJECTIVE FUNCTION
In Definition 3.2 we discuss how low dimensional softmax language models learn a linear projection of p. only solving tasks that lie in the row span of word embeddings Φ. Although & defines tasks that language model features can solve, the standard cross-entropy objective does not lend a simple closed form expression for optimal 4. This motivates the construction of our Quad objective, that has two nice properties: (1) the optimal feature map f is a linear function of p., and thus can solve some natural tasks, and (2) the optimal कh has an intuitively meaningful closed-fon -frm solution.
1 JopTon 011 = - "0ppip + +1 T012
Equad, (0, 2) = E [-0Towl
1
(10)
wrpij
Equad (S, Q) = E
puudd (f(),, +)]
(11)
ampL
The Quad objective is very similar to the cross-entropy objective from Equation (2) with the log
partition function replaced by quadratic function, inspired in part by Assumption 4.1 We can
derive the optimal solution 46 that depends the digen-decomposition of substitutability matrix.
Definition 5.1
The substitutability matrix is defined to be Ω = E
[P]. Pi.
7 ERRXY | If
srpp
Ω = US is the eigendecomposition, then Ua E RVVxd is matrix of top d eigenvectors of
The matrix O² captures substitutability between pairs of words. Words and w' are substitutable they have identical conditional probabilities for every context s E S and thus can replace occur- rences of each other while still providing meaningful completions. By definition, these words satisfy ["w]] = _2'(w'). Such pairs of words were called "free variants" in the work on distributional se- manties (Harris 1954), and capture the notion of synonyms in the distributional hypothesis. We now derive expressions for the optimal solution of the Quad objective described in Equation The proof of all results from this section in Section E.5
Theorem C.1. The optimal solution f*, 4' arg min, o Equad 2) satisfies
4- BUT full rank B E Rddd
f (s) = (4 n T p. = CUJ P.ja- full rank C E Rdxd
If de is fixed, then the optimal solution is $ (*) = ( -
Theorem 5.2 Let f" x- arg min j. Equad(S (). Then x- BUJ, full rank B e Rdxd Also, for classification task 'T that is (T, B)-naturall w.r.t. Φ², have b (J*) ST.
Thus f excels on natural tasks w.r.t. D², which turn, is the best d-dimensional projection of Q* Thus words w, w' E W that are synonyms (hence substitutable) will satisfy の = oi,,. fulfilling the desired property for word embeddings discussed in Definition 3.2 We train using the Quad objective and compare its performance
similarly trained language model, finding Quad to be reasonably effective. The goal of testing Quad not to obtain state-of-the art results, but to demonstrate that theoretical insights can aid the design of provably effective algorithms.
14

---

Published as a conference paper at ICLR 2021
D MORE ON NATURAL TASKS
The discussions in this section may not be formal and precise in places, they are meant to provide more intuition for some of the definitions and results.
D.1 SENTENCE COMPLETION REFORMULATION = NATURAL TASK
We provide informal justification for why the sentence completion reformulation can be formalized as being able to solve using linear classifier over p E RV. The analysis will also end up providing some intuitions for т and B in Definition 3.T and Theorem [41 In particular, show that task that is amenable to the sentence completion reformulation will be (T, -ntturll, with ((Bayes-Error((T) i.c. т is small if the Bayes error for the task error, and B ((a(Miiicatve)) ) is inversely proportional to the probability mass of the set of indicative words for the task. This is formalized in Proposition D.2
will
Linear classifier over P.
Consider a binary classification task T and that can be solved with a sentence completion refor- mulation after adding prompt in Section 3.1 for e.g. sentiment classification
be solved by adding a prompt "This movie is" at the end of every movie review and use the completions to solve the task. Recall that pT is the distribution over S x {+1} for the task T. We abuse notation and use PT to denote the distribution over inputs where a prompt is added to each to input, for e.g. "I loved the movie." is transformed to "I loved the movie. This movie is". For any pT. let (y = 1[s) and pT (y = -1[s) denote the conditional probabilities of the sentiment of review (with an added prompt) being positive and negative respectively. By law of total probability we can write this conditional probability as
pT = 1|s)=
Σ WEW
Pr(y = (((,s ()) = Σ Pr w)) pT.(w)
(12)
WEW
For any task T we can roughly partition the vocabulary set W into the following
Indicative words Windicative: w can be an indicative completion for the task, like "good", "boring", "trash" etc, after movie review like ="I loved the movie. This movie is". In this case the sentence completion reformulation can be interpreted as the following: the completion w after a review s is sufficient content of the review s to predict the label if we know the completion w. This can be formalized as Pr(y = [| w)) R Ply 1[w) for some fixed distribution P for indicative completions w.
determine the sentiment of the review, i.c. we do not need to know the
Irrelevant words Wireetvant w can be an irrelevant completion for the task, like "a", "very", "not". In this case the completions, on the other hand, do not reveal anything more about the sentiment for the review than itself, i.e. Pr(y = 1[(s, w)) R PT (y = 1|s) for irelevant completions
Thus from Equation
12
we get
PT (
1|s)
Σ
Pr(y = 1|(s, w)) p.j. (.j.(w)+
Σ
Pr(y (ss,, w)) (w)
we Wastariii
Σ
Ply = 1/w) p.js(w)
Σ
pT pT = 111)) (1]((w)
WE Wantiaiio
WE Wimntooe
Σ
v1(10) p.pp (w)
(w) + Py - 1|s)
Σ
(1p(w)
W NEEWase WaWaiisae
w WiWneeaan
2 P/,, + pr 1.s)p. (1 (imelevan))
where U1 E RV is defined as t2 (w) = P(y 1/w) for w E Winicctive and U1 (w) 0 for w W Wirelvaant- Similarly we can define 1 E R with v_l(w) = P(y = -1/w) for w E W Iidicatives 0- (w) = for w E Wimelevant- From the earlier calculation, and similar one for y - 1, we get
1
1
PT (3) b(s) 1 (M (Wirelevant)
2. p.
2. p... for b {=1}
p. (W,,(WWiee (Mndicativ))
15

---

Published as a conference paper at ICLR 2021
If assume P (Vindicative) = ((Windicative)) is roughly the same for all s, ie. probability mass of indicative words following modified review is approximately the same, then we get
pr = 1|s) PT(y = 1s)) = pyp. where by a a(Windvcaive))
(0) v--) (13)
Thus we can approximately express the difference in conditional probabilities of the 2 classes as linear function of p.. While it intuitively clear why knowing pr 1[s) pr (y 1|s) is useful for solving the task, we show precisely why the next part.
Interpretation for and B
Based on the above discussed, we will show that the task T from earlier is (r, B)-natural according to the Definition 3.T and will also give us an interpretation for and B. First we show that the following predictor from Equation (13) is effective for task T
gT (T( = pr 1|s) pr 1[s) v.p'ls
(14)
We reuse the notation from Equation
3 and define the task loss for any predictor g S - IR as
(g) ((.))~T (((g(s),)))
(15)
Furthermore let Bayes-Error(?? := inf g:S- (p [1{g()) # y denote the Bayes error of the task T, i.e the optimal 1 error achievable on the task.
Proposition D.1. For any task T and for the hinge loss e, Er(9T) < Bayes-Error(T), where @r (s) = pr (y = 1|s) - pr = 1|s).
Thus if task is easily solvable, i.e. has small Bayes error, then it will be solvable by the predictor @r (s). Since we argued above that sentence reformulation implies that gr(s) is a linear function of Pl.., now show that T is natural task formalized in Definition 3.1
Proposition D.2 (Informal). Task T that can be reformulated as sentence completion task (de- scribed above) is a (T, B) natuural task w.r.l. the hinge loss, with the follow parameters
S 4 Bayes-Error(T) and B Here Bayes-Error(T) is the Bayes error of task T and a (Vinticative) is the total mass of the indica- tive words for the task.
If the task T can be reformulated sentence completion, then T is (T, -atturaal where
T is small the task is unambiguous, i.e. it has small Bayes error
B is small if the probability mass of the set of indicative words Windicative is large, i.e. the task depends on a large set of frequent words
Thus the upper bound in Theorem 4.1 is smaller if the task can be reformulated as sentence comple- tion task with a large and frequent set of completions, and we can ever hope to solve it well (Bayes error small). The proofs for the above propositions in Section D.I
D.2 NICE PROPERTIES OF WORD EMBEDDINGS
We argue here that if the word embeddings O satisfy certain nice properties, then (r, B)-nuturall tasks of interest will be (r', n'aturtural w.r.t. b, where we will provide informal quantifications for the nice properties and tasks of interest that lead to a small value for r' and B'. The nice property will be related to & capturing the semantic meaning (synonym structure) of words and tasks of interest will be those that try to distinguish word completion (in the sentence completion reformulation) with very different meanings, i.c. trics to distinguish more coarse-grained semantic notions rather than very fine-grained ones. Note that the results here are informal and qualitative, rather than quantitative.
Consider task T that is (T, nattural task and let 8" E RV be the classifier such that (T( }, v²) S J and ||o*|s< 1100 < B. We want to find properties of and v" that will make T to be (7', natural w.r.t. such that 7' and B' are not too large
'Note that the converse is trivially trie, i.c. (r, naanaturall task w.r.t. T is also (r, Ba-nutural.
16

---

Published as a conference paper at ICLR 2021
We will show that T is (r', B3) -natural w.r.t & by finding classifier v such that o E RV, ||olloo I B' and ((P: ). v) I T'. First we define Ps. == Φ E RVV to be the projection matrix for the row-span of & and P lv Ps be orthogonal projection matrix. We will show that the classifier v = Pa suffices for our case, under some intuitive conditions on o* and Φ.
To compute B', we first look at the Loe norm of P40"
B' = | = P| cc = l|0 - Ps l S ||00 + +| 001 S
B | U2
To find the upper bound r', we upper bound the classification loss of Pa. substitutability matrix 12 smp
We first define the Then
E
Pi.Pi
T
similar the in Definition 5.1
a( ((o'pi(::))) pis:
[(((P>>") v)
E (8.))pT
= E (8.0)~p7
<(a) E
e(00 [ Di
+ E 11(0" P P
(ew)rpp
SNpT
((p(pp))) + E
( P
SupT
<10) ++ V enupT
(w (on P =T+
E suppr
Ps p.jp..
T "P v
100) T + V V T
Po pf ((1) T |P v
P 2 Po1l2
where (a) follows from 1-Lipschitz property of l, (b) from Jensen's inequality and that (TP S T. (c) from the definition of substitutability matrix Spt and (d) by definition of spectral norm of a symmetric PSD matrix.
Thus we have shown that T is (r', )nnauuuural w.r.t ④ where
= P v² llav P Pp B'= B |P
(16)
We will now show that if & captures the notion of synonyms, then P P will be small leading to r' being small. Furthermore we also shed some light on what means for P| o*112 to be small, which will in turn make B' small and r' smaller. We do so with the following arguments, 1) 22 captures semantic meaning of words and thus its top eigen-directions will capture more dominant semantic concepts, 2) if & captures the "top-d" directions of meaning, i.e. the top-d eigen-directions of Opp then 12, Ps 2l (1/d), 3) additionally v" cares about the "top- d' directions of meaning, i.e. top-d eigen-directions of Ω then || o" ||2 will be small. We expand on these points below
1. Substitutability matrix (22;)) captures semantic meaning: We use a similar argument to the one in Section 5.2 right after Definition 5.1 that is based on distributional semantics (Harris) 1954). Harris 1954) posits that meaning for elements (words) can be derived from the environments (contexts) in which they occur. Thus Harris (1954) argues that words that occur in almost identical set of contexts have the same meaning, i.e. are synonyms. On the other hand, if two words share some contexts but not all, then they have different meanings and the amount of difference in meaning roughly corresponds to amount of difference in contexts. In our setting, the similarity of words w and w' can then be determined by the probabilities assigned to them by different contexts s. particular, if p. (w) P.,((v' (w') for all or most supp(p7), then and w/ have essentially the same meaning w.r.t. the distribution of contexts pr and the closer [p', (w))sesmp((P)) and [p]], (w'))==spppi))) sespp((p)) are, E [p],P.j. TJ (w)) VVV is not hard show that O7, (w) 2pr (w') is equivalent to (w)) (1''(" Vs PT. where 22 (w) is the row of 22, corresponding to word w. To show this, we can define Bi, e R[ to be embedding of w that looks like B.n = (p), / (s))((())))))))) (#))sesup((r)) It is easy to P.j. (wi (w.)p, = 227 (W1, W2). Thus B... = Bw = O2 (w) = SNuPT
the closer the meaning of w and w' are. For the substitutability matrix 2 SypT
see that BT, Bwr
E.
17

---

Published as a conference paper at ICLR 2021
Pp (w") is straightforward to see. For the converse,
P (w) = Rip (w') - Ppr (w, w) - P (w', = 2 (w, w') = Ppr (w', B. B.. Bur Bun Bio Bi
(17)
(18)
Thus 2 indeed does capture the synonyms structure between words, and the top eigen-directions of it capture the most significant "semantic meaning" directions.
2. + has nice properties: if o roughly respects this synonym structure by aligning with the top-d eigen-directions of Ω Spp, have
Ps 22 P+ll 112 S X4+1
Σ
X2 pT
)5
d tr(2pp)
(19)
d+1-
E + 1 eMpT
[(Pij.Pi.
(20)
d ++
From Equation (16 we then have 7' ST
1 vd
3. Tasks of interest: more likely for a classifier ex to separate words with big differences in meaning rather than small differences. For e.g., it is more likely for a task to separate word completions "good" and "bad" rather than "good" and "nice". Since top eigen-directions of S2 capture more dominant semantic meanings, this could correspond to v* aligning with the top eigen- directions of p In combination with the above property about b this could suggest that P| v" ||2 small, thus leading r' and B' being small.
Note that they above arguments are informal and qualitative, and we leave exploring desirable prop- erties of & more formally future work.
D.3 PROOFS FOR SECTION D.I
Proposition D. Let po(s) Pr(y = bls) for b E {+1}, Pmin(s) minse (+ Pb(s), Pmoz (s) (ax (8( and (s) arg max Pb(s) denote the Bayes optimal predictor. We first. notice that there is simple well-known closed form expression for the Bayes risk
Bayes-Error(T)
[1 (g) øy]
(5.2)~p7
1 arg max po + 1 bE(+1}
E (a.0)~P
E pupp
[Pman (8))
We now analyze the hinge loss of the predictor 9pr defined in Equation 14 Note that since Gpr (s) < 1, the hinge loss E(gpT (s). = (1 ydpr (s))+ = 1 - ygpr (s) for every 8, Thus the total loss is
Jpr (s)
E (P
[11 V9pr (s))+]
E
[[1 y9pr (8))
(8.0)~pT
(a)
pp (s) (1 ggpr(s) + (s) (1 + 9or (s))
E [1 (pi(s) p-1(s))gpr()) (s)]
S-PT
11 (pi(s) p=1())) = E
+ (p) +p-1(s)) (pi(s) ps)
sppT
S
E [4p (4pp(s)p-1(s) 4 E. [pmán (5) (s)Pmax(s) s~pT
4 sppT
Pinin (s)] Bayes-Error(T)
where (a) follows by splitting the expectation over y/s, (b) follows from the definition of 9pT (s) in
Equation (14 and (c) follows from Pmax (8) < 1. This completes the proof.
18

---

Published as a conference paper at ICLR 2021
Proposition D.2
Let B
a: Minicaative )
We first note the following using the definition of
from Equation 13
max U1 (w) u-1((w) = B max [P(y | a(Windicative)
1/w) P(y -1/ww) S B
wEW
(21)
To find the value of r that makes the task (T, B))-tturall (Definition 3.1). we observe the following
min vER".llolssB
& (T(P v) =(a) (1p%),))))) (..5)~'
(b) E [((T(())))]] (g(T)
(s.))~T
<(c) Bayes-Error(T)
where (a) follows from the calculation in Equation follows from Proposition D.I
21 (b) follows from Equation 13 and (c)
E PROOFS
E.
PROOF SKETCH
We first present a sketch of the arguments that help show our main results, theorems 4.1 and 4.2
The subsections after the next one contain the full proof' for strengthened versions of these results.
E.I.1 PROOF SKETCH FOR ARBITRARY LANGUAGE MODELS: THEOREM
Here we want to show guarantees for features {p.,)) on a (r, B) -nauura task T. From the definition of natural tasks, know
FF E
|000 S s.t. ((pp)p),
(22)
We wish to upper bound the classification error e, ((p.)11))) and do so using the following sequence of inequalities.
((p = inf (T((P)u).)) - l (1.1 e (TP.pjtt VERV
(T((.1))) 0*) - -T((P)
E
[(v
(p/1s - PP1,))3]
SNPT
E [10 (p.je P.j
E [[v* T (p.1s - -P%1))3
[(v" (p.1. D. [)3]
V
sppt
1 sNpT
8NpL
l v²) - v*)
w Epr (A (A(p))))0*
E supL
[(v (p.)p P.).)) (23)
Vor Expt (A (
U+ Epz. (A (p.
Q1 (8))
(2) covariance from por - pt. Use tramsferability coellicient
O33 (x*)
covariance cross entropy
Classification loss covariance Use Lipschitzness and Jensen's insquality
Use (modified) Pinsker's inequality
where Ep(g)
E [s( is the uncentered covariance of g w.r.t. distribution p s~p
As. as defined in Section 5.1 We upper bound (r([p.s)) - by upper bounding each of a1 (v*), (2(1"), (g as Tollows
Classification loss + prediction error covariance: αi (x=" upper bounded by using Lipschitzness of the loss used in the definition of ly, e.g. hinge loss or logistic loss, and then followed by an application of Jensen's 's inequality
Lemma E.8
01 (v) < for all E RV
Error covariance from PT PL: (zz(v*) handles the mismatch in distributions py and Ps. over which the classification loss and cross cross-entropy losses are measured respectively. is upper bounded by the transferability coefficient
19

---

Published as a conference paper at ICLR 2021
Lemma E.10 Lemma E.
az(w) v for all RV
Error covariance -> cross-entropy loss (arbitrary language models): This is arguably the most important step that connects the error prediction the cross-entropy loss. For the arbitrary language model case, this is proved using Pinsker's inequality and taking expectation over the distribution Pt.
Lemma E.3
as( S v 2211011a (Eremi (PP (xem(P)))) for all RV
E.1.2 PROOF SKETCH FOR SOFTMAX LANGUAGE MODELS: THEOREM 4.7
Here we want to show guarantees for features Øp, = { on a (r, -natuuuall task T w.r.t Φ. From the definition of natural tasks w.r.t. 4, we know
FF + T
||00 B.t. e (1PT)),,
(24)
Note that the difference here is that v" is in the span of 4 rather than an arbitrary vector in RK We wish to upper bound the classification error ey (7(1p))()) and do so using the following sequence of inequalities.
(7((%/(t))) - inf (((()1(()))))>
XE
w= v= NER" {) -
inf
$ (((())())) 0²) E ( S (a az(2") az((")
(25)
where the first inequality follows because v" is in the span of & and second inequality follows from Equation (23 The bounds for 0n (v*) and 0² (v") are the same as arbitrary language models. The main difference is the bound (xs(v" which will be stronger bound for softmax models.
Error covariance cross-entropy loss (softmax language models): For softmax lan- guage models, we need to prove a modified version of Pinsker's inequality specifically for softmax models. This version will show bound that only works when v" is in the span of & and the evaluated model Py(s) computes softmax using d as well.
Lemma E.4
as(v)=
/
| (Exem (fps())) f
inf(((pr())) Vo + R
Thus we suffer the suboptimality of the language model {pf())} w.r.t. the best softmax model (p)) rather than the absolute best language model {p. This is done using the softmax variant of Pinsker's inequality in Lemma E.4 We now present the detailed proofs for all results.
E.2 PROOFS FOR ARBITRARY LANGUAGE MODELS
Theorem B.1
Strengthened Theorem 4.1). Let {p.aa} be language model that is c-optimal, i.e. E (p. Gzent for some > O. For a classification task T that is (T, B) B)-natural, we have
2
è (1p ST+
<(P: (p))))
For classification task T that is (r, B) B)-naturral w.r.f. 4, have
2Bac
e ((p s er ((Tp, S
(PT::: (p.1))
Proof. The proof has two main steps that we summarize by the following two lemmas. The first one upper bounds the downstream performance natural tasks with the covariance of errors.
20

---

Published as a conference paper at ICLR 2021
Lemma E.2 For language model (p.1.). Tf is (r, B.aturaal
Eep (A (1 )0 7(T: (p.1o))
(T(1p.1)) ST+
sup VR
T is (T, B)-naturll w.r.t. E RdxV
v Epe (A 1.) (plp: (P.1a))
((+P))))) ST+
sup n= AER eell SB
where 7()) and (-) are from Definition B.1
The second lemma upper bounds the covariance of enror with the suboptimality of the language model.
Lemma
E.6 For a language model [p-1a} and classifier v RV
v. J'Epp (A 11)1)))000 21/101130 (Crone ((on((PP)))) Erem))
[p]. Pip)) PT defined in Section
where Epr. (A (p.,)))
E supi.
We prove both the above lemmas in Section E.6 We first use these prove the main result.
Combining the two lemmas, we get the following inequality
Epr. (A (p.1,))
(((p)))) <(a) +
sup vERV SB
7 (PT; (P.1ao)
2 21100 (Exem((PP (p-1a)
lx
(0)
+
sup
M(P: (P.1))
vERV 1| SB
2B <PTi (p.1.})
+
where (a) uses first part of Lemma E.2 (b) uses Lemma E.6 and (c) uses the c-optimality of {p.))) This proves the first part of the result. The second part can also be proved similarly.
Ep.. (A (P.1)))0
(T((pp1)) 5(0) +
sup v= T V 10 <B
V
(PTT (P.1a))
(Esem ((p.)=7) 2 21101/3
(tems
<(b)
sup v= V || B
Ta (pT: (p.1a))
211 (Eremt (1p.1=)
(xent)
<(() +
2
ST
sup VERY SB
To (P (p.1.))
V
Yo (pri (p))))
where (a) uses second part of Lemma E.2
uses Lemma E.6 and (c) uses the c-optimality of
(P.p). The proof of the lemmas can be found in Section E.6
Theorem 4.1
Let {p.pa} be a language model that is e-optimal, i.e. (xcm(((p.)a)) Eu I for some For a classification task that (r, -nantrall,, we have
2B
(T( S 5 V
7(P))
21

---

Published as a conference paper at ICLR 2021
Proof. This follows from the first part of Theorem B.1
also show that ((P:: (p(1)))-'
(PT))-1 For that the following lemma that prove in Section E.6
Lemma E.9 For any g S + RD and pr As. have P (g) - - (ppp(1))) (1) (g)- 5
(PT)
Instantiating this for g = (pppa and using Equation 7 we get 7(Pi (D 3 (PT)
which completes the proof.
E.3 PROOFS FOR SOFTMAX LANGUAGE MODELS
Theorem 5.1 (Strengthened Theorem For a fixed 4, let be features from c-optimal d- dimensional softmax language model, i.e. Lawk], 1) (rem(()) S 6, where Exem is defined in Equation For classification task T that (T, B) Bnn-tturall w.r.l. D, have
2
&T (1P s (T(p))) S
1(P: ppf)
Proof. Instantiating Lemma B.2 for P.f. Ps(s), we get
Ep, (A (pri)))
(((9p)())) ST+
sup
V
(o (P)(,)))
AERV
|| SB
sup
AT (A (psw:)))
| XII
7(p: bpp)
sup
X Epe (#D (P(()t))
| All SB
(PT bpp)
where (a) follows from Equation 9 that says 7(T: kpp) To (pT: {Pr(a)) }). We now prove similar result for the second term in the following lemma that we prove in Section E.6
Lemma E.7 For a fixed & and softmax language model wiature f and X Rd
A E Epp. (DA())))) )a - T 1113. (Erend (f, (T))
where Epe (DA()))))) = E (Tp]() pp )) (Fp)() 4p. as defined in Section
Bf( (4)) and the e-optimality almost completes the proof. The only thing remaining to show is that & (T(()())) s (Tpp) which follows from the following sequence.
Using Lemma E.7 directly gives us ( = er( ((((
To (e(pr:0pp (pf)
(T((PrO))) = inf (T((Pr()), u) x xERY DEE
inf (((pFro TT) x,,))
vER DEE
inf (r((ppr()). (A.b)) = ()(pp)) DER"
For a fixed tb let I be features from an e-optimal d-dimensional softmax language model, i.e. Exene J. ②) ge ))) S where Erm( is defined in Equation (4 For classification task T that is (T, B) -nantrra w.r.t. 4. we have
Theorem 4.2
2B
ET (1p)(9)) s (r(Dp)) ST+ 11 7(PT)
22

---

Published as a conference paper at ICLR 2021
Proof. This result follows directly from Theorem 5.1 if also show that (pr: Upp )-1 (PT) just like in the proof of Theorem 4.1 For that we again use Lemma E.9 with g
$ (P(()))) and Equation and this completes the proof.
E.4 PROOFS FOR SECTION 4.3
We first show why Assumption 4.1 is approximately true when word embeddings are gaussian like.
Lemma E.1. Suppose word embeddings &w independent samples from the distribution ((u, 2). Then for any 0 R such that X² 00 >0 ((1) we have that log Zo) - _.00 Σ0 0' μ log(V)) I
with probability 1 - 6 for = Ô e
and ö = =1
(V))).
sw = A μ log EE e (dw-ni) ) thus can simply deal with the case where &n are sampled from (N Σ). Furthermore the only random variable of interest X = T Φn which a gaussian variable /O 0 Σ0) = -V ^²). Thus the problem reduces to showing that for V samples of X ~ (0, A²). log(Z) ()) is concentrated around
Proof. We first note that (og(Zo) = log
X + log(V) where = _.. exp(Xw). This can be proved similarly to the proof of Lemma 2.1
in Arora et al. (2016). It is easy see that X=~~((UAA3) [exp(X..) = ex² However the variable
exp( X w neither sub- gaussian nor sub-exponential and thus standard inequalities cannot be used directly. We use the same technique as Arora et al. (2016) to first observe that E[Z] = Vet and
Var [ < _[exp(2Xx)) = Ve2x². After conditioning on the event that X I {)oo((V) and
applying Berstein's inequality just like in Arora et al. (2016) completes the proof.
We next prove Lemma 4.3 that establishes a linear relationship between Φpy and f (under Assump- tion 4.1 and also the guarantees for on natural tasks.
Lemma 4.3 Under Assumption 4.1 any feature map S + Rd satisfies 4py(s) = Af(s) + for all
E.
Proof. Assumption 4.1
gives that log(Zo 10T AO + 0 b + c. We prove this lemma by matching the gradients of log(Z,, and the quadratic function on the R.H.S.
V 2o)) = VoZo
Emew e
Po poww)ow = 4po
Zo
Zo
WEW
Whereas the gradient of the quadratic part is Vol A0 + 0 b+c AO + b. Matching the two for 0 (() gives Dpf (s) 4P Af(s) +b.
Corollary 4.1
Using Lemma 4.3 for any e-optimal f, defined in Theorem 4.2 for classification tasks that are (r, natuuall w.r.t & we have lr(f) & T O(ve).
Proof. The main idea is that Lemma 4.3 gives us that Φpy(s) Af(s) b and thus any linear function of op, will also be a linear function ((s). From Theorem 5.1 (or Theorem 4.2), we also know that bpy will do well on T, i.e. ( S T O(Bvc). We formalize the intuition as
(7(2p,)
(AT. inf ( (A,6)) inf (Af + (A. inf AER2,6
XERd
inf (T(, ()) e(()
WER .y
This shows that ( < (7(0p)) S + O(BVE) and completes the proof.
"Noee that here we assume that we learn both linear classifier and an intercept for a downstream classifica- tion task. All results in the paper essentially remain the same with intercept in the definition of classification loss.
23

---

Published as a conference paper at ICLR 2021
E.5
PROOFS FOR SECTION
Theorem
C.1
The optimal solution
,", = arg min 1 Equaa D) satisfies
4- BUJ full rank B Bxddd
(*) ( //2" Pi = CU pi for fill rank C e Rdxd
If & is fixed, then the optimal solution is (s) = (FD.)) )- 1'/0""""
Proof. From Equations 110 and 11 know that, Lquad, (0, ②) -0 T + 2 01 and equad (f, +) = IE (equad,s (s),, P). For fixed #p, we define fī(s) arg mingeR" Lquad, (0, $).
srppr
We use the first-order optimality condition to get (s), by using the fact that V Vol Vulquad s (0, .)
-. + 2bT0. Setting the gradient to zero, get fā(s) (D4 pp To get the optimal $* for this objective, plug in this expression for fā in lquad and hnd $= arg ming (s.
Equar b)= gmp
E
(f E srup'
(6) pp' + +221 1 ((()))
1 +
IE
- -((DT) }1pp))))
4p.
11100
T (DD )
xmpp"
2
Pi|s T T T(D ) 1p... + 2/1s
T
TT ((D T) -- + (2 )
E
smpp
1
EE srup"
2 Pi
4 (D4 )
E eoup'
tr P. x ( -'ppppj))
tr
4 ④ (40) T) 1-14
E
pl.ei 7
2
Smp
1
7))
1
($
T)
> E
Φ. arup'
Pipple
2
2
where Ω* the substitutability matrix defined in Definition 5.1 the above objective reduces to ( b) = (VVT) ! reduces to learning an optimal V such that
Let $
NTV be the SVD. Then And hence learning the optimal 4*
arg min VERVX"VTVE
(VV ,9*
We will now show that the best such matrix is the matrix of top d eigenvectors of Ω², i.e. V Ua (cf. Definition Here we will assume that the eigenvalues of S2" are all distinct for simplicity of presentation. First note that (VVT, "VT22")
|VV s $ 112.. where ? US²UT, with U, Un and S define in Definition 5.1 This can be shown by the following sequence of steps
(VVT. ( (n(VV trV VV "2) > VVVVV
= VV)
= v
vv
=
Furthermore, we notice that VV T
| 1
|| VV u 213 as shown below
|| ||
IVV 12 2tr vv
It will be clear later that the optimal solution will have as high rank possible 4. All inverses can be replaced by pseudo-inverses for low-rank matrices.
24

---

Published as a conference paper at ICLR 2021
||2*
||
IVV
T 2V vv
||2*
|
VV
.
||
VV
||2
1
VV
Thus we get
arg mig VERV VE
(VV (2*)
| VV V 1 VVR
arg min
Note that VV r has columns that are columns of Ω projected on the space spanned by columns V. It folklore that the best such subspace V the subspace spanned by the top d eigen- vectors of Ω²z which is the same as top eigenvectors of f²,, thus giving us V*V* = UAUI Thus we get V = U.. for M = U V
This tells us that the optimal solution d will have SVD of the form d' = NT thus giving us 4* BU for matrix B NT M Rdxd This directly gives f fū. ( th ) p'p' N V P. = CU p/ for = N M
E.6 PROOFS FOR SUPPORTING LEMMAS
Lemma E.2. For a language model (p.1s). T (r, -attural.
v' Epr. ( (m 7(Ti (p.1o))
(T((P,))) +
sup vER" ||llI. SB
If is (T, B)-ntrrl w.r.t. $ RaxV
v Epe (A 11 (PT: (p.1))
((((p))))) ST+
sup v= AERV ||llo SB
where 7(-) and (-) are from Definition B.1
Proof. We note the following upper bounds on (f(pp,))) and (7((%p11)))
(r((p/a)) inf vERV
(Er(1p. v)
inf (er((p(s),)) VER"
(26)
|
inf w= NERV (7(10p1)))=
{))
w= T inf
(er((p(s):o)
(27)
"ea" DER,
1ell
When (r, -atturll by Definition 3.1 we know that inf
[ 2) S We now I
VERY
upper bound & ((p. J. inequality in Lemma E.8
using Lemma R.8 Taking infimum w.r.t. R.
||0| S B from the
(T(P))))))) (T(()))))) + VOT (A
inf (T(()))))))
inf (T(((()))))
sup vERY 11
/TTEp (A
VERV
vERR <B
||ll| SB
This, combined with Equation 26 gives us
(((p)))))
sup V
TTpp Epr ( (A(p))))
(28)
25

---

Published as a conference paper at ICLR 2021
Using Lemma E.10 and the definition of r(Pr: {p.(.)) in Equation 7 get that
v. Epr (A (1 <
Epp (A (p.1o)) Epp (p Exe (A ( #11 (o
(o E'E (A
Epp (A (p11))
(29)
(PT: (p.(o))
We have thus successfully transferred the bound from the distribution pr to PL. Combining this with Equation (28 completes the proof of the first part of the lemma.
We now prove the second part of the lemma where we only assume that 7 is (T, B) natural w.r.t. Q. Here instead take the infimum over classifiers in the span of Φ in Lemma E.8 to get
inf D= AER DER, vll SB
{(T((p))).)) v
inf v= AER DERR v
& ((((ptc):
sup w= AER" ||o|| SB
VUT (A (1p))))
(30)
This, combined with definition of (T, B) -ntuurll task w.r.t. & and Equation 27 gives us
(T((pp)))) T
sup
v ET (A
(31)
v= AERY ||e|| SB
For the last term, for any v = + X, E RR
we notice that
T Epr ( (p...) 2 = A AT ( (p. 1 T A= A
A
Epr (1A (p)r))
(a)
Epz (D (p Σ + 1 ppp (D) (x)) 11/ (1 Epe (DA(p)))A))
A Ep ( (p-1.1)
Ep (p)a))0
70 (PT: (P.12))
70 (PT; (P.12))
This combined with Equation
31 we get
v Epe (A (p,,)))
((0p:)))) ST+
inf AERV
7o (pT: (P.1o))
|ell $B
Lemma E.3 (Pinsker's inequality). For discrete distributions q,9" E Av, let q. RV be the corresponding vector of probabilities. Then we have
10 (q 9²)) S 22KKI
max [1
Proof. This basically follows from Pinsker's inequality which upper bounds the total variation dis- tance between distributions by their KL-divergence
max 10 (q q" q q²112 = 2 .g) 22DxL 1
We remind the reader that for embedding matrix & Rd Ps.4 softmax TT))
Lemma E.4 (Softmax variant of Pinsker's inequality). Consider a matrix & RRdXV with d < V. For any discrete distribution q" E By and softmax distribution P8, = softmax( T Ay for 0 E Rdd, let q", Pa E IRV be the corresponding vector of probabilities. Then we have
max
10 T (Po.0 -q") s (DLL(PO,, (')- inf DEL (po Eld
2
(32)
26

---

Published as a conference paper at ICLR 2021
Pinsker's inequality (Lemma
E.3
the other hand, gives
10' (po. 9-)) - 22Dk
max
1 11
Proof. Define the loss p(0) == DKL q²). The statement in Equation 32
to prove reduces to
( - -00)) /2 p(() (m(0)- inf
max o
p(0r)
(33)
ErERd
To prove this, compute the gradient and hessian of p(0) w.r.t. O. We can simplify p(0) as follows
cd th
p(0) (pL
E (logpp,,((w))) wrig"
log
Σ
el
wmyg
-0 -000q0 + -og
Σ
e ou
=- -0000 log (Zo)
The gradient is
Vio
Vp(0) = [-0TT D8 + logl Za)] = =-0-" +
7e
T oTo on
Zg =-0- + En g qua Dw
=- +
Zu
=qq" + Apo.0
Similarly the Hessian can be computed
e
v - ((V)))((V((0) = VI-2eq
+ 2pp.0l =
Σ
Pa,a (w)
=
Zo
wEW
WEW
e
dv
c dw
Ow
Σ
TT
Ooa Qw'
Zo
a0 wEW
E wm-poo0
E sp--po,
[dw]]
E mropy..
[(ho)]
CoV [pw)
wo-poo Cov [p denotes the covariance of the word embeddings &w when measured w.r.t. the distribution pa.a- This directly gives that p²) > 0, since the covariance is always psd, and thus is convex in 0.
Where Cov
We return the statement in Equation (33) that need prove. With the expression for gradient of at hand, we can rewrite Equation 33 as trying to prove
(VTP)(() S T T \ /2 p(())
inf ,(Or)
(34)
e" 0"ER""
Furthermore, using the definition of the Hessian, it is not hard to see for some X, 0 E Bd that A 2 x())> Cov [A [Twl
E [(A TT < | A/12. Thus we evoke wr-Po
Lemma E.5 with f = and I. = || A1I to prove Equation (34) and thus completing the proof.
Intuitively Lemma E.5 exploits the smoothness of the function to argue that small suboptimality (i.e. being close to optimal solution in function value) is sufficient to guarantee small norm of the gradi- ent, property that well-known in the optimization literature. We now present this lemma
O
Lemma E.5. If function E Rd + R and X E Rd satisfy A { s L, L e Rd (L- smoothness in the direction of X) and if inf deRd E(0), then X V((0)) (T(((0) e*)
27

---

Published as a conference paper at ICLR 2021
Proof. This is variant of classical result used in optimization and prove it here for complete- ness. For any 17 E R we have
{(0) e" 2(0) ((0) L(O (m/)
(b) {(0) ((() + (Ve(0), + n {
r
(c) 7 T VE(O))
2
where (a) follows from the definition of infimum and (b) follows from Taylor's expansion for some i E 10-n2, 0) and (c) follows from the smoothness condition in the statement of the lemma. Picking Viio) gives us {(0) - e" 2 A V2 thus completing the proof.
Lemma E.6. For a language model {p.1s} } and classifier E ERY
v V' ( (A) S 21/00130
- -Gee)
where Epe. (g) = E [ and A (s)))(()) = p.1s - Pi1s are defined in Section B
sepp
Proof.
We first note that
P.ja (w)
plss (w) )1
(xem(((P.))) fxeml fxem((1)) ((p1,))=
E srrp woopli
E
log
E srup..
(PL
(35)
We bound v Epr. (4 (m) ) below
(00 (P p(i)) (p
v T E'Epp (4 (A (p.
EE supL
(a) |1011 IE
(2D.L(P.... P.1a)]
smpL
(b) 21101133 (frem((P - bxem (P.).))
where (a) follows from Lemma
E.3 (Pinsker's inequality), (b) uses Equation 35
Lemma F.7. For a fixed $, a softmax language model with features f and X E Rd.
A Ep (P 1) s 2 T AlI (Crem (f, 1) (1))
where Epr. (D) (prcn))
E sepp..
(Tp -bp (OPF() - tpp"
defined in Section B
Proof. We start by nothing that A Ep (+D = T
E
( -
- -pp)))
sOPL
E [|] (p)()) 4p =
E s
[1 A) (Pf(s) p.),))3]
SP-PL
We will use the variant of Pinsker's inequality from Lemma [exem,
E.4 to bound each term on the right (f(s),
hand side. Notice that exemt (f, ②) au ( sppr.
(f
②)
inf ( })
BERd
X ETpp ( (p
E Srup
[1( X) ( p1
<(a) 21 T 11
E anopl.
DKI(Pis:ProN
inf DKL (P Po.o) BERa
x12 < 2
E
Cremt (f(s),((ss),)))))))) 1) BER4
inf bxemt (0, $)
smpt
< 2
2l2 (Erem (S, from (T))
where (a) follows from Lemma
E.4 This completes the proof.
28

---

Published as a conference paper at ICLR 2021
E.6.1 CLASSIFICATION LOSS TO COVARIANCE OF ERROR
Lemma E.S. For any task T and classifier v E RRV and predicted probabilities {p.1a}
er((p.t) v) S (7((Pt,)).)
E
(p-l0- -P.
suuT
=c7((P/s)(e) + ++vee" VOT Epy (A (p.,))
where Exp (g)
[ and A (p))(()) = = poja are defined in Section SmpT
E
Proof. The following sequence of inequalities proves it
(T((P(s), (P.))PT
E
[e(w Pll s(a)
[( Pi.: + +100
(p. +1wJ(pp
P.oo]]
(a.1)>pP
v (D. - -P|po) 27
(6)
((v Pia: +
E (S.
E supT
(pTTe PIN)( - p()]T]
(T((P))a)))) +1/0 1/07 N/1
E eupp
= =r((f v) + Vu ".'Ep (A
where (a) follows from I-lipschitzness of é, (b) follows from Jensen's inequality.
E.6.2 HANDLING DISTRIBUTION SHIFT
Lemma E.9. For any g S + RD and pr As, have | (g) (1)- Ep) (g))pp 2 ((PT)-
Proof. By definition of r(PT), we have that
Epp [=())g(s) Σ (((())(s)g(())
8NPL
> (PT) Σ PT (5) x(P)) snup
E (9(s))(s) - (g)
>(PT)
Thus 7(PT) Epz (g) Epr (g) - - finishes the proof.
222p (g)2ipz (g) Epy (g) (g) which is equivalent to ((P)) ID > Ep, (g) -
(p) and hence 7(P)) 21 -
- Epp (g) Epp (g) - This
Lemma E.10. For matrices X,Y Y E D s.s. X,Y & 0 and Y is full rank, we have that
max xa IY XY la for any norm ||. GERP,O<[|a||s.
Yoa
T Xa Proof. Note that a Ya completes the proof
is independent of the scaling of a. The following sequence of inequalitics
T Xa Yaa
Xa
T Xa
max oERP 0<
max aEED
max aERD (Yta)T (Y (rta)
la
SA a
a Ya
max
Xa
max
(Y-1
X(Y-11) 1b)
aERD y faalma=1
bE 0.Ro_bba1==
b Y XY b
#
bERD,
29

---

Published as a conference paper at ICLR 2021
F EXPERIMENT DETAILS
For all experiments we use the 117M parameter "small" GPT-2 model proposed in Radford et al. (219) and implemented in HuggingFace (Wolf et al., 2019). Linear classification experiments (except for fine-tuning baseline in Table T are performed on fixed output features from GPT-2.
We note that the binary SST-2 dataset used in all experiments is comprised of complete sentences, and there are 6,920 train examples and 1,821 test examples. In particular, this dataset is smaller than the version included with the GLUE benchmark (Wang et al. 2018). This smaller version of SST-2 better fits the sentence completion hypothesis we propose.
F.1 SOLVING DOWNSTREAM TASKS USING J AND Opj
The features f from GPT-2 for any input sequence (wi...., wx) is the output embedding of the final token WN at the final layer, where N is the input length and can be different for different inputs. This is also the embedding that is directly multiplied by the word embeddings to get the softmax distribution for language modeling, as in the theoretical setting. To use prompt, the same prompt is added at the end of all inputs and the features are extracted for this modified input.
We use the Logist ceegressionnv class from the scikit-learn package to fit linear classifiers to all fixed features (i.e., no finetuning). We use the liblinear solver and one-vs-rest loss function unless catastrophically fails (e.g., close to random performance) on particular multi-class task. In that case, use the stochastic average gradient (SAG) algorithm with multinomial loss. We use 5-fold cross validation for all experiments and test values for the regularization parameter C between 1e-6 and 1c4 for small datasets (i.e., fewer than 10K examples) and between 1c-3 and 1c3 for larger datasets.
Details about word subsets: For all of the results presented in Table we use a pre-trained GPT- 2 model. For SST, we use the prompt "This movie is when indicated. For AG News, we use the "This
prompt "This article about " when indicated.
We compute the conditional probability of selecting a subset of words to complete the sentence. For AG News, this subset is: 'world', 'politics', 'sports', 'business', 'science', 'financial', 'market', 'foreign', 'tecnnoloyy, 'international', 'stock', 'company', 'tech', "technologies'. For SST, this subset is: ")),, ":C', 'great', 'charming', 'flawed', 'classic", 'interesting'. 'boring'. 'sad', 'happy', 'terrible', 'fantastic', "exciting". 'strong'. For AG News, the class words we are: 'foreign', 'sports', 'financial', 'scientific' For SST, the class words we use are ":)' and "(""
We account for BPE tokenization by using the encoding of the word directly and the encoding of the word with a space prepended. We then filter to use only words that encode to a single BPE token.
Tests on additional datasets: We also test the performance of pre-trained GPT-2 embeddings and the conditional mean embeddings 4py on the DBPedia (Auer et al.) (007), Yahoo Answers (Zhang et al. 2015), TREC (Li Roth 2002), IMDb (Maa et al.) 2011). Customer Review (CR) (Hu & Liu 2004), and MPQA polarity (Wilson& & Wiebe 2003) datasets in Table We limited the training set size to 250K for larger datasets (i.c., DBPedia and Yahoo Answers). For CR and MPQA, follow Zhang et al. (2015) and average the performance across 10 random 90-10 train-test splits of the dataset.
We find that Φy, consistently has comparable performance to f across non-sentiment and sentiment downstream classification tasks. We include baseline results bag of n-grams (BonG) for most tasks and the mLSTM model (Radford et al. 2017) for sentiment tasks. BonG performs quite well on the larger datasets, but not as well on smaller datasets, due to the high dimensionality of features.
For sentiment tasks, adding prompt almost always boosts performance. We also demonstrate that much of the performance can be recovered by only looking at "positive" and "negative" or ":)" and ":(" as class words. Using these 2-dimensional features is even more sample-efficient than the standard 768-dimensional ones.
Link to code: https://github..om kams aad1 /t////attemmn hemat 1-expl peeeploooo oo-aaooondowwwsrrrra--ccco
30

---

Published as a conference paper at ICLR 2021
Table 2: GPT-2 performance without fine-tuning on downstream task sets with &t classes. We provide the performance of bag of n-grams (BonG) as an approximate baseline for these tasks. AG News, DBPedia and Yahoo performances were reported in ZZhng al. (201)), and the other tasks were reported in Khodak et al. (2018). We also include results from mLSTM (Sentiment Neuron) (Radford et al., 2017) for the sentiment-related classification tasks (SST, IMDb, CR, and MPQA) with numbers reported from Khodak et al. (2018). Furthermore, we include results for BERT (Devlin (et al., 2018) features without fine-tuning, where we use the output features for the first position of an imput for linear classification. An asterisk indicates add standard sentiment prompt "The sentiment is" to each input, but for AG News we used the prompt "This article is about". We also tested the performance of the conditional probability distribution over "positive" and "negative" as well as ":)" and "x(" on the sentiment-related tasks with and without the prompt.
Task
(s(s) (p)(() P.y.- pos,neg P.p= )
BonG
mLSTM BERT
Non-sentiment
AG News
4
990.7
84.6
92.4 (n 5)
88.9
AG News
91.1
88.2
89.9
DBPedia
14
97.2
88.2
98.6 (n 5)
98.7
Yahoo
10
69.2
56.7
68.5 n = 5)
65.0
TREC
6
93.6
87.8
89.8 (n 3)
90.6
Sentiment
SST
2
87.5
83.3
74.9
78.7
80.9 (n = 2)
91.8
85.8
SST*
2
89.4
87.3
80.8
79.1
84.1
SST fine
5
49.2
43.5
37.5
39.2
42.3
(n = 3)
52.9
43.5
SST fine*
5
49.4
48.0
41.5
40.2
43.3
IMDb
2
88.1
82.7
73.8
76.2
89.8 (n = 3)
92.3
82.2
IMDb*
88.4
85.3
81.8
80.9
84.0
CR
2
86.8
84.6
74.9
80.0
(n 3) 78.3
91.4
85.5
CR
87.9
87.1
82.5
79.4
84.6
MPQA
2
86.0
79.2
75.6
70.7
85.6 (n = 3)
88.5
87.3
MPQA
87.88
86.1
80.3
71.4
88.1
We also include results using the pre-trained BERT base cased model (Devlin al. 2018 Wolf et al. 2019), using the embedding at the first token as input to the downstream task. We also tried using the mean embedding and last token embedding and found that the first token embedding is often the best. Moreover, the first token embedding is what is extracted in the traditional usage of BERT on downstream tasks, though we note that it rare use BERT without fine-tuning.
F.2 FINETUNING EXPERIMENTS
As strong baseline, we finetune the GPT-2 features along with learning linear classifier for the SST and AG News classification tasks and report accuracy numbers in Table We use maximum sequence length 128 BPE tokens for downstream inputs SST-2 and maximum length 400 BPE tokens for AG News inputs. We use the end of sentence token the padding token. The datasets are described below.
1. AG News has 108K train examples, 12K dev examples, 7600 test examples. We split the train set for AG News into train and dev (90-10) and use the same test set as the non- finetuning experiments.
2. The sentence version of SST-2 has 6,920 train examples (same as non-finetuning), and 810 examples for dev and test each (split the original test set in half).
3. Fine-grained SST-2 has 8,544 train examples (same non-finetuning). and 1,105 examples each for the dev and test data (split the original test set in half).
To select the best hyperparameter configuration, we run a grid search over learning rate and batch size. We train each model for 10 epochs. For all datasets, test learning rates 5e-5, 1e-4, and
31

---

Published as a conference paper at ICLR 2021
Table 3: Comparing Quad features to cross-entropy features for GPT-2 trained on the IMDb un- labeled corpus (Mass et al. (011). In this experiment we fix + to be the word embeddings from prertained GPT-2 model for the cross-entropy objective. For the Quad objective, we initialize & to be the SVD of the pre-trained embeddings. An asterisk indicates that we added the prompt "This movie " to each input.
Task
(() (xent)
$pp(() (xent) f(x) (Quad)
SST
82.1%
79.9%
77.3%
SST*
83.1%
81.1%
80.7%
Table 4: Comparing Quad features to cross-entropy features for GPT-2 trained on the Amazon corpus. An asterisk indicates that added the prompt "This movie to each input. Note that the validation loss was still decreasing at the time of measurement.
Task
f(s) (xent)
4py (s) (xent)
f(s) (Quad, learned +)
SST
89.4%
89.7%
79.2%
SST*
89.7%
89.2%
84.3%
3c-4. For both version of SST-2, we try batch sizes 8, 16, and 32, and for AG News, we try batch sizes 8, 12, and We note that the longer sequence length AG News inputs required us to use parallelization across multiple GPUs to simulate larger batch sizes, which made batch size 32 prohibitively expensive to test.
We take the hyperparameter configuration that achieves the best performance on the dev set and then perform fine-tuning using those settings with three different random seeds: 8, 33, and 42. We then report the average performance on the test set in Table
We perform the hyperparameter grid search over the standard datasets and then perform fine-tuning using the best settings the dataset with task task-specific prompts added. For SST-2, we use the prompt "This movie is " and for AG News we use "This article is about
F.3 TESTING QUAD OBJECTIVE
We test two models with the same parametrizations, trained using Quad objective and
other trained with the standard cross-entropp objective using the unlabeled IMDb corpus (Maas et al. 201T) and the Amazon product review corpus (McAuley et al., 2015). We slightly modify the standard architecture GPT-2 to generate Tables and First we add a single linear layer (that is trained) top of the output features of the standard Transformer architecture. Furthermore, instead of tying the input and output word (token) embeddings, we learn them separately so that f and & are independent functions; this more in line with out theoretical setup. We fix the input embeddings and the positional embeddings to be the parameters from the pre-trained GPT-2.
For Quad, we initialize Φ the output embeddings, using the singular vectors of the pre-trained word embeddings 4. For the cross-entropy models, we initialize & to be the full pre-trained word embeddings +, because we found that initializing with the singular vectors harmed performance. Given our parameterization, initializing with the singular vectors is as expressive as initializing with the pretrained embeddings themselves; however it potentially lends better optimization landscape and speeds up training for our new objective Quad. As described in Section 5.2 minimize the following objective
Equad 1)
E.
[(ss) ow + (8,w)
2 7 (TTTf(s)
(36)
where (s, w) are sampled from the text corpus. The implementation of the Quad loss is the same as the standard cross-entrope loss, the main difference being the second term: it is | f(s) 12 for
Quad instead of the log-partition function log Eww, (e)) d., in the cross-entropy objective.
32

---

Published as a conference paper at ICLR 2021
True log partition learned quadratic function
DBPedia
Predicted log partition
100
100
-50
0-
1000
-50
200
-100
-200
-100
-200
-100
-100
100
TREC
MPQA
CR
-40
0-
Predicted log partition
-60
-50
0-
-80
-50
-100
1100
100
-120
-120
-100
-80
-60
-50 True log partition
-100
0
-100
-50
0
50
True log partition
True partition
Figure 2: Fit of the learned quadratic function to the log partition function various datasets for features computed by the full, pre-trained GPT-2. We also plot the x line for reference. These plots meant verify Assumption 4.1
Because IMDb is smaller dataset, we fix at its initialization and only train f to generate Table 3 When training on the Amazon dataset, we initialized & the same way we did for the IMDb dataset, but we allowed I and to both be trained, since more data was available. To train the models,
the standard learning rate schedule in in Radford et al. (211). To learn model on IMDb, we use context size 512 BPE tokens, and for the Amazon reviews dataset (McAuley al. 2015), use the standard context length of 1,024 BPE tokens.
We observe that training using Quad, in both cases, yields comparable performance to the language model the SST task, but always slightly worse. According to the theory, features f(s) from Quad should learn P.j. on a subspace, just like 4p, from cross-entropy models, thus making the comparison between these two important. Furthermore, adding a prompt consistently improves performance for both objectives. While Quad did not beat the cross-entropy either case, good performs at least demonstrates that insights from the theoretical analysis can translate to practical algorithms. We leave exploring the gap performance between Quad and cross-entropy and more extensive evaluation of Quad for future work.
F.4 LEARNING THE QUADRATIC. APPROXIMATION OF THE LOG-PARTITION FUNCTION
In Assumption 4.1 we assert that there quadratic fit for the log partition function, which allows us to show in Lemma 4.3 that a linear relation holds between and hpy. We validate these theoret- ical findings by fitting quadratic function to the log partition function for subset of embeddings from the IMDb, SST, and AG News datasets (Figure T.. Here, describe how learned A, b and c. To ensure is symmetric and positive semi-definite as required, we parametrize A U
As defined earlier, the partition function Za = Ew gTow owr and ppp = En. o TTow 0 E IRd. We minimize the following objective function:
2a Der for any
33

---

Published as a conference paper at ICLR 2021
Logebic Loss vs Cross Entropy for 6pne
Logoöic Loss Cross Entropy for 4pts
0.70
0.70
0.65
0.64
0.60
0.62
038
0.56
4.00
4.25
4.50
4.75
5.00
5.25
5.50
5755
35
4.0
4.5
5.0
5.5
Crows Entrrpy
Crosa Entrapy
(a) Trained on IMDb Maas et al.
2011 (b) Trained on Amazon (McAuley et al. 2015
Figure 3: Logistic loss of conditional mean features on the SST-2 task for various checkpoints of GPT-2 architecture trained on IMDb and Amazon. The reported cross-entropy is measured on the validation set. The red trend shows the fit of square-root function, which is what the upper bound in Theorem 4.2 looks like.
2 C(U,b, =E \ (108(Z0) - Yo 0T To 0000000--0 c)
A2 bpp U0 61r)
(37)
In practice, we train only on the regression loss (i.e., A = 0, A2 = 1) for the most promising results. Note that the regression term trying to learn a linear relationship between between and 4pp that Lemma 4.3 aims to prove. This ends up learning a matrix A = UU and vector b that also satisfy the quadratic form log( Zo) from Assumption 4.1
We use 20,000 examples from mix of IMDb, SST, and AG News embeddings as the training set. Thus we sample by sampling from the aforementioned datasets and set the feature map from pretrained GPT-2. We use the Adam (Kingma & Ba, 2014) optimizer with learning rate 1e-3 for U and learning rate 1e-4 for b and c. We decay the learning rate every 50 steps by factor of 0.1. We use the U obtained after 8 epochs of training. We further demonstrate the quality of the learned by plotting the true log partition and estimated log partition function for embeddings from other datasets in Figure 2
f(s). f being
F.5 EXPERIMENTALLS CHECKING THEOREM
Theorem 4.2 can be informally summarized as stating that an suboptimality in the cross-entropy of d-dimensional language model propagates ve increase in the logistic loss. We note that the B, and 7(Pr) factors are fixed for given pre-training corpus and downstream task, so we can empirically test if this square root relationship holds practice. In particular, Theorem 4.2 says
(T(pp) V 2B2 (m(PT)) (exemi (f. ②) ((
(38)
Of these, r, B,B, (pT) and ft are independent of the language model (f,, b) and only depend on the task T and language modeling distribution. Thus we can rewrite this as (T(ppy) S c + (f /vesem ②)- b for suitable constants b, E R.. The left hand side, ( is the logistic loss conditional mean features from language model (f, b) on task T and Exem (f. +)) is the cross- entropy loss of the language model, both of which can be measured in practice.
We train a 117M parameter GPT-2 model from scratch on the IMDb and Amazon corpora, described in Section F3
We maintain checkpoints during training, and for each checkpoint, measure the cross-entropy of the model on the validation set as well as the performance of the conditional mean features 4py on SST-2. Plotting these values together yields Figure 3
We furthermore fit a square root trend, shown in red, to these points. We learn a, b, such that y R3 av/z + c, where y = (1b,) is the logistic loss and x Exem(f. .) is the cross-entropy loss. For this, perform grid search over 100 evenly spaced valid values of b, and for each b,
34

---

Published as a conference paper at ICLR 2021
perform linear regression vz - b to find and We choose the a, b, that maximizes the r-value of the regression. While Theorem 4.2 only provides an upper bound on the logistic loss, this experiment shows that some square-rood root trend is observable in practice.
35


Reviews:
[{"review_id": "T6O2wO-1p9L", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper presents an explanation of why pretraining on language modeling (LM) helps performance on downstream text classification tasks. The explanation relies on formulating classification tasks as next word prediction tasks (i.e. language modeling). They use their theoretical results to design the Quad objective and experiment with it on SST and AG news, finding that it performs close but slightly worse than standard cross entropy training of classifiers.\n\nOverall, this work contributes an interesting framework for analysis. However, I have one large conceptual concern about the framework. Central to the proposed explanation is the ability to formulate text classification tasks as next word prediction, possibly with a prompt appended to the input (e.g. for sentiment analysis of movie reviews, “This movie is “). In a trivial sense, this is always possible: We can simply append the task definition to the end of an input as a question (e.g. “Is the sentiment of the review positive?”) and check the probabilities of “yes”/”no”. Then predicting the answer to this prompt is equivalent to performing the task, and a perfect LM is of course able to perform the task perfectly. This formulation makes Sections 3 and 4.1 feel trivially true. Though, to the authors’ credit, they do have to do additional work to extend an LM that is eps-optimal in next word cross entropy (i.e. on average) to optimality on the specific task formulation. \n\nHowever, the authors don’t mention this trivial reformulation strategy and instead base their argument on the existence of heuristic words (e.g. for sentiment analysis, the probability of “:)” or “:(“ after a review). This strategy introduces the potential for spurious correlations: The heuristics might be strongly correlated with the task in general, but might be off due to other factors like sarcasm. Additionally, relying on these single-word heuristics seems a bit off to me, as many text classification tasks don’t readily admit single words that encapsulate the task definition or label semantics. There’s an argument, then, that the theory described here doesn’t apply to these tasks (i.e. they’re not (t, B)-natural), but what’s frustrating about this argument is that this theory doesn’t provide us a way to distinguish which tasks fall in the category of single-word predictable or how to find such words other than trial and error. \n\nIt is very likely I am misunderstanding something about this paper. I am not sure what it means for a task to “lie in the row span of word embeddings”.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ljbfz8JWK8D", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper studies why language model pre-training has been such an effective technique in improving downstream performance across a wide range of NLP tasks recently.  In particular, it considers language models which compute a probability distribution over the next word in a text, given the previous context.  Then, taking inspiration from recent work that shows that many downstream tasks can be reframed as sentence completion tasks, it defines a “natural task” as one on which a sparse linear model over the output of the “true” language model (next word probability distribution, conditioned on context) attains strong performance.  Theoretically, it shows that language models which are close to the “true” language model are guaranteed to attain strong performance on natural tasks.  Empirically, it demonstrates that several NLP tasks are “natural”.\n\n\n**Strengths**\n- The paper is generally quite clearly written and the claims are well-validated.\n- The definitions, models, and assumptions in the paper are intuitive and clear (e.g., natural task).\n- The analysis which builds on these definitions/models/assumptions provides meaningful theoretical insight into why language model pre-training may be so beneficial for downstream training.  It provides a nice theoretical framework for thinking about the connection between language models and downstream tasks, which future work could build on.\n- The empirical validation is thoughtful and relatively thorough.  Even though the results don’t show that the proposed loss function and proposed “conditional mean features” give improvements over baselines, the empirical results show that the basic assumptions and definitions in the theoretical analysis are relatively realistic.  For example, Figure 4 validates Assumption 4.1 (log-partition function is roughly quadratic in theta), and Table 1 shows many real tasks are approximately “natural”.  Furthermore, when there is a gap between the empirical results and the theoretical results (e.g., validation of Lemma 4.3 at the end of Section 4), the paper makes these limitations clear, which I appreciated very much as a reader (paper does not over-claim its contributions).\n\n**Weaknesses**\n* Unclear if there are real practical applications to the insights from this paper.  Neither the proposed “Quad” loss function, nor the theoretically inspired “conditional mean features”, perform better than the baselines.\n* The current analysis doesn’t apply directly to BERT, which is trained to predict masked words in a sentence, instead of the next word.  Furthermore, BERT doesn’t predict these masked words using a linear softmax model over a contextual embedding for the whole sentence, which is the assumed structure for the softmax language models considered in the analysis.  (This limitation is acknowledged in the conclusion, which is good).\n* The paper doesn’t explain why learning a linear model directly on the context embeddings f(s) performs better than using the contextual mean embeddings.\n  * One idea I had here: Could you define a natural task as one for which there exists a sparse linear model over the *logits* of p*( . | s) which performs well, instead of a model directly over p*( . | s)?  Due to the very “flat” portions of the softmax function, there can be meaningful differences between the logits corresponding to 2 different words, but the LM probabilities for those words are extremely similar (and thus, harder for a linear model to distinguish).  With this definition, a linear model of the logits is also a linear model over the context embeddings f(s) directly.\n* There are some points in the paper that could be made clearer.\n  * I think it should be discussed earlier (in intro/related work) why the paper focuses on language models which do next word prediction via linear softmax models over fixed dimensional context embeddings, and that BERT is out of scope.\n  * I think there should be more discussion about the implications of Proposition 2.2.  As I understand it, this result shows that any part of p_f(s) orthogonal to row-span(Phi) doesn’t affect the cross-entropy of the language model (first order optimality condition would still be satisfied).  However, this doesn’t necessarily imply that p_f(s) will be in span(Phi) for all contexts s.  In particular,  the architecture of the embedding model f likely constrains f in such a way that makes it impossible for p_f(s) to be in span(Phi) for all contexts s.  Furthermore, at the end of section 3 it should be better explained why the assumption that p_f(s) is in span(Phi) for all s implies that Definition 3.2 should only consider sparse models v which are in this span as well (decompose v = v_in + v_out (component of v in the span, and orthogonal to the span), v^T p = (v_in + v_out)^T p = v_in^T p).\n  * I found the discussion in Section 4.1 pretty confusing.  In particular the part that argued why B = O(1/alpha).\n\nOverall, I really enjoyed reading this paper, and found it to be quite insightful.  It provided me with a much more thoughtful explanation for why language model pre-training improves downstream task performance, beyond simply “it helps learn good general representations of language using large amounts of unlabeled text data” (my previous reasoning).  As a result, I recommend acceptance for this paper.\n\nNIT: \n- Grammar last sentence of Section 1.1 (“…analyze the efficiency *of* …”)\n- Proposition 2.2: Maybe write “\\forall s \\in S” instead of “\\forall s ~ p_L”.\n- Section 3: ...append a prompt like “This movie is” (the final quotation mark is on the next line).\n- Equation (5).  Use “sup” instead of “max”.\n- Discussion in Section 4.1\n- I think Figure 4 should be explained in more detail (in caption and/or text).\n- Using capital and lower case tau in Theorem 4.2 is confusing notation.\n- Similarly, using bold and not-bold B in Theorem 5.2 is confusing notation.\n- After definition 5.1, what does Omega[w] = Omega[w’] mean?\n- In Table 1, can you explain more explicitly (in caption and text) what “subset” and “class words” means? Also, can you add a column where a dense linear model over p_f(s) is used?", "extra_scores": {}, "extra_text": {}}, {"review_id": "7s1JHlVw-Ua", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary.\n\nThis work tries to understand why features from trained language models can be used to solve classification tasks effectively. A language model (LM) in the analysis is modeled as a feature map $f : S \\rightarrow \\mathbb{R}^d$, a word embedding $\\Phi \\in \\mathbb{R}^{d \\times V}$, and a trained language model is thought of as $\\epsilon$-optimal in terms of its cross-entropy (from the true distribution over S). The work shows that for classification tasks approximately solvable by linear functions over the distributions of the next token the best linear classifier based on $(f, \\Phi)$ will suffer error of $O(\\sqrt{\\epsilon})$. The authors also propose an additional assumption where the log partition function is quadratic in $f$ based on which some improvements can be obtained. Being inspired by this assumption, a new objective function Quad where the partition function is directly replaced by a quadratic of $f$ is proposed. Experiments seem to support key assertions in the theoretical analysis.\n\nStrengths.\n\n1. The authors’ approach to a well-posed question seem original to me. In particular, some proposed concepts such as the refined transferability coefficient, conditional mean features, substitutability matrix might be useful for future studies.\n2. The article is precise, well-written and cautious in its tone. The accompanying experiments are informative and supportive of the main theoretical claims. I enjoy the overall journey the authors presented and would love to see more well-reasoned articles like this in ICLR.\n\nWeaknesses.\n\n1. The presentation can be improved by allocating more space to ideas in the Extensions section. This part seems more creative (perhaps a little less coherent but expected for “a mathematical exploration”) but is too compressed as it stands. (See Suggestion1)\n\nMinor issues.\n\n1. Consider to replace “partial sentences” with prefixes, which is more technically accurate.\n2. The many notations involving $p$ has inconsistent meanings for their subscripts. I would suggest to consolidate them to reduce confusion. For example, consider to use notations of this form $p(w|s; \\theta)$. Perhaps boldface for when it is viewed as a vector. Similarly for $\\ell$. (See Suggestion2)\n3. Below section 2.1, “trained to learn” -> “trained to fit”.\n\nSuggestions.\n\n1. I think a moderate revision reducing some (parallel) elaboration on “unconstrained language model” should provide the space needed for Extensions (and other novel ideas). After all, I found results on unconstrained LMs somewhat trivial and I suspect that you hope to use it only as an instructional tool. Perhaps a serial layout would save space and even improve the perceived emphasis.\n2. I strongly suggest a review of the notations and to adopt a more consistent scheme. One trick I found useful is to follow the notational convention of a textbook or a classic paper. The current notation has too much overloading and variability. \n\nQuestions.\n\n1. It seems to me that the central question lacks strong practical motivation. The NLP community seems to move to prefer a _natural_ answer (in the form of a generated sentence) instead of a label (from a classifier). As you have argued, many classification tasks can be framed as predicting the next token (perhaps in the presence of a prompt). What is your opinion? \n2. How realistic is the $\\epsilon$-optimality condition on cross-entropy for LMs? Can you comment on any associated sample complexity bounds? \n\n// Post-rebuttal update:\nThank you for replying to my questions. I am still concerned about the sample complexity associated with $\\epsilon$-optimality in cross-entropy (even in the trivial case, and perhaps impossible for some low-dim representations) as LMs are over a countably infinite extended alphabet.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7c16GW4v4AU", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThis work relates a pre-training performance with a downstream performance for tasks that _can_ be reformulated as next word prediction tasks. The authors show that for such tasks, if the pre-training objective is $\\epsilon$-optimal, then the downstream objective of a linear classifier is $\\mathcal{O}(\\sqrt{\\epsilon})$-optimal.\n\n**Strengths**\n- To the best of my knowledge, this is the first work that _mathematically_ justifies the connection between the pre-training objective and the downstream performance.\n- The proof technique (pre-training performance $\\to$ covariance of pre-training errors $\\to$ covariance of downstream errors $\\to$ downstream performance) is itself interesting. If the paper is accepted I am looking forward to seeing a high-level proof sketch in the main part as is done in Section 2.1 of [Arora et al. (2015)](https://arxiv.org/abs/1502.03520). A three-line explanation at the end of Section 4.1 seems a bit scarce to me.\n- The paper is well written: it gives an appropriate context, presents the main theoretical results, and verifies _some_ of the claims experimentally.\n\n**Major concern**\nIf I understand correctly (and please correct me if I am wrong), in Theorem B.1, the ratio between the downstream error $\\ell_\\mathcal{T}(\\\\{p_{\\cdot\\mid s}\\\\}) - \\tau$ and the pre-training error $\\ell_\\text{xent}(\\\\{p_{\\cdot\\mid s}\\\\})-\\ell_\\text{xent}^\\ast$ is _hidden_ in the $\\gamma(p_{\\mathcal{T}}; \\\\{p_{\\cdot\\mid s}\\\\})$ coefficient. Let me elaborate on this:\n1. In Lemma D.1 (with the help of Lemma D.9) you show that $\\frac{1}{\\gamma(p_{\\mathcal{T}}; \\\\{p_{\\cdot\\mid s}\\\\})}$ is an upper bound for the ratio $\\frac{\\boldsymbol{v}^\\top\\Sigma_{p_{\\mathcal{T}}}\\boldsymbol{v}}{\\boldsymbol{v}^\\top\\Sigma_{p_{L}}\\boldsymbol{v}}$\n2. The latter ratio seems proportional to the ratio $\\frac{\\ell_\\mathcal{T}(\\\\{p_{\\cdot\\mid s}\\\\}) - \\tau}{\\ell_\\text{xent}(\\\\{p_{\\cdot\\mid s}\\\\})-\\ell_\\text{xent}^\\ast}$. I am not sure on this---my intuition is based on your Lemma D.2 and the fact that for a $p_{\\cdot\\mid s}$ with full support a non-precise reverse version of [Pinsker's inequality](https://en.wikipedia.org/wiki/Pinsker%27s_inequality#Inverse_problem) holds.\n\nIn a nutshell, aren't you showing that\n$$\\text{downstream error}=\\mathcal{O}\\left(\\sqrt{\\text{pre-training error}\\cdot\\frac{\\text{downstream error}}{\\text{pre-training error}}}\\right)\\qquad ?$$\n\n**Issues**\n- Why don't you verify the main claim---$\\epsilon$-optimality in pre-training propagates as $\\mathcal{O}(\\sqrt{\\epsilon})$-optimality on downstream---empirically?For this, you may want to vary the language modeling performance (e.g. by pruning the language model) and then verifying that the downstream loss increase is indeed $\\mathcal{O}(\\sqrt{\\text{pre-training loss increase}})$. I believe such an experiment will definitely make the submission stronger.\n- I don't see why your theory does not generalize to a _masked_ language modeling (MLM). Why do we need to treat $s$ as the left-context only? Given the success of MLM as a powerful pre-training objective, please consider formulating your claims in a more general way.\n\n**Minor issues**\n- At the beginning of Section 2.3, $p_{\\mathcal{T}}$ is introduced as a distribution over $\\mathcal{S}\\times\\\\{\\pm1\\\\}$, but later (e.g. in formula (5)), it is used as a distribution over $\\mathcal{S}$ only. Please clarify/fix this.\n- What is the \"margin of task $\\mathcal{T}$\" mentioned on p.5? Is it a margin of an SVM classifier that solves $\\mathcal{T}$?\n\n**Limitations**\n- The authors admit that their work is limited to a particular type of downstream tasks. Indeed, it is not clear how one can reformulate e.g. linguistic tasks (like POS-tagging or dependency parsing) as a next word prediction task.\n\n**Update (after the author's response)**: During the rebuttal, the authors clarified my major concern, as well as provided additional experiments that verify the main claim of the paper. I am totally satisfied with the author's response, hence I am changing the score 6 $\\to$ 7", "extra_scores": {}, "extra_text": {}}, {"review_id": "3cDis2cmU1g", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of review:\n\nThere have been lots of interests to understand why self-supervised learning approaches such as the next word prediction task learn a useful representation for downstream tasks. This paper provides a mathematical framework to understand this question. One novel finding of this paper is that the distribution of the next word, conditional on the context, can provide a strong discriminative signal for the downstream task. In particular, using a carefully selected subset of \"prompt\" words, the authors observe that learning a linear predictor over the next word distributions of these words achieves performance close to a pre-trained GPT-2 model.\n\nSetting and Main Result:\n\nThis paper focuses on classification tasks, and the bulk of the work goes into how to model the next word distributions as features or representations. For this purpose, the authors introduced the definition of a \"natural\" task. Informally , a task is defined as natural if, just by using the next word distributions as features, the downstream task can be solved with a small loss.\n\nResult 1: Under the above assumption over the downstream task, this paper provides a bound on the empirical loss of the downstream prediction task. This bound consists of two parts:\n- The first part measures how \"natural\" the task is, that is, how well can the task be solved using the next word distributions as features.\n- The second part measures the difference between the \"empirical\" next word distributions and the \"optimal\" next word distributions.\n\nResult 2: The authors further extend this result to word embedding features, which are obtained by a weighted average of word embedding vectors based on the next word distributions.\n\nThere are several follow-up results built on these two results, such as a new loss objective for predicting the downstream task, but to the best of my understanding, these two results are the main claims of this paper.\n\nA key parameter that occurs in obtaining the above results is a worst-case coefficient that bounds the distributional shift between language model distributions of the training dataset and that of the downstream task. Intuitively, this parameter arises from translating the \"natural\" task assumption, which only guarantees transfer on average to the downstream task.\n\nPros:\n\n- A new framework for understanding why learning how to predict the next word helps the downstream task. This paper finds that the next word distributions of a subset of \"prompt\" words contain discriminative signals and are good features. This seems to be a novel finding and may help inspire future work in this important direction.\n\nCons:\n\n- The main result (Thm 4.1) applies to next/conditional word distributions that are very close to the optimal distribution. It is unclear to me how the authors are going to justify this \"assumption\". Should we expect the empirical distribution to converge to the true distribution when there are an infinite amount of samples?\n\n- Secondly, this main result depends on the worst-case coefficient, which is also unclear to me. For the transferability coefficient proposed in Section 5.1, is it possible to measure it in experiments? How large should we expect this coefficient to be?\n\nWriting:\n\nThe writing is overall clear and easy to follow, although it took me quite some time to map out the definitions of various notations. Many of the notations look cumbersome and I suspect that there is still room for making the notations more accessible for new readers.\n\n\nDetailed comments:\n\n- P2, Sec 1.1: \"analyze the efficiency language model features\" -> analyze the efficiency of language model features\n\n- P2, Sec 2: you started introducing these notations without explaining what they mean. For example, the $p^{\\star}$ notation is also defined in Sec 2.1.\n\n- P2, Sec 2: \"where $p^{\\star}_{\\cdot | s}$ is used as a vector on the left and distribution on the right\". What does this sentence mean?\n\n- P3, Sec 2.2: \"... achieve lower test perplexity than traditional n-gram models\" Why is this true? Could you add a reference?\n\n- P5, Sec 4.1: \"The result suggests that small test cross-entropy (hence test perplexity)...\" Same question as above.\n\n- P6, Sec 4.3: \"In fact, $f$ almost always performs better than ...\" This part seems intriguing despite the linear relationship shown in figure 1. Could you discuss this more here?\n\n- P8, Table 2: The results from using Quad look worse than the above two. Could you explain the significance of this result again?\n\n- P24, Figure 2: What are the x and y axis, and what does each dot mean in this figure?", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper attempts to provide a theoretical explanation for benefit of language model pretraining on downstream classification task. In this regard, the authors provide a mathematical framework which seems to indicate that the distribution of the next word, conditional on the context, can provide a strong discriminative signal for the downstream task. The reviewers found the formulation insightful, interesting, and novel. Also reviewers enjoyed reading the well written paper and appreciated its cautious in its tone. As correctly pointed out by reviewers, the proposed framework might not directly align with techniques used in practice. Applicability of the framework to other pre-training approaches is limited.  Also, there are some unresolved concerns about $O(\sqrt{\epsilon})$ assumption still. Nevertheless, reviewers reached a consensus that the framework would be beneficial for the community and attract follow-up works. Thus, I recommend an acceptance to ICLR. Following reviewer suggestion, it is strongly recommended that extensions section be expanded in the revised version using the extra page.
Decision: Accept (Poster)