A M ATHEMATICAL E XPLORATION OF W HY L AN -
GUAGE M ODELS H ELP S OLVE D OWNSTREAM T ASKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Autoregressive language models pretrained on large corpora have been successful at solving downstream tasks, even with zero-shot usage.
However, there is little theoretical justiﬁcation for their success.
This paper considers the following questions: (1) Why should learning the distribution of natural language help with downstream classiﬁcation tasks?
(2) Why do features learned using language modeling help solve downstream tasks with linear classiﬁers ?
For (1), we hypothesize, and verify empirically, that classiﬁcation tasks of interest can be reformulated as next word prediction tasks, thus making language modeling a meaningful pretraining task.
For (2), we analyze properties of the cross-entropy objective to show that (cid:15) -optimal language models in cross-entropy (log-perplexity) learn features that are O ( √ (cid:15) )
-good on natural linear classiﬁcation tasks, thus demonstrating mathematically that doing well on language modeling can be beneﬁcial for downstream tasks.
We perform experiments to verify assumptions and validate theoretical results.
Our theoretical insights motivate a simple alternative to the cross-entropy objective that performs well on some linear classiﬁcation tasks
I NTRODUCTION
The construction of increasingly powerful language models has revolutionized natural language processing (NLP).
Using gigantic text corpora and a cross-entropy objective, such models learn to predict a distribution over the next word after a given context.
The learned representations are useful for many other tasks, either as initializations (Ramachandran et al., 2017; Howard & Ruder, 2018) or as a source of contextual word embeddings (McCann et al., 2017; Peters et al., 2018).
Although representations previously needed ﬁne-tuning to solve downstream tasks, recent models (Radford et al., 2019; Brown et al., 2020) have demonstrated strong performance even without ﬁne-tuning.
Since next word prediction is a powerful test of language understanding, at some intuitive level it is believable that doing well on language modeling can help with many diverse downstream tasks.
At the same time, it is quite intriguing how even small decreases in test perplexity can lead to improved downstream performance.
Attempting to understand this phenomenon naturally raises the following questions: (a) why should learning the next-word prediction task result in representations useful for downstream tasks?
(b) what is the role of inductive biases of the models and algorithms used in this empirical success?
Given the nascency of deep learning theory, however, it is very challenging to say anything mathematically precise about (b) for even 2-layer networks, let alone deep networks.
Given these difﬁculties, this paper takes a stab at the mathematical study of (a).
We do so by exploring if and how quantitative improvements on downstream NLP tasks can be mathematically guaranteed as the language model distribution approaches the true one in some metric.
As a ﬁrst cut analysis, we restrict attention to classiﬁcation tasks and the striking observation that they can be solved fairly well using linear classiﬁers on language model features without ﬁne-tuning.
Although we are forced to treat models as black boxes, just ﬁrst-order optimality conditions reveal interesting properties of the learned features, leading to an understanding of their success on interesting linear classiﬁcation tasks.
We also use theoretical insights to construct and empirically test a simple objective that we call Quad, that provably learns useful features for classiﬁcation tasks.
We now summarize our contributions.
With the observation that classiﬁcation tasks of interest can be phrased as sentence completion tasks, we deﬁne natural classiﬁcation tasks , in Section 3, as those that can be solved as linear functions of the conditional distribution over words that can follow a given context.
Section 4 presents our main results, theorems 4.1 and 4.2, that mathematically quantify the beneﬁt of language model features for solving natural tasks.
We show that a (cid:15) -optimal language model (in cross-entropy) will do O ( √ (cid:15) )
-well on natural tasks; Theorem 4.2 proves a stronger result for low dimensional softmax models using a new tool that we call conditional mean features (Deﬁnition 4.1), which we show in Section 6 are good features for the downstream task.
The usefulness of the language model features themselves is demonstrated by arguing that there is a linear relationship between them and the conditional mean features.
Finally, we extend our theory to construct a new mathematically motivated objective Quad with useful guarantees in Section 5.2, and we report its good performance on linear classiﬁcation tasks in Section 6.
1. R ELATED WORK
Embedding methods: Prior to language models, large text corpora like Wikipedia (Merity et al., 2016) were used to learn low dimensional embeddings for words (Mikolov et al., 2013b;a; Pennington et al., 2014) and subsequently for sentences (Kiros et al., 2015; Arora et 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 et al., 2019; Liu et al., 2019) learn from auxiliary tasks, such as sentence completion, and are among the top performers on downstream tasks.
Language models for downstream tasks:
We are interested in language models, including n -gram models (Chen & Goodman, 1999), and more recent models (Xu & Rudnicky, 2000; Bengio et al., 2003) that use neural networks to compute low dimensional features for contexts and parametrize the next word distribution using softmax.
Language models have shown to be useful for downstream tasks as initializations (Ramachandran et al., 2017; Howard & Ruder, 2018) or as learned feature maps (Radford et al., 2017; McCann et al., 2017; Peters et al., 2018).
The idea of phrasing classiﬁcation tasks as sentence completion problems is motivated by recent works (Radford et al., 2019; Puri & Catanzaro, 2019; Schick & Sch ¨ utze, 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 a suite of downstream tasks as question-answering tasks (McCann et al., 2018) or text-to-text tasks (Raffel et al., 2019) and symbolic reasoning as ﬁll-in-the-blank tasks (Talmor et al., 2019).
Our work exploits this prevalent idea of task rephrasing to theoretically analyze why language models succeed on downstream tasks.
Theoretical analysis: 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) show that, in some regimes, the word2vec algorithm implicitly factorizes the PMI matrix.
The theory of Noise Contrastive Estimation (NCE) is used by Dyer (2014) to understand word embedding methods and by Ma & Collins (2018) to prove parameter recovery for negative sampling methods that learn conditional models.
A latent variable log-linear model is proposed in Arora et al.
(2016) to explain and unify various word embedding algorithms.
Theoretical justiﬁcation is provided for sentence embedding methods 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., 2019; Tosh et al., 2020b;a; Wang & Isola, 2020) and inpainting-based methods (Lee et al., 2020), which analyze the utility of representations learned in the self-supervised regime for downstream tasks.
Our work is the ﬁrst to analyze the efﬁcacy language model features on downstream tasks.
L ANGUAGE MODELING AND OPTIMAL SOLUTIONS
We use S to denote the discrete set of all contexts, i.e. sentences and partial sentences in natural language, W to denote the vocabulary of words, with V = |W| .
∆ A denotes the set of all distributions on A .
We use p, p L ∈ ∆ S to denote probability distributions over S , and p ·| s , p ∗·| s ∈ ∆ W to denote conditional distributions given a context s .
When clear from context, we use p ·| s , p ∗·| s ∈ R V as vectors of probabilities.
For v ∈ R V , we index the coordinate for w ∈ W by v [ w ] ; thus p ∗·| s [ w ] :=
p ∗·| s ( w ) , where p ∗·| s is used as a vector on the left and distribution on the right.
Let φ w ∈ R d denote a d -dimensional embedding for word w ∈ W ; word embeddings are stacked into a matrix Φ ∈ R d × V with columns being { φ w } w ∈W .
We use f : S → R d for a feature map that maps contexts to d -dimensional embeddings, e.g. f ( s ) can be the output of a transformer or RNN for an input context s ∈ S .
For embeddings { θ s } s ∈S , we use { θ s } to denote a function g with g ( s )
=
θ s .
2. U NCONSTRAINED LANGUAGE MODELING USING CROSS - ENTROPY
A language model is trained to learn the true distribution of a text corpus.
Given a context (e.g., a partial sentence s ∈ S ), it predicts a distribution p ·| s over the word to follow.
For example, given the context “The food was ”, the model could place high probabilities on words “delicious”, “expensive”, “bland”, etc.
We use p L to denote the true distribution over the context set S in the language modeling corpus and p ∗·| s to denote the true conditional distribution on words for s .
A standard approach is to minimize the expected cross-entropy loss between p ∗·| s and the model prediction p ·| s .
[Equation]
The optimal unconstrained language modeling recovers the distributions p
∗·| s exactly.
Proposition 2.1.
The unique minimizer of (cid:96) xent ( { p ·| s } ) is p ·| s = p
∗·| s for every s ∼ p L .
2. S OFTMAX PARAMETRIZED LANGUAGE MODELING
Recent models parametrize the conditional distribution p ·| s as a softmax computed using low dimensional embeddings.
Such models are often more sample efﬁcient and achieve lower test perplexity than traditional n -gram models.
For an embedding θ ∈ R d , the softmax distribution over W using word embeddings Φ ∈ R d × V is p θ ( w ) =
e θ (cid:62) φ w /Z θ , where Z θ = (cid:80) w (cid:48)
∈W e θ (cid:62) φ w (cid:48) is the partition function.
Just like p ∗·| s , we can interpret p θ ∈ R V as a vector of probabilities.
The context embeddings are parametrized by a model architecture of choice (e.g. transformer (Vaswani et al., 2017)) as a feature map f : S → R d .
These embeddings then induce the softmax distribution p ·| s = p f ( s ) .
We can now write the cross-entropy in a more familiar form
[Equation]
We rewrite it as (cid:96) xent ( f, Φ) =
E s ∼ p
L [ (cid:96) xent ,s ( f ( s ) , Φ)] , where (cid:96) xent ,s ( θ, Φ) is the cross-entropy loss for a context s that uses embedding θ .
Analogous to Proposition 2.1, we would like to know the optimal d -dimensional feature map f ∗ and the induced conditional distribution p f ∗ ( s ) 1 .
Proposition 2.2.
For a ﬁxed Φ , if f ∗ ∈ arg min f : S→
[Equation]
xent ( f, Φ) , then Φ p f ∗ ( s ) =
Φ p ∗·| s for every s ∼ p L .
Unlike Proposition 2.1, p f ∗ ( s ) ∈
R V is only equal to p ∗·| s ∈ R V in the subspace of the rows of Φ ∈ R d × V .
Thus smaller values of d will guarantee learning p ∗ ·| s on smaller subspaces.
This is proved by observing that for s ∈ S , ∇ θ (cid:96) xent ,s ( θ ) = − Φ p
∗·| s + ∇ θ Z θ Z θ = − Φ p
∗·| s + Φ p θ .
Applying the ﬁrst-order optimality condition for f ∗ ( s ) , i.e. ∇ θ (cid:96) xent ,s ( f ∗ ( s ))
= 0 gives the desired result.
2. D OWNSTREAM CLASSIFICATION TASKS
We focus on developing theory only for classiﬁcation tasks.
A binary classiﬁcation task 2 T is characterized by a distribution p T over S×{± 1 } , where the input s is a sentence from S and the label y is {± 1 } .
Given any feature map g : S → R D (for any D ), we solve task T by ﬁtting a linear classiﬁer v ∈ R D on top of g ( s ) .
The classiﬁcation loss is written as (cid:96) T ( g, v ) =
E ( s,y ) ∼ p T (cid:2) (cid:96) ( v (cid:62) g
( s ) , y ) (cid:3) , where (cid:96) is a 1-Lipschitz surrogate to the 0-1 loss, like the hinge loss (cid:96) (ˆ y, y )
= (1 − y ˆ y )
+ or the logistic loss (cid:96) (ˆ y, y )
=
log(1 + e − y ˆ y ) .
The loss incurred by a representation function g is deﬁned as
[Equation]
For embeddings { θ s } s ∈S , classiﬁcation loss is thus (cid:96) T ( { θ s } , v ) =
E ( s,y ) ∼ p T [ (cid:96) ( v (cid:62) θ s , y )] .
U SING LANGUAGE MODELS FOR CLASSIFICATION TASKS
Sections 2.1 and 2.2 imply that both unconstrained and softmax language models aim to learn either p ∗·| s exactly or a projection Φ p ∗·| s .
Thus we ask: why should learning p ∗·| s be useful for solving downstream tasks?
To answer this, we use a thought experiment where we have access to an oracle that provides p ∗·| s for any s .
We then demonstrate how to solve a classiﬁcation task through reframing the task as a sentence completion problem and using p ∗·| s to get completions to predict the label.
Classiﬁcation task as sentence completion: For an input s in a task, we propose to predict the label by using p ∗·| s to predict words that could follow s .
In particular, for a movie review sentiment analysis task, we can use the oracle to compare probabilities of “:)” and “:(” after an input and predict sentiment based on which is higher.
However, p ∗·| s will likely place higher probability on words that start sentences, like “The”, instead of discriminative words useful for the task.
To allow a larger set of grammatically correct completions, we can additionally append a prompt like “This movie is ” at the end of a review and query probabilities of indicative adjectives like good, bad, interesting, boring etc.
This approach also works for the AG news dataset (Zhang et al., 2015) which contains news articles from 4 categories: world, science/technology, sports, and business, where a prompt like “This article is about ” can be added.
We further theoretically motivate the usage of prompts in the discussion of Theorem 4.1, and we experimentally verify in Section 6 that probabilities of a small subset of words are sufﬁcient to solve these tasks.
Solving tasks using a linear function of p ∗·| s : The above process of solving tasks is actually a subcase of learning a linear function of p ∗·| s ∈ R V .
In the case of movie reviews, we can let w + = “:)” and w − = “:(” and compute the sign of p ∗·| s ( w + ) − p ∗·| s ( w − ) to predict the label.
This strategy can be expressed as a linear separator on p ∗·| s ∈ R V , with a classiﬁer v ∈ R V s.t.
v [ w + ]
= 1 , v [ w − ]
= − 1 and v [ w (cid:48) ]
= 0 for all other w (cid:48) ∈ W .
The prediction function then is v (cid:62) p ∗·| s > 0 .
This also works when we append a prompt; we can assign positive weights in v to adjectives like “good” and negative weights to adjectives like “boring”.
Strength of sentiment in different adjectives (e.g., “good” vs “amazing”) can be captured through different weights for words.
To solve other tasks similarly, we can select a different set of words for each class.
We verify experimentally that SST and AG news tasks can be solved by a linear function of probabilities of just a small subset of words in Section 6.
We formalize this intuition below to deﬁne natural tasks.
Deﬁnition 3.1.
A classiﬁcation task T is ( τ, B ) -natural if
[Table]
[Equation]
[Table]
[Equation]
Thus a natural task T is one that can achieve a small error of τ by learning an (cid:96) ∞ -norm boundedlinear classiﬁer on top of features p ∗·| s ∈ R V .
Low dimensional softmax models, however, only learn p ∗·| s in the subspace of Φ , as per Proposition 2.2.
Thus we cannot ensure guarantees on all natural tasks and we are interested in tasks that this subspace can solve.
Deﬁnition 3.2.
Task T is ( τ, B ) -natural w.r.t.
Φ ∈ R d × V if
[Table]
[Equation]
[Table]
[Equation]
This is a subset of natural tasks where the classiﬁer v on p ∗·| s is restricted to be in the span of Φ , i.e. v = Φ (cid:62) λ
= [ φ (cid:62) w λ ] w ∈W ; thus φ w ≈ φ w (cid:48) implies v [ w ] ≈ v [ w (cid:48) ] .
We argue that v being in the span of Φ may not be a restriction if Φ satisﬁes some nice properties.
For instance, if interesting tasks in practice tend to assign similar weights to pairs of words that are synonyms, then Φ just needs to assigns similar embeddings for those pairs.
In fact, Section 5.2 describes a carefully designed objective that can learn word embeddings where synonyms have similar embeddings.
G UARANTEES FOR LANGUAGE MODELS ON NATURAL TASKS
We now show guarantees for features from language models on natural tasks.
For an unconstrained model, we use the learned p ·| s ∈ R V as features and for softmax model f , we show guarantees for Φ p f ( s ) ∈ R d .
Since we cannot practically hope to learn the optimal solutions described in propositions 2.1 and 2.2 , we only assume that the language models are (cid:15) -optimal in cross-entropy.
To deﬁne (cid:15) -optimality in the two settings, we ﬁrst deﬁne optimal cross-entropies.
[Equation]
where (cid:96) ∗ xent is the absolute minimum achievable cross-entropy, while (cid:96) ∗ xent (Φ) is the minimum achievable cross-entropy by a d -dimensional softmax language model using Φ ; clearly (cid:96) ∗ xent ≤ (cid:96) ∗ xent (Φ) .
4. U NCONSTRAINED LANGUAGE MODELS
We show guarantees for a language model that satisﬁes (cid:96) xent ( { p ·| s } ) − (cid:96) ∗ xent ≤ (cid:15) .
An important consideration is that the language model distribution p L of contexts is often a diverse superset of the downstream distribution p T (see Section 2.2), thus requiring us to show how guarantees of p ·| s ≈ p
∗·| s on average over the distribution s ∼ p
L transfer to guarantees on a subset p T .
In the worst case, all of the (cid:15) error by { p ·| s } occurs on sentences from the subset p T , leading to pessimistic bounds 4 .
In practice, however, the errors might be more evenly distributed across p L , thus bypassing this worst case bound.
As a ﬁrst step, we present the worst case bound here; stronger guarantees are in Section 5.1.
The worst-case coefﬁcient γ ( p T ) , deﬁned below, captures that p T is a γ ( p T ) -fraction of p L .
[Equation]
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 classiﬁcation performance, thus formalizing the intuition that better language models will be more useful for downstream tasks.
Theorem 4.1.
Let { p ·| s } be a language model that is (cid:15)
-optimal, i.e. (cid:96) xent ( { p ·| s } ) − (cid:96) ∗ xent ≤ (cid:15) , for some (cid:15) > 0 .
For a classiﬁcation task T that is ( τ, B ) -natural, we have
[Equation]
Discussion: The suboptimality (cid:15) propagates gracefully as O ( √ (cid:15) ) to a downstream task.
While τ can be thought of as a small constant like 0.01, the norm bound B captures the margin of task T when solved linearly using { p ∗·| s } .
Intuitively, for (cid:96) T ( { p ∗·| s } , v ) to be smaller than τ , B needs to be large enough so that v (cid:62) p
∗·| s = Ω(1) .
Thus if the words of interest in the support of v , described in
Section 3, have total probability mass of Ω( α ) in p ∗·| s , then
B ∼ O
(1 /α ) .
It is thus desirable for a task T to depend on a larger and more frequent set of words.
A task that depends on probabilities of rare words, will have a high value of B .
Adding a prompt, as described in Section 3, can broaden the set of indicative words, thus potentially decreasing B .
A key step in the proof is to bound the difference in prediction on s for classiﬁer v as | v (cid:62) ( p ·| s − p ∗·| s ) | ≤ (cid:107) v (cid:107) ∞ (cid:107) p ·
| s − p
∗·| s (cid:107) 1 ≤ (cid:107) v (cid:107) ∞ (cid:113) 2( (cid:96) xent ,s ( p ·
| s ) − (cid:96) xent ,s ( p ∗·| s )) , using Holder’s and Pinsker’s inequalities respectively.
4. S OFTMAX LANGUAGE MODEL WITH CONDITIONAL MEAN FEATURES
We now describe guarantees for a softmax language model with feature map f that satisﬁes (cid:96) xent ( f, Φ) − (cid:96) ∗ xent (Φ) ≤ (cid:15) ; suboptimality is measured w.r.t.
the best d -dimensional model, unlike Theorem 4.1,.
Note that Theorem 4.1 can be applied here to give a bound of (cid:96) T ( { p f ( s ) } )
≤ τ + O ( B (cid:112) (cid:15) + (cid:15) ∗ Φ ) on ( τ, B ) -natural tasks, where (cid:15) ∗ Φ = (cid:96) ∗ xent (Φ) − (cid:96) ∗ xent is the suboptimality of the best d -dimensional model.
This ﬁxed error of O ( B (cid:112) (cid:15) ∗ Φ ) (even when (cid:15) = 0 ), however, is undesirable.
We improve on this by proving a stronger result speciﬁcally for softmax models.
Inspired by
Proposition 2.2 that shows Φ p f ∗ ( s )
=
Φ p ∗·| s , our guarantees are for features Φ p f ( s ) ∈ R d ; we call these conditional mean features.
Deﬁnition 4.1 (Conditional Mean Features) .
For a feature map f : S → R d and Φ ∈ R d × V , we deﬁne conditional mean features g f, Φ : S → R d , where g f, Φ ( s ) =
Φ p f ( s ) , where p f ( s ) ∈ R V .
The result below will show that conditional mean features g f, Φ are guaranteed to do well on natural tasks w.r.t.
Φ , thereby suggesting a novel way to use softmax features f for downstream tasks.
We also test g f, Φ ( s ) on downstream tasks in Section 6 and ﬁnd that they perform comparably to f .
We now present the result for softmax language models that has the similar implication as Theorem 4.1, but with above-mentioned subtle differences.
The proof (Section D.2) is similar to that of
Theorem 4.1, but crucially requires showing a d -dimensional version of Pinkser’s inequality.
Theorem 4.2.
For a ﬁxed Φ , let f be features from an (cid:15) -optimal
d -dimensional softmax language model, i.e. (cid:96) xent ( f, Φ) − (cid:96) ∗ xent (Φ) ≤ (cid:15) .
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , (cid:96) T ( g f, Φ )
[Equation]
We experimentally ﬁnd a O ( √ (cid:15) ) -like behavior in practice in Section E.4.
g f, Φ ( s ) IS A LINEAR FUNCTION OF f ( s )
Theorem 4.2 shows that g f, Φ is useful for linear classiﬁcation.
However, using feature map f directly is more standard and performs well in practice (see Section 6).
Here we argue that there is a linear relation between f ∗ and g f ∗ , Φ if word embeddings Φ satisfy a certain property, which we show implies that tasks solvable linearly with g f, Φ are also solvable linearly using f .
Our main assumption about word embeddings Φ is that the logarithm of the partition function Z f ( s ) is quadratic in f ( s ) .
Assumption 4.1.
There exists a symmetric positive semideﬁnite matrix A ∈ R d × d , a vector b ∈ R d and a constant c ∈ R such that log( Z θ )
= 12 θ (cid:62)
A θ + θ (cid:62) b
+ c for any θ ∈ R d .
If word embeddings were distributed as Gaussians, i.e. V columns of Φ are sampled from N ( µ, Σ) independently, it is not hard to show (Lemma D.5) that log( Z θ )
≈ 12 θ (cid:62) Σ θ + θ (cid:62) µ
+ 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 is constant for all (cid:107) θ (cid:107) 2 = 1 , our quadratic assumption is a bit stronger.
However, empirically we ﬁnd the ﬁt to be very good, as evident in Figure 1.
Under the above assumption, we can show a linear relation between f and Φ p f .
Lemma 4.3.
Under Assumption 4.1, feature map f satisﬁes g f, Φ ( s ) =
A f ( s ) + b , ∀ s ∈ S .
Corollary 4.1.
Under same setting as Lemma 4.3 and Theorem 4.2, (cid:96) T ( f ) ≤ τ + O ( B √ (cid:15) ) .
Thus we get that f itself is good for natural linear classiﬁcation tasks.
However, in practice, the linearity between f and g f, Φ holds only approximately when tested on features from the pretrained
GPT-2 language model Radford et al.
The ratio of the residual norm of the best linear map to the norm of f , i.e. r = E s [Equation] [Equation] 2 E
[Equation]
g f, Φ ( s ) (cid:107), is measured for different distributions p ( r =means perfect ﬁt).
These ratios 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 .
In fact, f almost always performs better than g f, Φ ; we leave exploring this to future work.
E XTENSIONS 5. B ETTER HANDLING OF DISTRIBUTIONAL SHIFT
The bounds in the previous section use the coefﬁcient γ ( p T ) to transfer guarantees from p L to p T and we deﬁne a more reﬁned notion of transferability here.
The coefﬁcient γ ( p T ) is independent of the learned model and assumes a worst case distribution of errors.
For the reﬁned coefﬁcient, we ﬁrst deﬁne the error made in predicted probabilities by a softmax language model f as ∆ { p f ( s ) } ( s ) =
p f ( s ) − p ∗·| s .
For any distribution p ∈ ∆ S , we deﬁne the covariance 5 of a function g : S → R D as Σ p ( g ) =
E s ∼ p (cid:2) g ( s ) g ( s ) (cid:62) (cid:3) .
The reﬁned transferability coefﬁcient is then deﬁned as
[Equation]
We now state the result for softmax language models; detailed results (including those for unconstrained language models) can be found in Section B.
Theorem 5.1 (Simpliﬁed) .
In the same setting as Theorem 4.2, (cid:96) T ( g f, Φ ) ≤ τ +
[Equation]
It is easy show that γ ( p T ; g f, Φ ) ≥ γ ( p T ) , so this is indeed a stronger bound.
The coefﬁcient γ ( p T ; g f, Φ ) measures of how average error by a language model f on p L can propagate to p T .
This can potentially be much smaller than γ ( p T ) due to some inductive biases of f .
For instance, if errors made by the model are random-like, i.e. ∆ { p f ( s ) } ( s ) ∼ ρ , independently of s , then Σ p L (∆ { p ·| s } )
≈ Σ p (∆ { p ·| s } )
≈
E
η ∼ ρ
[ ηη (cid:62) ] , making γ ( p ; { p ·| s } )
≈ 1 .
Independence prevents language modeling error from accumulating on contexts from p T , bypassing the worst case transfer of γ ( p T ) .
5. Q UAD : A NEW OBJECTIVE FUNCTION
In Deﬁnition 3.2 we discuss how low dimensional softmax language models learn a linear projection of p ∗·| s , only solving tasks that lie in the row span of word embeddings Φ .
Although Φ deﬁnes tasks that language model features can solve, the standard cross-entropy objective does not lend a 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 ∗·| s and thus can solve some natural tasks, and (2) the optimal Φ ∗ has an intuitively meaningful closed-form solution.
[Equation]
The Quad objective is very similar to the cross-entropy objective from Equation (2), with the log partition function replaced by a 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 .
Deﬁnition 5.1.
The substitutability matrix is deﬁned to be Ω ∗ :=
[Equation]
Ω ∗ =
USU (cid:62) is the eigendecomposition, then U d ∈ R V
×
d is matrix of top d eigenvectors of Ω ∗ .
The matrix Ω ∗ captures substitutability between pairs of words.
Words w and w (cid:48) are substitutable if they have identical conditional probabilities for every context s ∈ S and thus can replace occurrences of each other while still providing meaningful completions.
By deﬁnition, these words satisfy Ω ∗ [ w ] =
Ω ∗ [ w (cid:48) ] .
Such pairs of words were called “free variants” in the work on distributional semantics (Harris, 1954), and capture the notion of synonyms in the distributional hypothesis.
Theorem 5.2.
Let f ∗ , Φ ∗ = arg min f, Φ (cid:96)
quad ( f, Φ) .
Then Φ ∗ = BU (cid:62) d , for full rank B ∈ R d × d .
Also, for a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ ∗ , we have (cid:96) T ( f ∗ ) ≤ τ . [Table] [Table]
Thus f ∗ excels on natural tasks w.r
[Table]
t. Φ ∗ , which in turn, is the best d -dimensional projection of Ω ∗ .
Thus words w, w (cid:48) ∈ W that are synonyms (hence substitutable) will satisfy φ ∗ w = φ ∗ w (cid:48) , fulﬁlling the desired property for word embeddings discussed in Deﬁnition 3.2.
We train using the Quad objective and compare its performance to a similarly trained language model in Section 6, ﬁnding
Quad to be reasonably effective.
The goal of testing Quad is not to obtain state-of-the-art results, but to demonstrate that theoretical insights can aid the design of provably effective algorithms
[Table]
E XPERIMENTS
Tasks using linear function of p ·| s : We validate our claims from Section 3 that classiﬁcation tasks can be solved by linear functions of p ·| s .
Table 1 demonstrates that on the Stanford Sentiment Treebank (SST) (Socher et al., 2013) and AG News (Zhang et al., 2015) tasks we can use the conditional probabilities p ·| s =
p f ( s ) from GPT-2 (Radford et al., 2019) of just 20 intuitively task-relevant tokens (see Section E.1) to solve tasks.
Even just one token per class yields non-trivial performance.
Furthermore, we validate the complete-the-sentence intuition in Section 3 by using the probabilities after adding a task speciﬁc prompt and observing improved performance.
g f, Φ and f are good features: We validate Theorem 4.2 by verifying that the conditional mean features g f, Φ ( s ) =
Φ p f ( s ) also linearly solve downstream tasks fairly well.
This performance is comparable to, but always worse than f ( s ) , as seen in columns 2 and 3 of Table 1.
We again ﬁnd that adding a prompt improves performance.
Section D.4 has results for a wider range of classiﬁcation tasks.
Evidence for Assumption 4.1 is provided by learning a quadratic function to ﬁt the log partition function of features from pretrained GPT-2 model (see Section E.3).
Figuredemonstrates that the ﬁt holds for its training and unseen data (e.g., WebText (Radford et al., 2019)).d
Quad objective: We compare downstream performance of the Quad objective to the cross-entropy objective by training GPT-2 on the IMDb dataset (Maas et al., 2011).
Table 2 shows that Quad performs comparably to g f, Φ for f learned by the cross-entropy objective, which ﬁt our theory since both are linear functions of p ∗·| s .
We reiterate that these experiments are intended only to demonstrate how theoretical insights can motivate the design of new objectives.
C ONCLUSIONS AND FUTURE WORK
We provide intuitive and mathematical explanations for the success of language model features on linear classiﬁcation tasks.
The sentence completion intuition helps argue that p ∗·| s can solve natural tasks linearly.
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 g f, Φ and f , it is insufﬁcient to explain the observed superiority of f over g f, Φ .
We leave exploring this limitation of our analysis to future work.
Guarantees for softmax models are for natural tasks w.r.t.
Φ , thus knowing the optimal d -dimensional word embeddings Φ ∗ for (cid:96) xent ( f, Φ) 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 we would like to show stronger guarantees by exploiting model and algorithmic inductive biases, as well as study the setting of ﬁnetuning language model features, lack of a good theory of deep learning is the current bottleneck.
A O VERVIEW
Section B is a more detailed version of Section 5.1 and Section C is a detailed version of Section 5.2.
Section D contains proofs for all results.
Section E contains many more experimental ﬁndings that consolidate many of our theoretical results.
Section E.1 provides the information about subsets of words used for results in Table 1 and also additional experiments to test the performance of pretrained language model embeddings f on more downstream tasks and also verifying that conditional mean embeddings Φ p f do well on these tasks.
In Section E.2, we present additional results for our Quad objective trained on a larger corpus and tested on SST.
Finally Section E.3 provides additional details on how A , b and c from Assumption 4.1 are learned and also further veriﬁcation of the assumption on more datasets.
B B ETTER HANDLING OF DISTRIBUTIONAL SHIFT
While the bounds above used γ ( p T ) to transfer from the distribution p L to p T , we deﬁne a more reﬁned notion of transferability here.
While γ ( p T ) only depends on p L and p T , the more reﬁned notions depend also on the learned language model, thus potentially exploiting some inductive biases.
We ﬁrst deﬁne the notion of error made in the predicted probabilities by any predictor p ·| s as ∆ { p ·| s } ( s ) =
p ·| s − p ∗·| s .
Thus for any softmax language model f we have ∆ { p f ( s ) } ( s ) =
p f ( s ) − p ∗·| s .
For any distribution p ∈ ∆ S , we deﬁne the covariance 6 of a function g : S → R D as Σ p ( g ) =
E s ∼ p (cid:2) g ( s ) g ( s ) (cid:62) (cid:3) .
We deﬁne 3 coefﬁcients for the results to follow
Deﬁnition B.1.
For any distribution p ∈ ∆ S , we deﬁne the following
[Equation]
We notice that Σ p (∆ { p
[Equation]
∼
[Equation]
Σ
[Equation]
∆
[Equation]
Φ
[Equation]
We are now ready to state the most general results.
Theorem B.1 (Strengthened Theorem 4.1) .
Let { p ·| s } be a language model that is (cid:15) -optimal, i.e.
(cid:96) xent ( { p ·| s } ) − (cid:96) ∗ xent ≤ (cid:15) for some (cid:15) > 0 .
For a classiﬁcation task T that is ( τ, B ) -natural, we have (cid:96) T (cid:0) { p ·| s } (cid:1) ≤ τ + (cid:115) 2 B 2 (cid:15) γ ( p T ; { p ·| s } )
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , we have (cid:96) T (cid:0) {
[Equation]
Theorem 5.1 (Strengthened Theorem 4.2) .
For a ﬁxed Φ , let f be features from an (cid:15) -optimal
ddimensional softmax language model, i.e. (cid:96) xent ( f, Φ) − (cid:96) ∗ xent (Φ) ≤ (cid:15) , where (cid:96) ∗ xent (Φ) is deﬁned in
Equation (4).
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , we have
[Equation]
Discussions: It is not hard to show that the coefﬁcients satisfy γ Φ ( p T ; { p ·| s } ) ≥ γ ( p T ; { p ·| s } ) ≥ γ ( p T ) and γ ( p T ; g f, Φ ) ≥ γ ( p T ) , thus showing that these results are strictly stronger than the ones from the previous section.
The transferability coefﬁcient is a measure of how guarantees on p L 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 T ) , the coefﬁcients in Deﬁnition B. depend on the learned models, either { p ·| s } or { p f ( s ) } , 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 f ( s ) } ( s ) ∼ ρ , independently of s , then Σ p L (∆ { p ·| s } )
≈ Σ p (∆ { p ·| s } )
≈
E
η ∼ ρ
[ ηη (cid:62) ] , making γ ( p ; { p ·| s } )
≈ 1 .
Independence prevents language modeling error from accumulating on contexts from p T , bypassing the worst case transfer of γ ( p T ) .
C Q UAD : A NEW OBJECTIVE FUNCTION
In Deﬁnition 3.2 we discuss how low dimensional softmax language models learn a linear projection of p ∗·| s , only solving tasks that lie in the row span of word embeddings Φ .
Although Φ deﬁnes tasks that language model features can solve, the standard cross-entropy objective does not lend a 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 ∗·| s and thus can solve some natural tasks, and (2) the optimal Φ ∗ has an intuitively meaningful closed-form solution.
[Equation]
The Quad objective is very similar to the cross-entropy objective from Equation (2), with the log partition function replaced by a 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 .
Deﬁnition 5.1.
The substitutability matrix is deﬁned to be Ω ∗ :=
[Equation]
p ∗·
[Equation]
s
[Equation]
If Ω ∗ = USU (cid:62) is the eigendecomposition, then U d ∈ R V ×
d is matrix of top d eigenvectors of Ω ∗ .
The matrix Ω ∗ captures substitutability between pairs of words.
Words w and w (cid:48) are substitutable if they have identical conditional probabilities for every context s ∈ S and thus can replace occurrences of each other while still providing meaningful completions.
By deﬁnition, these words satisfy Ω ∗ [ w ] =
Ω ∗ [ w (cid:48) ] .
Such pairs of words were called “free variants” in the work on distributional semantics (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 (11).
Theorem C.1.
The optimal solution f ∗ , Φ ∗ = arg min f, Φ (cid:96) quad ( f, Φ) satisﬁes
[Equation]
If Φ is ﬁxed, then the optimal solution is f ∗ ( s ) =
(ΦΦ
[Equation] [Equation]
[Equation]
Φ p ∗·| s .
Theorem 5.2.
Let f ∗ , Φ ∗ = arg min f, Φ (cid:96)
quad ( f, Φ) .
Then Φ ∗ = BU (cid:62) d , for full rank B ∈ R d × d .
Also, for a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ ∗ , we have (cid:96) T ( f ∗ ) ≤ τ .
Thus f ∗ excels on natural tasks w.r.t.
Φ ∗ , which in turn, is the best d -dimensional projection of Ω ∗ .
Thus words w, w (cid:48) ∈ W that are synonyms (hence substitutable) will satisfy φ ∗ w = φ ∗ w (cid:48) , fulﬁlling the desired property for word embeddings discussed in Deﬁnition 3.2.
We train using the Quad objective and compare its performance to a similarly trained language model in Section 6, ﬁnding
Quad to be reasonably effective.
The goal of testing Quad is not to obtain state-of-the-art results, but to demonstrate that theoretical insights can aid the design of provably effective algorithms.
D P ROOFS D. P ROOFS FOR UNCONSTRAINED LANGUAGE MODELS
Theorem B.1 (Strengthened Theorem 4.1) .
Let { p ·| s } be a language model that is (cid:15) -optimal, i.e.
(cid:96) xent ( { p ·| s } ) − (cid:96) ∗ xent ≤ (cid:15) for some (cid:15) > 0 .
For a classiﬁcation task T that is ( τ, B ) -natural, we have
[Equation]
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , we have
[Equation]
Proof.
The proof has two main steps that we summarize by the following two lemmas.
The ﬁrst one upper bounds the downstream performance on natural tasks with the covariance of errors.
Lemma D.1.
For a language model { p ·| s } , if T is ( τ, B ) -natural,
[Equation]
If T is ( τ, B ) -natural w.r.t.
Φ ∈ R d × V ,
[Equation]
where γ ( · ) and γ Φ ( · ) are from Deﬁnition B.1.
The second lemma upper bounds the covariance of error with the suboptimality of the language model.
Lemma D.2.
For a language model { p ·| s } and classiﬁer v ∈ R V ,
[Equation]
where Σ p L (∆ { p ·| s } )
=
E s ∼ p L (cid:104) ( p ·| s − p ∗·| s )( p
[Equation]
| s − p ∗·|
[Equation]
) (cid:62) (cid:105) as deﬁned in Section B.
Combining the two lemmas, we get the following inequality
[Equation]
where ( a ) uses ﬁrst part of Lemma D.1, ( b ) uses Lemma D.2 and ( c ) uses the (cid:15) -optimality of { p ·| s } .
This proves the ﬁrst part of the result.
The second part can also be proved similarly.
[Equation]
[Table]
[Equation]
[Table]
where ( a ) uses second part of Lemma D.1, ( b ) uses Lemma D.2 and ( c ) uses the (cid:15) -optimality of { p ·| s } .
The proof of the lemmas can be found in Section D.5.
Theorem 4.1.
Let { p ·| s } be a language model that is (cid:15)
-optimal, i.e. (cid:96) xent ( { p ·| s } ) − (cid:96) ∗ xent ≤ (cid:15) , for some (cid:15) > 0 .
For a classiﬁcation task T that is ( τ, B ) -natural, we have
[Equation]
Proof.
This follows from the ﬁrst part of Theorem B.1 if we can also show that γ ( p T ; { p ·| s } ) − 1 ≤ γ ( p T ) − 1 .
For that we use the following lemma that we prove in Section D.5.
Lemma D.3.
For any g : S → R D and p T ∈ ∆ S , we have (cid:107) Σ p L ( g ) −
12 Σ p T ( g )Σ p L ( g ) − 12 (cid:107) 2 ≤ γ ( p T ) −
Instantiating this for g = ∆ { p ·| s } and using Equation (7), we get γ ( p T ; { p ·| s } ) − 1 ≤ γ ( p T ) − 1 , which completes the proof.
D. P ROOFS FOR SOFTMAX LANGUAGE MODELS
Theorem 5.1 (Strengthened Theorem 4.2) .
For a ﬁxed Φ , let f be features from an (cid:15) -optimal
ddimensional softmax language model, i.e. (cid:96) xent ( f, Φ) − (cid:96) ∗ xent (Φ) ≤ (cid:15) , where (cid:96) ∗ xent (Φ) is deﬁned in
Equation (4).
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , we have (cid:96)
[Equation]
(cid:115)
[Equation]
Proof.
Instantiating Lemma D.1 for p ·| s = p f ( s ) , we get
[Equation]
where ( a ) follows from Equation (9) that says γ Φ ( p T ; g f, Φ )
=
γ Φ ( p T ; { p f ( s ) } ) .
We now prove a similar result for the second term in the following lemma
Lemma D.4.
For a ﬁxed Φ and a softmax language model with features f and λ ∈ R d ,
[Equation]
where Σ p L (Φ∆ { p f ( s ) } )
=
E s ∼ p L (cid:104) (Φ p f ( s ) − Φ p ∗ ·| s )(Φ p f ( s ) − Φ p ∗ ·| s ) (cid:62) (cid:105) as deﬁned in Section B.
Using Lemma D.4 directly gives us
[Equation]
, and the (cid:15) -optimality almost completes the proof.
The only thing remaining to show is that (cid:96) T ( { p f ( s ) } ) ≤ (cid:96) T ( g f, Φ ) which follows from the following sequence.
[Equation]
Theorem 4.2.
For a ﬁxed Φ , let f be features from an (cid:15) -optimal
d -dimensional softmax language model, i.e. (cid:96) xent ( f, Φ) − (cid:96) ∗ xent (Φ) ≤ (cid:15) , where (cid:96) ∗ xent (Φ) is deﬁned in Equation (4).
For a classiﬁcation task T that is ( τ, B ) -natural w.r.t.
Φ , we have
[Equation]
Proof.
This result follows directly from Theorem 5.1, if we can also show that γ ( p T ; g f, Φ ) − 1 ≤ γ ( p T ) − 1 just like in the proof of Theorem 4.1.
For that we again use Lemma D.3 with g = Φ∆ { p f ( s ) } and Equation (9) and this completes the proof.
D. P ROOFS FOR S ECTION 4.
We ﬁrst show why Assumption 4.1 is approximately true when word embeddings are gaussian like.
Lemma D.5.
Suppose word embeddings φ w are independent samples from the distribution N ( µ, Σ) .
Then for any θ ∈ R d such that λ 2 = θ (cid:62) Σ θ = O (1) we have that | log( Z θ ) − 12 θ (cid:62) Σ θ − θ (cid:62) µ − log( V ) | ≤ (cid:15) with probability 1 − δ for (cid:15) =
˜ O (cid:16) e λ 2 √ V (cid:17) and δ = 1 − exp( − Ω(log 2 ( V ))) .
Proof.
We ﬁrst note that log( Z θ )
= log (cid:16)(cid:80) w e θ (cid:62) φ w (cid:17) =
θ (cid:62) µ + log (cid:16)(cid:80) w e θ (cid:62) ( φ w − µ ) (cid:17) , thus we can simply deal with the case where φ w are sampled from N (0 , Σ) .
Furthermore the only random variable of interest is X w = θ (cid:62) φ w which is a gaussian variable N (0 , θ (cid:62) Σ θ ) =
N (0 , λ 2 ) .
Thus the problem reduces to showing that for V samples of X w ∼ N (0 , λ 2 ) , log( Z ) is concentrated around λ 2 + log( V ) where Z = (cid:80) w exp( X w ) .
This can be proved similarly to the proof of Lemma 2.
in Arora et al.
It is easy to see that
E
[Equation]
) [exp
[Equation]
X w )]
=
e λ 2 .
However the variable exp( X w ) is 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 ﬁrst observe that E
[ Z ] =
V e 12 λ 2 and
Var[ Z ] ≤ E [exp(2 X w )] =
V e 2 λ 2 .
After conditioning on the event that X w ≤ 12 λ log( V ) and applying Berstein’s inequality just like in Arora et al.
(2016) completes the proof.
Lemma 4.3.
Under Assumption 4.1, any feature map f : S → R d satisﬁes g f, Φ ( s ) =
A f ( s ) + b , for all s ∈ S .
Proof.
Assumption 4.1 gives us that log( Z θ )
= 12 θ (cid:62)
A θ + θ (cid:62) b + c .
We prove this lemma by matching the gradients of log( Z θ ) and the quadratic function on the R.H.S.
[Equation]
Whereas the gradient of the quadratic part is ∇ θ [ 12 θ (cid:62) A θ + θ (cid:62) b + c ]
=
A θ + b .
Matching the two gives us g f, Φ ( s ) =
Φ p f ( s ) =
A f ( s ) + b .
Corollary 4.1.
Using Lemma 4.3, for any (cid:15) -optimal f , as deﬁned in Theorem 4.2, for classiﬁcation tasks that are ( τ, B ) -natural w.r.t.
Φ we have (cid:96) T ( f ) ≤ τ + O ( √ (cid:15) ) .
Proof.
The main idea is that Lemma 4.3 gives us that g f, Φ ( s ) =
A f ( s ) + b and thus any linear function of g f, Φ will also be a linear function of f ( s ) .
And from Theorem 5.1 (or Theorem 4.2), we know that g f, Φ will do well on T , i.e. (cid:96) T ( g f, Φ )
≤ τ + O ( B √ (cid:15) ) .
We formalize the intuition as
[Equation]
This shows that (cid:96) T ( f ) ≤ (cid:96) T ( g f, Φ ) ≤ τ + O ( B √ (cid:15) ) and completes the proof.
D. P ROOFS FOR S ECTION C
Theorem C.1.
The optimal solution f ∗ , Φ ∗ = arg min f, Φ (cid:96) quad ( f, Φ) satisﬁes
[Equation]
= BU
[Equation]
d , for full rank B ∈
[Equation]
×
d f
[Equation]
=
[Equation]
p ∗·| s , for full rank C ∈ R
d × d
If Φ is ﬁxed, then the optimal solution is f ∗ ( s ) =
(ΦΦ (cid:62) )
− 1 / 2 Φ p ∗ ·| s .
Proof.
Given that, (cid:96) quad,s ( θ, Φ) =
− θ (cid:62) Φ p
∗·| s + 12 (cid:107) Φ (cid:62) θ (cid:107) 2 , for a ﬁxed Φ , we deﬁne f ∗ Φ ( s ) =
arg min θ ∈ R d (cid:96) quad,s ( θ, Φ) .
We use the ﬁrst order optimality condition to get f ∗ Φ ( s ) , by using the fact that ∇ θ (cid:96) quad,s ( θ, Φ)
= − Φ p
∗·| s + ΦΦ (cid:62) θ .
Setting the gradient to zero, we f ∗ Φ ( s ) =
(ΦΦ (cid:62) )
− 1 Φ p ∗·| s .
To get the optimal Φ ∗ for this objective, we plug in this expression for f ∗ Φ ( s ) and observing that the optimal solution also satisﬁes Φ ∗ = arg min Φ (cid:96) quad ( f ∗ Φ , Φ) .
[Equation]
Let Φ =
NT V (cid:62) be the SVD.
Then the above objective reduces to (cid:96) quad ( f ∗ Φ , Φ) = −
[Equation]
V V
[Equation]
, Ω ∗ (cid:11)
And hence learning the optimal Φ ∗ reduces to learning an optimal V ∗ such that
[Equation]
We will now show that the best such matrix is the matrix of top d eigenvectors of Ω ∗ , i.e. V ∗ =
U d .
Here we will assume that the eigenvalues of Ω ∗ are all distinct for simplicity of presentation.
First we note that (cid:104) V V (cid:62) , Ω ∗ (cid:105) =
(cid:107) V V (cid:62) Ω ∗ 12 (cid:107) 2 F , where Ω ∗ 12 = US 12 U (cid:62) , with U , U d and S deﬁne in Deﬁnition
[Equation]
This can be shown by the following sequence of steps
[Equation]
Furthermore, we notice that (cid:107) V V (cid:62) Ω ∗ 12 (cid:107) 2 F =
[Equation] [Equation]
− (cid:107)
[Equation]
12 −
[Equation]
V
[Equation] [Equation]
2 F as shown below
[Equation]
Thus we get
[Equation]
Note that V V (cid:62) Ω ∗ 12 has columns that are columns of Ω ∗ 12 projected on the space spanned by columns V .
It is folklore that the best such subspace V ∗ is the subspace spanned by the top d eigenvectors of Ω ∗ 12 , which is the same as top d eigenvectors of Ω ∗ , thus giving us V ∗ V ∗(cid:62) =
U d U (cid:62) d .
Thus we get V ∗ =
U d M for M =
U (cid:62) d V ∗ .
This tells us that the optimal solution Φ ∗ will have SVD of the form Φ ∗ =
N ∗ T ∗ V ∗(cid:62) , thus giving us Φ ∗ = BU (cid:62) d for matrix B = N ∗ T ∗ M (cid:62)
∈ R d × d .
This directly gives f ∗ =
f ∗ Φ ∗ =
(Φ ∗ Φ ∗(cid:62) )
− 1 Φ ∗ p
∗·| s = N ∗ T − 1 V ∗(cid:62)
p ∗·| s = CU (cid:62) d p ∗·| s for C
= N ∗ T ∗− 1 M (cid:62) .
D. P ROOF FOR SUPPORTING LEMMAS
Lemma D.1.
For a language model { p ·| s } , if T is ( τ, B ) -natural,
[Equation]
p
[Equation]
|
[Equation] [Equation]
If T is ( τ, B ) -natural w.r.t.
Φ ∈ R d × V ,
[Equation]
where γ ( · ) and γ Φ ( · ) are from Deﬁnition B.1.
Proof.
We note the following upper bounds on (cid:96) T ( { p ·| s } ) and (cid:96) T ( { Φ p ·| s } ) .
[Equation]
When T is ( τ, B ) -natural, by Deﬁnition 3.1 we know that
[Equation]
{ p ∗·
[Equation]
v ) (cid:105) ≤ τ .
We now upper bound (cid:96) T ( { p ·| s } , v ) using Lemma D.8.
Taking inﬁmum w.r.t.
v ∈ R V , (cid:107) v (cid:107) ∞ ≤ B from the inequality in Lemma D.8.
[Equation]
This, combined with Equation (12)
[Equation]
gives us
[Equation]
Using Lemma D.9 and the deﬁnition of γ ( p T ; { p ·| s } ) in Equation (7), we get that
[Equation]
We have thus successfully transferred the bound from the distribution p T to p L .
Combining this with Equation (14) completes the proof of the ﬁrst part of the lemma.
We now prove the second part of the lemma where we only assume that T is ( τ, B ) -natural w.r.t.
Φ .
Here we instead take the inﬁmum over classiﬁers in the span of Φ in Lemma D.8 to get
[Equation] [Equation]
This, combined with deﬁnition of ( τ, B ) -natural task w.r.t.
Φ and Equation (13) gives us
[Equation]
For the last term, for any v = Φ (cid:62) λ, λ ∈ R
d we notice that
[Equation]
This combined with Equation (16), we get
[Equation]
Lemma D.2.
For a language model { p ·| s } and classiﬁer v ∈ R V ,
[Equation]
where Σ p L ( g ) =
E s ∼ p
L [ g ( s ) g ( s ) (cid:62) ] and ∆ { p ·| s } ( s ) =
p ·| s − p
∗·| s are deﬁned in Section B
Proof.
We ﬁrst note that
[Equation]
We bound v (cid:62) Σ p L (∆ { p ·| s } ) v below
[Equation]
where ( a ) uses Holder’s inequality, ( b ) uses Pinsker’s inequality, (
c
Lemma D.3.
For any g : S → R D and p T ∈ ∆ S , we have (cid:107) Σ
[Equation]
L
[Equation]
g
[Equation]
− 12 Σ
[Equation]
( g
[Equation]
Σ
[Equation]
L
[Equation]
g
[Equation]
−
[Equation]
2 ≤ γ ( p T ) −
Proof.
By deﬁnition of γ ( p T ) , we have that
[Equation]
Thus 1 γ
[Equation]
Σ
[Equation]
g
[Equation]
(cid:60) Σ
[Equation]
g
[Equation]
and hence
[Equation]
Σ p L ( g ) −
12 Σ p T ( g )Σ p L ( g ) − 12 , which is equivalent to 1 γ ( p T )
I D (cid:60) Σ
p L ( g ) −
12 Σ p T ( g )Σ p L ( g ) − 12 .
This ﬁnishes the proof.
Lemma D.4.
For a ﬁxed Φ , a softmax language model with features f and λ ∈ R d ,
[Equation]
where Σ p
[Equation]
f
[Equation]
s
[Equation]
Φ p
[Equation]
as deﬁned in Section B.
Proof.
We start by nothing that
[Equation]
We will use Lemma D.6 to bound each term on the right hand side
[Equation]
which essentially bounds the norm of the gradient of (cid:96) xent ,s at f ( s ) when f ( s ) is an almost optimal for s ∈ S .
Notice that
[Equation]
where ( a ) follows from Lemma D.6.
This completes the proof.
Lemma D.6.
For s ∈ S and embedding f ( s ) ∈ R d , we have
[Equation]
Proof.
Since we are assuming Φ to be ﬁxed
[Equation]
we will abuse notation and say (cid:96) xent ,s ( θ ) :
= (cid:96) xent ,s ( θ, Φ) .
All gradients are w
[Equation]
r
[Equation]
t. θ .
Before we get to the main proof
[Equation]
we compute the gradient and hessian of (cid:96) xent ,s ( θ ) w
[Equation]
r
[Equation]
t. θ .
The gradient is
[Equation]
Similarly the Hessian can be computed
[Equation]
[Equation]
Where Cov w ∼ p θ [ φ w ] denotes the covariance of the word embeddings φ w when measured w.r.t.
the distribution p θ .
This directly gives us that ∇ 2 (cid:96) xent ,s ( θ ) (cid:60) 0 , since the covariance is always psd, and thus (cid:96) xent ,s is convex in θ .
Using the closed form expression for ∇ (cid:96) xent ,s , we note that the quantity we wish to upper bound can be rewritten as | λ (cid:62) (Φ p f ( s ) − Φ p ∗·| s ) | 2 = | λ (cid:62) ∇ (cid:96) xent ,s ( f ( s )) | 2 .
Furthermore, using the deﬁnition of the Hessian, it is not hard to see for some λ, ˜ θ ∈ R
d that λ (cid:62) ∇ 2 (cid:96) xent ,s (˜ θ ) λ =
Cov w ∼ p ˜ θ
[ λ (cid:62) φ w ] ≤
E w ∼ p ˜ θ
[( λ (cid:62) φ w ) 2 ] ≤ (cid:107) Φ (cid:62) λ (cid:107) 2 ∞ .
We use the following lemma that can exploit the above observations.
Lemma D.7.
If a function (cid:96) : R d → R and λ ∈ R d satisfy λ (cid:62) ∇ 2 (cid:96) (˜ θ ) λ ≤ L, ∀ ˜ θ ∈ R d ( Lsmoothness in the direction of λ ) and if (cid:96) ∗ = inf θ ∈ R d (cid:96) ( θ ) , then | λ (cid:62) ∇ (cid:96) ( θ ) | 2 ≤ 2 L ( (cid:96) ( θ ) − (cid:96) ∗ )
We ﬁrst use this lemma for (cid:96) xent ,s to complete the proof with L = (cid:107) Φ (cid:62) λ (cid:107) 2 ∞ .
The lemma gives us that | λ (cid:62) ∇ (cid:96) xent ,s ( f ( s ))
| 2 ≤ 2 (cid:107) Φ (cid:62)
λ (cid:107) 2 ∞ ( (cid:96) xent ,s ( f ( s )) − (cid:96) ∗ xent ,s ) ≤ 2 (cid:107) Φ (cid:62)
λ (cid:107) 2 ∞ (cid:15)
s .
Combining this with the expression for the gradient computed earlier, we get | λ (cid:62) (Φ p f ( s ) − Φ p ∗·| s ) | 2 ≤ 2 (cid:107) Φ (cid:62) λ (cid:107) 2 ∞ (cid:15) s , thus completing the proof of the main lemma.
We now prove the lemma.
Proof of Lemma D.7.
This is a variant of a classical result used in optimization and we prove it here for completeness.
For any η ∈ R we have
[Equation]
where ( a ) follows from the deﬁnition of inﬁmum and ( b ) follows from Taylor’s expansion and ( c ) follows from the smoothness condition in the statement of the lemma.
Picking η = λ (cid:62) ∇ (cid:96) ( θ ) L gives us (cid:96)
[Equation]
θ
[Equation]
− (cid:96) ∗ ≥ 12 L | λ (cid:62) ∇ (cid:96) ( θ ) | 2 , thus completing the proof.
Lemma D.8.
For any task T and classiﬁer v ∈ R V and predicted probabilities { p ·| s } (cid:96) T ( { p ·| s } , v ) ≤ (cid:96) T ( { p ∗ ·| s } , v ) + (cid:113) v (cid:62) Σ p T (∆ { p ·| s } ) v where Σ p T ( g ) =
E s ∼ p T [ g ( s ) g ( s ) (cid:62) ] and ∆ { p ·| s } ( s ) =
p ·| s − p
∗·| s are deﬁned in Section B.
Proof.
The following sequence of inequalities proves it
[Equation]
where ( a ) follows from 1-lipschitzness of (cid:96) , ( b ) follows from Jensen’s inequality.
Lemma D.9.
For matrices X , Y ∈ R D × D s.t.
X , Y (cid:60)
and Y is full rank, we have that max a ∈ R D , 0 < (cid:107) a (cid:107)≤ λ a (cid:62)
Xa a (cid:62) Y a = (cid:107) Y − 12 XY − 12 (cid:107) 2 for any norm (cid:107) · (cid:107) .
Proof.
Note that a (cid:62)
Xa a (cid:62) Y a is independent of the scaling of a .
The following sequence of inequalities completes the proof
[Equation]
E E XPERIMENT D ETAILS
For all experiments, we use the 117M parameter “small” GPT-2 model proposed in Radford et al.
(2019) and implemented in HuggingFace Wolf et al.
We use the standard learning rate schedule and architecture provided in the initial publication.
To learn a model on IMDb, we use a context size of 512 BPE tokens, and for the Amazon reviews dataset McAuley et al.
(2015), we use the standard context length of 1,024 BPE tokens.
E. S OLVING DOWNSTREAM TASKS USING f AND Φ p f
Details about word subsets: For all of the results presented in Table 1, we use a pre-trained GPT2 model.
For SST, we use the prompt “This movie is ” when indicated.
For AG News, we use the prompt “This article is 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’, ’ﬁnancial’, ’market’, ’foreign’, ’technology’, ’international’, ’stock’, ’company’, ’tech’, ’technologies’.
For SST, this subset is: ’:)’, ’:(’, ’great’, ’charming’, ’ﬂawed’, ’classic’, ’interesting’, ’boring’, ’sad’, ’happy’, ’terrible’, ’fantastic’, ’exciting’, ’strong’.
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 ﬁlter to use only words that encode to a single BPE token.
For AG News, the class words we use are: ’foreign’, ’sports’, ’ﬁnancial’, ’scientiﬁc’.
For SST, the class words we use are ‘ :)’ and ‘ :(’.
Tests on additional datasets: We also test the performance of GPT-2 frozen embeddings f and the conditional mean embeddings Φ p f on the DBPedia (Auer et al., 2007), Yahoo Answers (Zhang et al., 2015), TREC (Li & Roth, 2002), IMDb (Maas et al., 2011), Customer Review (CR) (Hu & Liu, 2004), and MPQA polarity (Wilson & Wiebe, 2003) datasets in Table 3.
We limited the training set size to 200K for larger datasets (i.e., DBPedia and Yahoo Answers).
For CR and MPQA, we created train-test splits with 75-25 percentage random splits of the data.
We ﬁnd that Φ p f ( s ) consistently has comparable performance to f ( s ) across non-sentiment and sentiment downstream classiﬁcation tasks.
We include results using a bag of n -grams (BoNG) and
Sentiment Neuron (mLSTM) (Radford et al., 2017), and we note that using 768-dimensional features is more sample efﬁcient than BoNG.
For sentiment tasks, adding a prompt 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-efﬁcient than the standard 768-dimensional ones.
We also include results using the pre-trained BERT base cased model (Devlin et al., 2018; Wolf et al., 2019), using the embedding at the ﬁrst token as input to the downstream task.
We also tried using the mean and last token and found that the token is best.
Moreover, ﬁrst is what is extracted in the traditional usage of BERT on downstream tasks.
[Table]
T ESTING Q UAD OBJECTIVE ON A LARGER DATASET
We test two models with the same parametrization and initializations, one trained using our Quad objective and another trained with the standard language modeling objective using the Amazon product review dataset (McAuley et al., 2015) instead of IMDB.
We slightly modify the standard architecture of GPT-2 to generate Tables 2 and 4.
We add a single linear layer after the Transformer to add expressivity.
Furthermore, instead of tying the input and output embeddings, we learn them separately so that f and Φ are independent functions.
We ﬁx the input embeddings and the positional embeddings to be the parameters from the pre-trained GPT-2.
We initialize Φ , the output embeddings, using the singular vectors of the pre-trained word embeddings Φ .
Given our parameterization, initializing with the singular vectors is as expressive as initializing with the pretrained embeddings Φ themselves; however it lends a better optimization landscape and speeds up training for our new objective.
We observe that even on a large dataset, training using Quad yields comparable performance to the language model on the SST task.
Furthermore, adding a prompt consistently improves performance for both objectives. Task f [Equation] s
[Equation]
KL
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
E
L EARNING THE QUADRATIC APPROXIMATION OF THE LOG - PARTITION FUNCTION
In Assumption 4.1, we assert that there is a quadratic ﬁt for the log partition function, which allows us to show in Lemma 4.3 that a linear relation holds between f ∗ and Φ p f ∗ .
We validate these theoretical ﬁndings by ﬁtting a quadratic function to the log partition function for a subset of embeddings from the IMDb, SST, and AG News datasets (Figure 1).
Here, we describe how we learned A , b and c .
To ensure A is symmetric and positive semi-deﬁnite as required, we parametrize A = UU T .
Let µ θ =
Φ p θ .
We minimize the following objective function:
[Equation]
In practice, we train only on the regression loss (i.e., λ 1 = 0 , λ 2 = 1 ) for the most promising results.
We use 20,000 examples from a mix of IMDb, SST, and AG News embeddings as the training set.
We used the Adam (Kingma & Ba, 2014) optimizer with learning rate 1e-3 for U and learning rate 1e-4 for b and c .
We decayed the learning rate every 50 steps by a factor of 0.1.
We found the ﬁt after 8 epochs of training.
We further demonstrate the quality of the learned ﬁt by plotting the true log partition and estimated log partition function for embeddings from other datasets in Figure
. E. E XPERIMENTALLY C HECKING T HEOREM 4.
Theorem 4.2 can be informally summarized as stating that an (cid:15) suboptimality in the cross-entropy of a d -dimensional language model propagates to a √ (cid:15) increase in the logistic loss.
We note that the τ, B, and γ ( p T ) factors are ﬁxed for a given pre-training corpus and downstream task, so we can empirically test if this square root relationship holds in practice.
[Figure]
We train a 117M parameter GPT-2 architecture from scratch on the IMDb dataset or the Amazon dataset, maintaining checkpoints during training.
For each checkpoint, we measure the crossentropy of the model on the validation set as well as the performance of the conditional mean features g f, Φ on SST-2.
Plotting these values together yields Figure 3.
We furthermore ﬁt a square root trend, shown in red, to these points.
We learn a, b, c such that y = a √ x − b + c , where y is the logistic loss and x is the cross-entropy loss.
To do so, we perform a grid search over 100 evenly spaced valid values of b , and for each value of b , we perform linear regression on √ x − b to ﬁnd a and c .
We choose the a, b, c that maximizes the r -value of the regression.
We note that Theorem 4.2 only provides an upper bound on the logistic loss, but we provide this experiment as validation that the square-root trend is observable in practice.