T HE I NDUCTIVE B IAS OF I N -C ONTEXT L EARNING : R ETHINKING P RETRAINING E XAMPLE D ESIGN
Anonymous authors
Paper under double-blind review
A BSTRACT
Pretraining Neural Language Models
NLMs) over a large corpus involves chunking the text into training examples which are contiguous text segments of sizes processable by the neural architecture
We highlight a bias introduced by this common practice: we prove that the pretrained NLM can model much stronger dependencies between text segments that appeared in the same training example than it can between text segments that appeared in different training examples
This intuitive result has a twofold role
First, it formalizes the motivation behind a broad line of recent successful NLM training heuristics, proposed for the pretraining and ﬁne
-tuning stages, which do not necessarily appear related at ﬁrst glance
Second
, our result clearly indicates further improvements to be made in NLM pretraining for the beneﬁt of Natural Language
Understanding tasks.
As an example, we show that including semantically related non
-neighboring sentences in the same pretraining example yields improved sentence representations and open domain question answering abilities.
This theoretically motivated degree of freedom for “pretraining example design" indicates new training schemes for self-improving representations.
I NTRODUCTION
Beyond excelling in their core task of pure language modeling, modern Neural Language Models (NLMs) show impressive zero- and few-shot abilities in more general Natural Language Understanding (NLU) tasks (Brown et al., 2020).
This implies that the training corpus contains the information required for performing such tasks, and moreover it implies that the common pretraining process grants the trained NLM some access to these higher level capabilities.
In this paper, we highlight a connection between the quality of the emergent NLU capabilities and a basic component in the NLM training scheme: the process of segmenting the corpus into training examples.
Speciﬁcally, NLMs self-train over huge training corpora (typically, billions to trillions of words).
A basic, automatic, operation in the training pipeline is to segment these corpora into training examples: contiguous text chunks of sizes processable by the neural architecture (typically, up to thousands of words).
We formalize an expressivity bias that this segmentation process introduces, to be referred to as the in-context bias , which directly affects the NLM’s ability to integrate cross-corpus information.
We show that the NLM can model much stronger dependencies between sentences that were shown together at least once in-context, i.e. , in the same training example, than between sentences that were never shown together in the same input.
This inductive bias may be good for language modeling, but it implies that NLU capabilities that involve integrating information from different examples across the corpus (see, e.g. , ﬁgure 1), are under-favored by design in the current setting.
Thus, if one sentence in the corpus can elucidate the meaning of another sentence (e.g., deﬁnes a hard concept or provides auxiliary information), our result implies that a model that saw them in different training examples will enjoy this elucidation less than a model that saw them in the same training example.
While standard approximation results examine the expressivity of an architecture over a single input, our theoretical approach pertains to the entire training process, and examines the expressive capacity of the resultant NLM with respect to the training set.
Therefore, our approximation result ties an optimization parameter (the learning-rate) to the regular NLM architecture expressivity parameters (depth, width).
Intuitively, sentences that were never shown in the same input can only access each other via the weights of the network during training.
The mechanism for “storing" information in the network involves a very small learning-rate term η ; our analysis formalizes and quantiﬁes an “expressivity toll" that the model pays when making use of such harder-to-access stored information.
We employ the tool of a function’s separation rank with respect to subsets of its variables, which quantiﬁes its ability to model input dependencies between these subsets.
The separation rank was employed for analyzing the dependencies modeled by convolutional (Cohen & Shashua, 2017), recurrent (Levine et al., 2018), and self-attention (Levine et al., 2020) networks with respect to a single input example.
In order to analyze an NLM’s ability to model dependencies between different training examples, we reﬁne the usage of this measure in two manners: (1) we introduce the ε -separation rank, which measures the effective ability of a function to model dependencies in a ﬁnite precision setting, and (2) we modify the separation rank such that it can account for the more intricate mechanism of mixing between variables that occurs in the sequential case.
Speciﬁcally, we upper bound the log of the separation rank of a depth L width d x self-attention based
NLM, with respect to two sentences that are shown in its input, by ˜ O
( d x L ) , and prove that this bound is tight.
On the other hand, we upper bound this measure with respect to two sentences that were never shown in the same input by ˜ O
( d x [ L − 0 .
5 log 3 ( η − 1 )]) .
Given common learning-rate values of η ∈ [10 − 6 , 10 − 4 ] , this implies a guaranteed “depth deﬁcit" of ∼ 6 layers for modeling dependencies between sentences that are not seen in the same training example.
After the presentation of our results, we point at empirical evidence that imply that this depth deﬁcit is more signiﬁcant, and may behave like a fraction of L .
We leave attempts to tighten the depth deﬁcit estimates to future work.
1. T HE IN - CONTEXT BIAS DRIVES A VARIETY OF EXISTING APPROACHES
Several recent works intuitively rely on the above formalized in-context expressivity bias in different manners, and signiﬁcantly improve both task-speciﬁc training and pretraining of NLMs.
Gao et al.
(2020) advance the frontier in k -shot learning via ﬁnetuning.
They show that by concatenating several related training examples per input, instead of using standard ﬁne-tuning practice of one example per input, the k -shot performance on sentence similarity tasks is considerably boosted.
Another example was pointed out in Humeau et al.
Thakur et al.
(2020): when training for sentence similarity tasks, including both sentences in the same input leads to a performance gain of around 10 points relative to separately encoding each sentence.
In the challenging setting of open-domain question answering, Izacard & Grave (2020) jointly attend to all documents that may contain the answer, and show large gains relative to prior methods that consider these documents in separate forward passes.
Turning our focus to methods that leverage the in-context bias for improved pretraining, the most straightforward effort is a body of work aimed at reducing the quadratic dependence of the Transformer computation on input sequence length (Tay et al., 2020).
While allowing for more text in-context during training, this does not improve the model’s ability to integrate text across different documents in the corpus.
The following approaches take a further step and enable direct cross-corpus connections during pretraining.
Lewis et al.
(2020) attend to related documents when maximizing the likelihood of a target document.
The scope of related documents is restricted by meta-data: taken from the same Wikipedia entry as the input, or published on the same date.
Guu et al.
(2020) expand the scope of the related documents, by training a Knowledge-Retrieval model that has access to the entire Wikipedia corpus.
They retrieve several related documents per target document, but condition on each related document independently.
Outside of the natural language domain, Rao et al.
(2021) train a Transformer based protein-LM that receives multiple related protein sequences in-context.
Their protein-LM surpasses previous methods which process one sequence per input by a wide margin, with signiﬁcant parameter efﬁciency.
1. L EVERAGING THE IN - CONTEXT BIAS FOR NLU ORIENTED TRAINING
Though the in-context bias is intuitive, the above subsection surveys recent advances that leverage it in non-trivial manners.
Having formalized the theoretical advantage for in-context integration of related text, the roots of the above successes can be uniﬁed, and importantly, new methods for tilting the pretraining bias towards NLU tasks are indicated.
Following the presentation of our theoretical results in section 2, we detail in section 3 two controlled setting exempliﬁcations of new methods that directly leverage the in-context bias.
Our ﬁrst experiment augments the Task Adaptive PreTraining (TAPT) setting of Gururangan et al.
(2020), in which an NLM that was pretrained on a general corpus continues pretraining (with its original objective) on the training set of an NLU task.
We perform TAPT on the SentEval sentence similarity benchmark (Conneau & Kiela, 2018), and during TAPT introduce the following augmentation: along with SentEval sentences, we simultaneously pretrain on related sentences from
Wikipedia, the general pretraining corpus.
The related sentences are found via k-Nearest Neighbors (kNN) search between the embeddings of SentEval examples and all Wikipedia sentences; we thus dub this approach kNN-TAPT .
Importantly, during kNN-TAPT, each input includes a training example from the task, appended in-context by its Wikipedia neighbors.
We demonstrate signiﬁcant gains of the kNN-TAPT over regular TAPT on SentEval sentence similarity tasks.
A dedicated ablation study shows the signiﬁcance of adding the general corpus neighbors in-context, versus in separate training examples, during kNN-TAPT.
[Figure]
Our second experiment introduces a task-independent pretraining phase, dubbed kNN-Pretraining.
As in kNN-TAPT, we group together sentences with similar sentence representations in the same training example, but in kNN-Pretraining we use only sentences from the general pretraining corpus.
This can be viewed as a sentence-focused variation of the above surveyed pretraining schemes in Lewis et al.
and Guu et al.
(2020), who operate on full documents (up to 512 each) and are therefore forced to make compromises ( e.g. , only two related documents in the same input).
Figure 1 shows that after regular pretraining for 200 K steps on Wikipedia, the zero-shot closed book performance of 3 different randomly initialized GPT2-medium models ( 345 M parameters) on open domain questions from Wikipedia (Kwiatkowski et al., 2019) is very low (correct on less thanquestions out of 20,000 in the evaluation set).
Adding kNN-Pretraining for 20 K steps raises performance signiﬁcantly (correct on roughly 250 questions in the evaluation set), reﬂecting the enhanced ability to integrate knowledge from related sentences, acquired via the in-context bias.
In summary, our main contributions are:
• We formally establish the in-context bias : information within pretraining examples is better represented than information integrated across pretraining examples.
• We ask and answer a new type of network expressivity question: how expressive is a network with respect to examples seen during its training process?
• We demonstrate that in-context bias motivated “pretraining example design" elicits better representations from the same data: kNN-Pretraining improves on several NLU tasks.
T HEORETICAL ANALYSIS : T HE IN - CONTEXT BIAS OF SELF - ATTENTION
In this section, we consider the entire NLM training procedure as a functional that receives an unlabeled training corpus and outputs a trained NLM.
Our analysis focuses on the corpus segmentation into training examples as a hyper-parameter of this functional.
We reduce the high-level notion of representing “cross-corpus correlations" to a quintessential case study: we quantify the NLM’s ability to model dependencies between two sentences that appear in the same training example (the in-context representation) and in different training examples (the sequential representation).
We believe that the in-context bias can be shown to exist in a broad range of architectures, but we focus on self-attention since almost all modern NLMs are based on the Transformer architecture of Vaswani et al.
Our theoretical framework is based on that of Levine et al.
; Wies et al.
(2021), who analyze a simpliﬁed, theoretically accessible, self-attention network.
They study the expressivity of this self-attention architecture with respect to its input, and use a measure of a multivariate function’s ability to correlate two subsets of its variable set, referred to as the separation rank.
The analyzed framework captures the connectivity of self-attention but omits its softmax and
ReLU non-linearities (see eq.
1 below).
We refer the reader to Levine et al.
; Wies et al.
(2021) for a discussion on the impact of these relaxations.
Essentially, they are shown to weaken the overall network power but still allow a meaningful comparison of the self-attention integration abilities.
Importantly, both works derive unforeseen theoretical conclusions from analyses of the separation rank measure for this architecture class, and then provide extensive empirical corroboration for their manifestation in common Transformer architectures, reinforcing the relevance of this setting.
In the following, we describe in section 2.1 the analyzed in-context and sequential self-attention representations of two sentences.
Then, in section 2.2, we present the separation rank, which we use in section 2.3 for quantifying the advantage of in-context representations versus sequential ones.
T HE ANALYZED IN - CONTEXT AND SEQUENTIAL REPRESENTATIONS
For an input sequence of N embedding vectors { x j ∈ R d x } Nj =1 , denote the function realized by the analyzed H -headed
depth-
L width- d x Transformer architecture at output location i ∈ [ N ] by: g i,L,d x W (cid:0) x 1 , ..., x N (cid:1)
∈ R d x , where W stands for learned parameters, recursively deﬁned:
[Equation]
where W is composed of Key, Query, Value and Output matrices:
∀ l ∈ [ L ] , h ∈
[ H ] ,
W K ,l,h , W Q ,l,h , W V ,l,h , ( W O ,l,h ) (cid:62) ∈ R
d a × d x , where we assume the standard choice d a = d x /H .
For a word from vocabulary of size V , w ∈
[ V ] , the translation into the Transformer dimension is done via a mapping E M V : [ V ] →
[Equation]
where M V ∈ R d x × V is the learned vocabulary embedding matrix, and (cid:0) M V (cid:1) w is its w th column, also referred to as the learned word embedding for w .
Overall, the function of the analyzed Transformer over a sequence of N words { w j ∈ [ V ] } Nj =1 can be written by composing eqs.
1 and 2:
[Equation]
i,L,
[Equation]
W ,
[Equation]
For simplicity of presentation, we examine two sentences S 1 and S 2 of equal length N : S 1 = { w j 1 } Nj =1 and S 2 = { w j 2 } Nj =1 .
The in-context representation simply concatenates both in the input:
[Equation]
For the sequential approach, we consider a setup in which sentence S 1 is inserted into the network at training step t and sentence S 2 is inserted into the network at training step t
+ 1 .
The output of the network at training step t is therefore: y i,L,d x W t ,M V t ( S 1 ) , where W t , M V t stand for all the learned weights before training step t .
Focusing on autoregressive NLMs for simplicity of presentation (the analysis holds for bidirectional NLMs as well), the log-likelihood loss is given by L ( S 1 )
= − (cid:80) Nj =1 log (cid:20)(cid:16) softmax (cid:110)(cid:0) M V t (cid:1)
[Equation] [Equation]
, and the gradient update for any learned weight θ ∈ {W t , M V t } is: θ t +1 ( S 1 ; η )
= θ t − η · ∂ L ( S 1 ) /
∂θ t , where η is the learning rate.
Accordingly, the analyzed sequential representation is the network output after training step t + 1 :
[Equation]
In practice, two relevant non-neighboring sentences are not necessarily shown in consecutive pretraining steps.
In comparison to the realistic scenario of S 1 and S 2 appearing at any training step, this simpliﬁcations tilts the representation in favor of modeling high correlations between S 1 and
S 2 .
Thus, by upper bounding the ability to correlate S 1 and S 2 in the setting of eq. 5 (as we do in section 2.3), we establish an inherent limitation of the network to access information that was stored in its weights via the gradient update mechanism.
In the next subsection, we present our approach for measuring a network’s ability to correlate two sentences seen during training, which we will use in order to separate between the in-context and sequential settings.
2. A MEASURE FOR MODELING IN - CONTEXT AND SEQUENTIAL DEPENDENCIES
In this section, we reﬁne the separation rank, used in prior work in order to analyze the dependencies between two sentences appended in-context.
In section 2.2.1 we present the separation rank and introduce a ﬁnite precision reﬁnement of it, referred to as the effective separation rank , which helps to elucidate the degradation in integration ability caused by the gradient update mechanism.
In section 2.2.2 we point at a structural problem in employing the separation rank in the same manner in which it was employed in prior work that analyzed only architecture expressivity, and introduce the the sequential separation rank , meaningful for
2.2. T HE EFFECTIVE SEPARATION RANK
The separation rank has been established as a measure of dependencies modeled by deep convolutional, recurrent, and self attention network (Cohen & Shashua, 2017; Levine et al., 2018; 2020).
For a function y ( A, B ) over variables A = { a j ∈ X} Mj =1 and B = { b j ∈ X} Mj =1 , the separation rank w.r.t.
( A, B ) is the minimal number of summands that together sum up to equal y ( A, B ) , where each summand is multiplicatively separable w.r.t.
( A, B ) , i.e. , is equal to a product of two functions – one that intakes only A variables and another that intakes only B variables.
Formally, the separation rank of y :
X 2 M → R w.r.t.
( A, B ) is deﬁned as follows:
[Equation]
If the separation rank of a function w.r.t.
( A, B ) is 1 , it is multiplicatively separable w.r.t.
( A, B ) , meaning it cannot take into account consistency between A and B .
The higher sep ( A,B ) ( y ) is, the farther y is from this situation, i.e. , the more it models dependency between A and B .
We will further make use of the effective separation rank :
[Equation]
In words, if a function has a high separation rank, but it can be approximated up to error ε by a function with a low separation rank, then it has a low ε -separation rank.
Prior works compare two functions by establishing the differences between their separation ranks .
In principle, these differences could manifest only in irrelevant magnitudes (if many of the summands in the separation rank deﬁnition are negligibly small for the function with the higher separation rank, for example).
The effective separation rank is key to our analysis because we rely on the fact that information on past examples is stored in the network weights in a small magnitude (due to a small learning-rate).
We show in section 2.3 that much of the integration between text segments from different training examples occurs in very small magnitudes due to high powers of the learning rate, limiting the effective integration, as measured by the ε -separation rank.
Our techniques for bounding the ε -separation rank are extendable to prior works, and while these did not examine the gradient update mechanism, their results can be reinforced due to the guarantees of this introduced measure.
T HE SEQUENTIAL SEPARATION RANK
Levine et al.
(2020), who were the ﬁrst to apply the separation rank to functions realized by Transformer architectures, studied classical architecture expressivity questions which apply only to the in-context representation.
Accordingly, they analyzed only the separation rank of g i,L,d x , deﬁned in eq. 1, and the input variables considered for calculating the separation rank were the word embedding vectors.
A fundamental difﬁculty arises when attempting to directly apply this method to the sequential representation: the word embedding vectors are learned parameters of the architecture.
In the sequential case, when the second sentence S 2 is introduced after the calculation at time-step t , the vectors used to describe it, if we were to follow prior practice, would already have depended on S 1 .
In order to meaningfully measure the integration ability of two sentences across examples in the presence of the gradient update mechanism, we introduce an auxiliary sentence association layer with new variables: a , b ∈ R d x , which explicitly associates each employed vocabulary embedding vector with the sentence index s ∈ { 1 , 2 } that invoked its usage:
[Equation]
where (cid:12) denotes element-wise multiplication.
We deﬁne the sentence association operation over the analyzed representations, denoted Z y ( a , b ) with y ∈ { y i,L,d x in-context ( S 1 , S 2 ) , y i,L,d x ,η sequential ( S 1 , S 2 ) }
(eqs.or 5), to be the application of the sentence association layer of eq.
8 to all uses of the input embedding layer during the computation of y .
Meaning, for both mechanisms, that chosen word embeddings are marked with the identity of the sentence that invoked them.
Finally, we deﬁne the following specialization of the separation rank measure to our setting, referred to as the sequential separation rank of y ∈ { y i,L,d x in-context ( S 1 , S 2 ) , y i,L,d x ,η sequential ( S 1 , S 2 ) } :
[Equation]
[Equation]
Clearly, when the introduced variables are vectors of 1 , the auxiliary layer in eq.
8 is the identity operation and so Z y ( 1 , 1 ) =
y for both representations.
More deeply, our expressivity questions query the ability of the in-context and sequential mechanisms to integrate two sets of variables, and
Z y captures the essence of this ability by explicating where each set enters the computation.
In the next subsection, we show that for the in-context case, analyzed in prior work, the introduced measure of the sequential separation rank is asymptotically equal to the previously employed measure of separation w.r.t.
a partition of the input word embeddings (Levine et al., 2020).
Thus, the properties of the existing framework are unchanged under the new deﬁnition.
At the same time, for the sequential case brought forth in this paper, the sequential separation rank considers both the effect of S 1 on the gradient-updated word embedding and the introduction of S 2 into the computation.
In the following section, we make use of both extensions to the separation rank in eqs.
7 and 9 in order to establish the in-context bias.
2. T HE EXPRESSIVE ADVANTAGE OF IN - CONTEXT LEARNING
We show below that the function computed by a self-attention based NLM when inserting sentences
S 1 and S 2 together in its input (the in-context representation) can model more elaborate dependencies between S 1 and S 2 than the function attained when showing S 1 in the input, modifying the network’s weights according to its loss, and then showing S 2 in a subsequent input (the sequential representation).
We begin by stating the following corollary, following from theorem 2 in Levine et al.
and proposition 1 in appendix A, which upper bounds the sequential separation rank of the in-context representation:
Corollary 1.
Let y ( p,i ) ,L,d x in-context be the p ∈
[ d x ] entry of the analyzed in-context representation deﬁned in eq.
Assume that L > log 3 d x .
Then ( ˜ O notation omits log terms: log d x , log L, log H ):
[Equation]
However, the ε -separation rank of the sequential representation is upper bounded by a lower term:
Theorem 1.
(See proof in appendix B).
Let y ( p,i ) ,L,d x ,η sequential be the p ∈
[ d x ] entry of the analyzed sequential representation deﬁned in eq.
Assume that all learned parameters and all gradients are bounded: ∀ θ ∈ {W , M V } :
<
Λ min ≤ | θ | , |
[Equation]
L
[Equation]
S 1 ) /
∂θ | ≤ Λ max , 2 N < d x , and that L > log 3 d x .
Then, ∀ ε > 0 :
[Equation]
Therefore, a gap between upper bounds on the ability to model dependencies between S 1 and S 2 is indicated.
Since the learning rate η is a small term, its log is negative and the gap is in favor of the in-context representation.
The following theorem guarantees that this gap is meaningful, by showing that the higher upper bound (of the in-context case) is tight in terms of effective rank:
Theorem 2.
(See proof in appendix C).
For y ( p,i ) ,L,d x in-context as deﬁned in corollary 1, there exists an assignment of the network weights for which the following holds:
[Equation] [Equation]
where ε =
˜ O (cid:18)(cid:0)
3 L + d x −
[Equation]
L (cid:1)
[Equation]
Notably, corollary 1 and theorem 2 show that for the in-context case, the sequential separation rank asymptotically equals the regular separation rank, validating the relevance of this measure.
We now provide a high level proof sketch that captures the manner in which the theoretical framework of sections 2.1 and 2.2 is used for establishing the above gap (full proof in the appendix).
For the in-context case, notice that each self-attention layer, deﬁned in eq.
1, is a degree 3 polynomial over its 2 N · d x inputs, rendering the whole network a degree 3 L polynomial.
We write this polynomial as a sum over many monomials, and by deﬁnition, the separation rank of any monomial composing the polynomial is 1 .
Since the separation rank of a sum of functions is upper bounded by the sum of their separation ranks, we upper bound the separation rank by the number of these monomials, yielding eq. 10.
The main difference in the sequential representation case is that the S 1 variables affect the computation only via the gradient, so their impact is expected to be limited.
However, considering that
S 2 ﬁrst encounters gradient updated vocabulary matrix entries (cid:0) m V
[Equation]
[Equation] [Equation]
, it appears that both S 1 and S 2 variables enter the self-attention stack via its input, similarly to the in-context case.
So the integration between S 1 and S 2 occurs right from the start, and indeed we show that the separation rank of both representations is similar.
However, since any function of S 1 is accompanied by the learning-rate η , the monomials for which there are many S 1 variables will be multiplied by high powers of η .
This causes many monomials to be negligibly small, and accordingly not to contribute to the ε -separation rank .
By combinatorial considerations we show that the number of monomials that are not attenuated by η (have sufﬁciently large magnitude) yields eq.
The above theorems establish that from an expressivity perspective, the small magnitude of commonly employed learning-rates hinders the ability to integrate information across different training examples.
Speciﬁcally, the established gap implies that the power of the joint representation of two sentences shown in different training examples is upper bounded by that of a network shallower by 0 .
5 log 3 ( η − 1 ) layers that has seen them in the same context.
Common learning-rate values are on the order of η ∈ [10 − 6 , 10 − 4 ] , implying a deﬁcit of ∼ 6 layers in the sequential case.
As shown in in Levine et al.
Tay et al.
(2021), in many practical regimes of network size depth is crucial for expressivity, reinforcing the implications of this gap.
The weaker upper bound, of the sequential case, is not guaranteed to be tight.
This means that theoretically, the sequential representation may in fact be much weaker than what we have proven, e.g. , that showing two sentences in the same context yields a representation that cannot be matched merely by showing them in separate contexts and adding a realistic number of layers.
However,
Roberts et al.
(2020) show evidence supporting our indicated link between architectural parameters and the in-context bias.
They show that when performing open domain question answering tasks (their deﬁned “closed book" setting), a large T5 model that sees only the question performs comparably to smaller models that are allowed to attend to the documents that contain the answer.
This directly implies a certain strength of the sequential mechanism, namely, that information which was seen during training can be accessed via the weights when the model is realistically stronger, as implied by our bounds.
Notably, the large T5 model is 2-4 times the depth of the contrasted smaller models ( 48 versus 12 - 24 layers), suggesting that the upper bound can be tightened to a fraction of L , or that factors that are beyond expressivity also contribute to the in-context bias ( e.g. , optimization, generalization).
Investigation of these aspects is left for future work.
K NN BASED PRETRAINING EXAMPLE DESIGN
Our theoretical analysis quantiﬁes the relation between the small magnitude of the learning rate, and the deﬁciency in the ability to model dependencies between different training examples.
Clearly, small learning-rates are critical for optimization purposes, so the formalized phenomenon should not be solved via high learning-rates during training.
Instead, our analysis makes it clear that if correlations between speciﬁc sentences are important for a given task, appending them in-context yields better representations for the task.
Below, we describe two controlled experiments that demonstrate the importance of this indicated “pretraining example design" degree of freedom.
In both experiments, correlated sentences are identiﬁed via kNN search in their RoBERTa-large sentence representation space (Reimers & Gurevych, 2019), performed using the FAISS library (Johnson et al., 2019).
3. K NN T ASK A DAPTIVE P RE T RAINING
The Task Adaptive PreTraining (TAPT) method, in which an NLM pretrains on the training set of an
NLU task, leads to impressive gains (Gururangan et al., 2020).
Notably, TAPT is most effective after the regular pretraining stage on a general corpus.
This implies that during TAPT, the model generates improved representations by integrating the task related text with the knowledge stored in its weights from the preceding general pretraining phase.
Under this premise, we postulated that performance will improve if we make relevant sentences from the general corpus more available to the model during the TAPT phase.
According to the above analysis, a simple and effective way to bias the model towards representing desired correlations between sentences is to append them in context.
We thus propose the kNN-TAPT phase, in which the training examples are composed of task examples, concatenated with their general corpus neighbors in embedding space.
We applied kNN-TAPT on the SentEval sentence similarity tasks.
Showing similar sentences from Wikipedia is expected to be particularly useful on these tasks, so this is a good experimentation ground to search for effects of the in-context bias.
For each SentEval example, we searched over 100 M Wikipedia sentences and appended in-context neighbors that have embeddings with over 0 .
8 cosine similarity to the SentEval
[Table]
harnesses example embedding, with a special token inserted between different sentences.
We continued until ﬁnding no more neighbors or reaching a maximum of 256 tokens in the RoBERTa vocabulary (Liu et al., 2019).
This search yielded 170 K examples, over which we continued training a pretrained
RoBERTa-base model for 5 epochs, using the ﬁrst epoch for learning-rate warmup and examining peak learning rates of { 1 , 3 , 5 , 7 } · 10 − 5 .
See appendix D for implementation details.
Table 1, shows zero-shot SentEval sentence similarity scores, attained by using the average word embedding of an inserted sentence as its examined sentence representation (shown by Reimers &
Gurevych (2019) to be most meaningful in zero shot).
All models were trained according to the above prescription, besides the baseline RoBERTa which was simply evaluated.
kNN-TAPT improves over regular TAPT, by over 1 point on average, implying that the Wikipedia neighbors are indeed useful to the TAPT stage.
We compared 4 kNN-TAPT variants as an ablations study.
Importantly all variants labeled with kNN-TAPT train on the same training data during the TAPT stage – the SentEval sentence similarity tasks training sets and their Wikipedia nearest neighbors, and differ only in the arrangement of the data into training examples.
The “neighbors" ﬂag relates a SentEval example to its actual neighbors from the kNN search, while the “random" ﬂag relates it to random Wikipedia sentences from the overall neighbors pool attained in the search.
The “in batch" ﬂag implies that related sentences were shown in the same batch, where every training example includes only one sentence from either SentEval or Wikipedia.
In contrast, the “in context" ﬂag implies that related sentences were shown in the same training example.
The weakness of “neighbors, in-batch" implies that the a-priori plausible approach of biasing the model to learn from these Wikipedia neighbors via placing them in the same batch is not nearly as effective as the theoretically motivated in-context approach.
Leading sentence representations employ in-batch techniques (see for example the contrasive setting of Gao et al.
(2021b))
, and this signal strongly suggests developing in-context parallels.
The fact that the original TAPT scheme outperforms the in-batch approaches implies that including the Wikipedia sentences in separate training examples is harmful.
We postulate that this is because training examples that have only
Wikipedia sentences actually dilute the original TAPT signal.
Indeed, by this view, the reason that “random, in-context" performs comparably to TAPT, is that it does not dilute the original TAPT signal – every training example includes a SentEval example.
Overall, the clear advantage of the “neighbors, in-context" kNN-TAPT variant encourages leveraging the in-context bias for TAPT in further tasks.
3. K NN P RETRAINING
We extended the above to more general kNN-Pretraining , designing pretraining examples with related non-neighboring sentences given only the general pretraining corpus.
kNN-Pretraining is also motivated by the kNN-LM results of Khandelwal et al.
(2019), who show signiﬁcant beneﬁts of using nearest neighbors in representation space at inference time.
Their results exemplify the potential impact of integrating cross-corpus related examples; our kNN-Pretraining approach provably biases the model to learn these correlations at pretraining time , via the in-context bias.
Speciﬁcally, we performed kNN search over Wikipedia sentences for every sentence in Wikipedia, and created each training example similarly to the protocol in the previous subsection.
During kNN-Pretraining, half of the batch contained regular pretraining examples and half contained the prepared kNN examples, in order to retain longer ranged LM abilities.
To examine the effect of kNNPretraining, we pretrained GPT-base and GPT-medium ( 110 M and 345 M parameters) architectures from scratch over Wikipedia in the regular pretraining scheme, and switched to kNN-Pretraining at two different points during pretraining ( 200 K and 400 K).
The training examples were of maximal size 256 , and the batch size was 128 for the GPT-medium models and 256 for the GPT-base models.
In order to directly probe the acquired ability to integrate non-neighboring sentences, we evaluated the resultant models on the very challenging setup of zero-shot closed-book open domain question answering.
In this setup, the unidirectional pretrained model decodes an answer conditioned on the given open ended question.
We evaluated the models on questions from the Natural Questions (NQ) benchmark (Kwiatkowski et al., 2019), using the same phrasing employed in Brown et al.
(2020), and employing the standard “open-domain” version as used e.g. by Lee et al.
; Asai et al.
; Roberts et al.
NQ is composed of questions that have answers within Wikipedia, our pretraining corpus.
kNN pretraining can imtuitively improve in cases where the passage containing the answer has elucidating nearest neighbors from across wikipedia that would help the model to better internalize the answer, such that it is more accessible to the model in zero shot.
As ﬁguredemonstrates, 3 baseline models, pretrained with the regular scheme, achieve very low F1 <
10 −scores on this task.
In contrast, kNN-Pretraining shows a low-scoring but signiﬁcant improvement.
To increase the credibility of the signal, we evaluated our models on the ﬁrst 20 K examples from the
NQ training set (we tested zero-shot performance, so the training set was not used earlier).
Indeed, the attained F1 scores are low, but they correspond to 100 s of correct answers that the kNN-Pretrained model provide after roughly 10 % of the overall training time, versus much less in the 3 randomly initialized baseline models.
Finally, we include in appendix E NQ scores of models of different sizes when starting kNN-Pretraining at different checkpoints, and in appendix F zero-shot scores on several
GLUE tasks, which demonstrate clear gains of kNN-Pretraining over the baselines.
A P ROOF OF C OROLLARY
Proposition 1.
Let S 1 = (cid:110) w j 1 (cid:111)
N j =1 and S 2 = (cid:110)
w j 2 (cid:111)
N j =1 be two sentences, g i,L,d x W the Transformer operation of y i,L,d x in-context ( S 1 , S 2 ) and M V be a vocabulary embedding matrix.
Then:
[Equation]
Proof.
Assume that sep ([ N ] , [2 N ] \ [ N ]) (cid:16) g
i,L,d x W (cid:17) =
R , then by deﬁnition there exist g 1 , . .
, g R : (cid:0) R d x (cid:1) N → R and g (cid:48) 1 , . . .
, g (cid:48)
R : (cid:0) R d x (cid:1) N
→ R such that for any (cid:8)
x j (cid:9) 2 N j =1 , g i,L,d x W (cid:0) x 1 , . . .
, x 2 N (cid:1) =
R (cid:88) r
= g r (cid:0) x 1 , . . .
, x N (cid:1) g (cid:48) r (cid:0) x N +1 , . . .
, x 2 N (cid:1)
Now, given a , b ∈ R d x , we can write:
[Equation]
[Table]
Clearly, this form of presenting Z y i,L,dx in-context ( S 1 ,S 2 ) ( a , b ) is separable with respect to ( a , b ) , and since it has R summands, we can conclude that:
[Equation]
Corollary 1 now follows from an upper bound on sep ([ N ] , [2 N ] \ [ N ]) (cid:16) g i,L,d x W (cid:17) given in Levine et al.
B U PPER BOUND FOR THE ε - SEPARATION RANK
Deﬁnition 1.
For an expression that can be represented as a sum of some terms,
[Equation]
denote by f + the corresponding sum, but with each term replaced by its absolute value, that is:
[Equation]
and note that by the triangle inequality it holds that:
[Equation]
Theorem 3.
Let y ( p,i ) ,L,d x ,η sequential be be the p ∈
[ d x ] entry of the analyzed sequential representation deﬁned in eq.
Assume that all learned parameters and all gradients are bounded: ∀ θ ∈ {W , M V } : Λ
[Equation]
| θ | , |
[Equation]
L ( S
[Equation]
/ ∂θ | ≤ Λ max for some 0 < Λ min ≤ Λ max , N < d x , η ∈ (0 , 1] , 2(1+ η ) d x η < 3 L , 2 (1 + η ) d 2 x < 3 L , In addition, assume that there exists M ≥ 0 for which it holds that Z + y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) <
M on its domain.
Then:
[Equation]
Proof.
Denote:
[Equation]
The proof outline is as follows:
We start by ﬁnding a representation of Z ( S 2 )
Θ t +1 ( a ,S 1 ; η ) as a sum of terms, where each term is separable with respect to ( a , b ) .
We then turn to ﬁnding a subset of these terms, denoted G , such that the sum of all terms in G is an ε -approximation of Z ( S 2 )
Θ t +1 ( a ,S 1 ; η ) .
Lastly, since it follows from the deﬁnition of the ε -separation rank and the construction of G that ε -sep ( a , b ) (cid:16) Z y p,i,H,L,dx,η sequential
[Equation]
S 1 ,S 2 ) (cid:17) is upper bounded by the cardinality of G (which is the number of summands in the approximation), we ﬁnd an upper bound to | G | , which is therefore an upper bound to ε -sep ( a , b ) (cid:16) Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) (cid:17) as well, which by deﬁnition is equal to ε -seq-sep (cid:16) y p,i,H,L,d x ,η sequential (cid:17) .
S
( S 1 ,S
Following Levine et al.
; Wies et al.
, Z ( S 2 ) Θ t +1 ( a ,S 1 ; η ) ( b ) can be written as:
[Equation]
d a
[Equation]
C
[Equation]
=
[Equation]
(13) where ˜ w j :=
g ( b ) j
+ ηf ( a ) j (cid:12) b (this form follows from g ( b ) j :=
E M V t (cid:16) w j 2 (cid:17) (cid:12) b is the entrywise product of b with the embedding of w j 2 prior to the t
th training step, and f ( a ) j := − ∂ L ( S 1 ; a ) ∂ ( M V t ) wj is the gradient update performed to w j 2 ’s embedding at time t ), the P ( c,h ) r c and Q ( c,h ) r c +1 terms are sums of products of the networks inner (i.e., non-embedding) weights which were also updated with respect to L ( S 1 ; a ) , and for convenience we denote j C ( L )+1 :
= i and P ( C ( L )
+1 ,h ) :
= P (0 ,h ) .
=
[Equation] [Equation]
(cid:88)
h
[Equation]
[ H ]
[Equation]
L
[Equation] [Equation]
(cid:88) r 1 ,...,r C ( L )+1 =
Q (0 ,h ) r 1 ,p ·  
[Equation] [Equation]
(cid:89) c
= (cid:68) P ( c,h ) r c , g ( b )
j c
[Equation]
( a ) j c (cid:12) b (cid:69) 
 ·  
[Equation]
Separating to vocab-gradient terms and vocab terms:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
Opening to indices:
[Equation]
Separating to weights and variables:
[Equation]
where:
[Equation]
:= (cid:88) h ∈ [ H ] [ C ( L )]
d
[Equation]
r 1 ,...,r
[Equation]
( L )
[Equation]
Q (
[Equation]
h )
[Equation]
1 ,p ·    
C ( L )+1 (cid:89) c
=
P ( c,h ) r c ,α c    
C ( L ) (cid:89) c
=
Q ( c,h ) r c +1 ,β c   

Compressing summation to count variable
[Table]
[Equation]

 where:
[Equation]
Pushing in summations on N , only the parity matters:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]

[Table]

[Table]
[Equation]

[Equation]
where each summand is separable with respect to ( a , b ) .
S TEP 2 — A N ε - APPROXIMATION OF Z y p,i,H,L,dx,η SEQUENTIAL ( S 1 ,S 2 )
Now that we have a representation of Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) as a sum of ( a , b ) -separable terms, we turn to ﬁnding a subsum that can approximate Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) up to an ε -precision.
First, let us deﬁne the set of all legal indices in the last sum:
[Equation]

[Equation]
 
And for G ⊆ D , denote:
[Equation]
which is the sum of all terms with indices in G .
Clearly, summing over all possible indices gives us the original expression:
[Equation]
Given ε > 0 , we wish to ﬁnd a subset of the indices, G ⊆ D , such that the sum of all terms whose indices are in G is an ε -approximation of Z y p,i,H,L,dx,η sequential
[Equation]
S 1 ,S 2 ) ( a , b ) .
That is, we are looking for
G ⊆ D such that for all a , b :
[Equation]
Note that since we assume that Z + D ( a , b )
=
Z + y p,i,H,L,dx
[Equation]
η
[Equation]
<
[Equation]
, we can get:
[Equation]
and it follows that it is enough for us to show that:
[Equation]
which is equivalent showing that:
[Equation] [Equation]
under the assumption:
[Equation]
which we make going forward.
This will ensure that Z G is an ε -approximation of
Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) ( a , b ) .
Now, we assume that ∀ θ ∈ Θ : θ, − ∂ L ( S 1 )
∂θ ∈ [Λ min , Λ max ] , and by Levine et al.
(2020), the P s and
Q s in eq.
13 are products of up to L matrices, so each of their coordinates is bounded in (cid:2)
Λ L min , Λ L max (cid:3) and we assume without loss of generality that Λ min ≤ 1 ≤ Λ max (otherwise we could have picked a smaller Λ min and a larger Λ max ).
Then for each ( N A , p , n , e ) ∈ D the following inequalities hold:
[Equation]
1] \ I P
[Equation]
α
[Equation] [Equation]
|
[Equation]
∈
[Equation] [Equation]
max =
[Equation]
N
[Equation]
d
[Equation]
L
[Equation]
2 max
[Table]
[Equation]

[Table]
C
[Equation]
L
[Equation]
d
[Equation] [Equation]
n
[Equation]
N
[Equation]


[Equation]
L
[Equation]
and therefore:
[Equation] [Equation]
N
[Equation]
max
Relaxing the parity constraints inside the brackets and recalling that Λ max ≥ 1 gives us an upper bound:
[Equation]
which we can further bound using lemmas 3 and 4 in Levine et al.
(2020) until we are left with:
[Equation]
On the other hand:
[Equation]
[Table]
N
[Table]
N
[Table]
[Equation]
Combining the upper and lower bound for λ N A , p , n · φ A, p , e · Ψ B, p , n , e , we get that for all G ⊆ D :
[Equation]
so in order to show that (14) holds, it sufﬁces to show that:
[Equation]
Now
[Equation]
we can limit ourselves to subsets G ⊆ D of the form:
[Equation]
and in this case we get that:
[Equation]
Let us deﬁne:
[Equation]
and note that:
[Equation]
and also:
[Equation]
where the inequality is due to lemma 3 in Levine et al.
Hence:
[Equation]
And after rearranging we get that for each T ≥ 0 such that:
[Equation]
In order for Z G ( T ) ( a , b ) to be an ε -approximation of Z y p,i,H,L,dx,η sequential ( S
[Equation]
S 2 ) (
, b ) .
[Table]
In the last step we have found a condition on subsets of indices, such that summing over any subset who meets this condition will yield an ε -approximation of Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) .
We will now ﬁnd a speciﬁc subset who meets this condition, and use it in order to bound ε
[Equation]
from above.
We will focus our attention on T s of the form:
[Equation]
for some s ∈ (cid:0) 0 , e − 1 .
5 (cid:3) which we’ll determine later.
By lemma 9 we have that:
[Equation]
so we can choose:
[Equation]

and since this upholds s ∗ <
e − 1 . 5 , we get that (16) indeed holds for ˜ G
( T ( s ∗ )) , and therefore that
Z G ( T ( s ∗ )) is an ε -approximation of Z y p,i,H,L,dx,η sequential ( S 1 ,S 2 ) .
Now, from lemma 8 we get:
[Equation]
and therefore:
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
B. L EMMAS FOR ESTIMATING THE NUMBER OF COEFFICIENTS
Remark 1.
For brevity and clarity we will use expressions of the form (cid:0)
K KM ,..., KM (cid:1) , regardless of whether K is divisible by M or not.
For the latter case, this expression should actually be:
[Equation]
(cid:19) expressions of the form (cid:0)
K KM ,..., KM +1 (cid:1) should be read as:
[Equation]
(cid:124)
[Equation]
(cid:125)
( K mod M )+1 times (cid:19) and expressions of the form (cid:0)
K KM ,..., KM − 1 (cid:1) should be read as: 
[Equation]
(cid:1) if K mod M ≡
[Equation]
[Table]
Lemma 1.
For all K, M ∈ N , the maximal value of (cid:0)
K a 1 ,...,a M (cid:1) is achieved when ∀ j 1 , j 2 ∈
[ M ] it holds that | a j 1 − a j 2 | ≤ 1 .
Proof.
Let a 1 , . .
. , a M be a sequence of non-negative integers such that a 1 + . . .
+ a M = K
and (cid:0)
K a 1 ,...,a M (cid:1) is maximal.
Assume towards a contradiction there exist j 1 , j 2 ∈
[ M ] such that a
[Equation]
−
[Equation]
> 1 ⇐⇒
[Equation] [Equation]
< 1 , then:
[Equation]
in contrary to the maximality of (cid:0)
K a 1 ,...,a M (cid:1) .
Therefore, ∀ j 1 , j 2 ∈ [ M ] , | a j 1 − a j 2 | ≤ 1 .
Lemma 2.
Let K, M be two ﬁxed natural numbers, η ∈ (0 , 1] and denote
[Equation]
Then for all n ∈ [ K ] ∪ { 0 } :
[Equation]
Proof.
We will prove by induction.
• Base case: n = 0 .
[Equation]
where the second equality is due to the fact that:
[Equation]
as an empty product
[Equation]
and:
[Equation]
Therefore, (17) is true for n = 0 .
• Induction step: Let n ≥ 0 such that (17) holds for n .
[Equation]
Thus, (17) holds for n + 1 , and the proof of the induction step is complete.
Hence, by induction, (17) is correct for all n ∈ [ K ] ∪ { 0 } .
Lemma 3.
Let K, M be two ﬁxed natural numbers, and η ∈ (0 , 1] .
Then the maximum of: (cid:18) K n
[Equation]
η n
[Equation]
n
[Equation]
. , a M (cid:19)(cid:18) K − n b 1 , . . .
, b M (cid:19) is achieved when:
[Equation]
and for all i ∈ [ M ] , (cid:12)(cid:12) a i − nM (cid:12)(cid:12) ≤ 1 and (cid:12)(cid:12) b
i − K − nM (cid:12)(cid:12) ≤ 1 .
Proof.
From lemma 1 we know that a multinumial coefﬁcient reaches its maximum when the sum is evenly distributed between all indices.
Since the a i s and b i can be chosen independently of each other given K and n , we may assume without loss of generality that no matter the value of n , for all i ∈ [ M ] , it holds that (cid:12)(cid:12)
a i − nM (cid:12)(cid:12) ≤ 1 and (cid:12)(cid:12)
b i − K − nM (cid:12)(cid:12) ≤ 1 .
Deﬁne S as in lemma 2.
Note that the function:
[Equation]
is Gaussian-shaped and therefore unimodal and has a unique maximum.
η x for η ∈ (0 , 1] is monotonically decreasing, and therefore their product is also unimodal.
Since we are only interested in solutions for n ∈ N , we can reduce our problem to ﬁnding the maximal n such that:
[Equation]
and from lemma 2 we have:
[Equation]
So we get that S ( n ) is monotonically increasing as long as (cid:4) n − 1 M (cid:5) M ≤ ηK − M 1+ η , and the largest integer for which this condition holds is:
[Equation]
Lemma 4.
Denote by N (cid:0) B dR (cid:1) the number of integer lattice points in B dR (the d -dimensional zerocentered ball of radius R ).
Then:
[Equation]
Proof.
Let I (cid:0) B dR (cid:1) be the set of all integer lattice points in B dR .
For x ∈ I (cid:0) B dR (cid:1) , deﬁne:
[Equation]
Let y ∈ (cid:83) x ∈I ( B dR )
C x , so there exists x ∈ I (cid:0) B dR (cid:1) such that y ∈ C x , and from the triangle inequality we get:
[Equation]
d and therefore y ∈ B
d R + √ d 2 .
Since y was chosen arbitrarily, we get that:
[Equation]
Note that V ol ( C x ) = 1 for all x ∈ I (cid:0) B dR (cid:1) and that for x , x (cid:48) ∈
I (cid:0) B dR (cid:1) such that x (cid:54) =
x (cid:48) , C x ∩ C x (cid:48) is a set of measure zero, hence:
[Equation]
Assume for convenience that d mod 2 ≡ 0 , so Γ (cid:0) d 2 + 1 (cid:1) =
(cid:0) d 2 (cid:1) ! , and Stirling’s approximation yields:
[Equation]
On the other hand, note that B d R − √ d 2 ⊆ (cid:83) x ∈I ( B dR ) C x , and therefore:
[Equation]
Lemma 5.
Let K
[Equation]
M be two ﬁxed natural numbers, s ∈ (0 , 1) a constant sensitivity parameter
[Equation]
and let”
[Equation]
Then:
[Equation]
Proof.
By Stirling’s approximation
[Equation]
it holds that:
[Equation]
and a slightly more accurate version is:
[Equation]
(cid:18) x + 1(cid:19) · (cid:16) x e (cid:17) x and by plugging this approximation to the deﬁnition of a multinomial coefﬁcient we get after some rearranging:
[Equation]
Note that using this approximation and some rearrangements
[Equation]
we get that the following are equivalent:
[Equation]
so we can characterize the set:
[Equation]
as ˜ T K,M (cid:39) T K,M .
We’ll start by ﬁnding a condition that will assure us that ( a 1 , . .
. , a M ) ∈ T K,M (i.e., we will characterize a subset of T K,M ).
First, note that by the AM-GM inequality,
[Equation]
and therefore:
M (cid:89)
[Equation]
Since we’re interested in a subset of T K,M , we can show that the last inequality holds when we replace the ﬁrst term in the left-hand side ( (cid:81)
Mi =1 (cid:113) a i + 16 ) with (cid:0) KM + 16 (cid:1)
M 2 (if the new inequality holds, (19) must hold as well), and we’re left with:
[Equation]
Now, for i ∈
[ M ] , let h
i
=
MK a i − 1 , so:
[Equation]
(cid:88)
i
[Equation] [Equation]
M
[Equation]
i (cid:19)
[Equation] [Equation]
and the last inequality becomes:
[Equation]
Observe that for all x ≥ − 1 , it holds that:
[Equation]
and therefore:
[Equation]
so it sufﬁces (again, we’re only interested in a subset of T K,M ) to show that:
[Equation]
Now, observe that:
[Equation]
so (21) becomes:
M (cid:88)
[Equation]
and we get that:
[Equation]
Let us now turn to ﬁnding a condition that will assure us that ( a 1 , . .
. , a M ) / ∈ T K,M (i.e., we will characterize a subset of the complement of T K,M ).
Note that: (cid:32) M (cid:89) i = (cid:114) a i + 1(cid:33) · exp (cid:32) M (cid:88)
i =
a i ln (cid:18) M
K a i (cid:19)(cid:33) =
(cid:18) K
M +  (cid:19) M· exp (cid:32) M (cid:88) i =
(cid:20) a
i ln (cid:18) M
K a i (cid:19)
+2 ln (cid:18)
Ma
i + M 6 K + M (cid:19)(cid:21)(cid:33) so if ( a 1 , . .
. , a M ) / ∈ T K,M , we must have that: (cid:18) K
M +  (cid:19) M· exp (cid:32) M (cid:88) i =
(cid:20) a
i ln (cid:18) M
K a i (cid:19)
+2 ln (cid:18)
Ma
i + M 6 K + M (cid:19)(cid:21)(cid:33) >
s − 1 · (cid:18) K
M +  (cid:19)
M⇐⇒
M (cid:88) i =
(cid:20) a
i ln (cid:18) M
K a i (cid:19)
+2 ln (cid:18)
Ma i + M 6 K + M (cid:19)(cid:21) >
ln (cid:0) s − 1 (cid:1) and using the same deﬁnition of h i as before, we get:
[Equation]
where the last equality is due to the fact that (cid:80) Mi
=1 h i = 0 .
Observing the ﬁrst order Taylor polynomial of the function:
[Equation]
at x = 0 with the remainder in the Lagrange form
[Equation]
we get:
[Equation]
for some ξ i between 0 and h i .
Note that f (cid:48)(cid:48) ( ξ i ) is monotonically decreasing with ξ i for ξ i ∈ ( − 1 , M − 1] , and using this fact and the fact that K ≥ 1 , (22) is lower bounded by:
[Equation]
So we can limit ourselves to looking at the cases where:
[Equation]
and we get that:
[Equation]
[Table]
[Equation]
Lemma 6.
Let K, M be two ﬁxed natural numbers, and s ∈ (0 , 1) a constant sensitivity parameter.
Then the number of multinomial coefﬁcients, (cid:0)
[Equation]
, which uphold:
[Equation]
is upper bounded by:
[Equation]
and lower bounded by:
[Equation]
Proof.
Let (cid:0)
K a 1 ,...,a M (cid:1) be a multinomial coefﬁcient for which it holds that:
[Equation]
and denote:
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
(cid:1)(cid:41)
[Equation]
(cid:1)(cid:41) then by lemma 5,
[Equation]
and therefore:
[Equation]
so in order to bound the cardinality of T K,M (which is the quantity we are interested in), we can ﬁnd an upper bound on the cardinality of T UK,M and a lower bound on the cardinality of T LK,M .
Let B ∈ { L, U } and a ∈ T BK,M , and denote:
[Equation]
For i ∈ [ M ] , denote x i :=
a i − KM .
So the problem has changed to ﬁnding the number of integer
M -tuples x 1 , . . .
, x M such that (cid:80) Mi
=1 x i = 0 and (cid:107) x (cid:107) ≤ R ( B ) , which is the number of integer lattice points x in the zero-centered M -dimensional ball of radius R ( B ) that uphold (cid:80) Mi
=1 x i = 0 .
Note that the intersection of a d -dimensional ball of radius R with the hyperplane H = (cid:110) y ∈ R d
cid:12)(cid:12)(cid:12)(cid:68)
y ,(cid:126) 1 (cid:69) =
(cid:111) is a ( d − 1) -dimensional ball, so we can assume that the number of integer lattice points in the zero-centered d -dimensional ball of radius R whose coordinates add-up to zero is ∼ 1 √ d · N (cid:0) B
d − 1 R (cid:1) , where 1 √ d is the cosine of the angle between H and one of the axis-aligned ( d − 1) -dimensional hyperplanes in R d .
In our case, R = R ( B ) and d = M , so by lemma 4, the cardinality of T UK,M is upper bounded by:
[Equation]
and the cardinality of T LK,M is lower bounded by:
[Equation]
Lemma 7.
Let K be a ﬁxed natural number, η ∈ (0 , 1] , and s ∈ (0 , 1) a constant sensitivity parameter.
Then number of integer n s such that (cid:0) Kn (cid:1)
η n ≥ s · (cid:0) K ηK 1+ η (cid:1) η ηK 1+ η is upper bounded by:
[Equation]
) − 1)
Proof.
Denote n := ηK 1+ η + x , so:
[Equation]
Recall that by Stirling’s approximation we know that:
[Equation]
So the number x s which uphold:
[Equation]
K
[Equation]
K
[Equation]
⇐⇒ ln
(cid:32) s − 1 · (cid:18) ηK 1 + η (cid:19)
ηK 1+ η + 12 (cid:18) K
+ η (cid:19) K
+ η + 12 (cid:33) ≥ ln (cid:32)(cid:18) ηK 1 + η + x (cid:19) ηK 1+ η + x +
[Equation]
After some rearranging, one can observe that (23) is equivalent to:
[Equation]
and since for all t > − 1 , t 1+ t ≤ ln (1 + t ) , the number of x s which uphold (24) is upper bounded by the number of integer x s for which it holds that:
[Equation]
Recall that for inequalities of the form ax 2 + bx + c ≤ 0 where a > 0 , the set of all values of x which satisfy this inequality is (cid:16) − b −√ b 2 − 4 ac 2 a , − b + √ b 2 − 4 ac 2 a (cid:17) and the number of integer values of x which satisfy this condition is approximately √ b 2 − 4 ac a (the interval’s length).
In our case, (25) gives us:
[Equation]
and therefore:
[Equation]
Lemma 8.
Let K
[Equation]
M be two ﬁxed natural numbers, η ∈ (0 , 1]
[Equation]
and s ∈ (cid:0) 0 , e − 1 . 5 (cid:3) a constant sensitivity parameter.
Denote:
[Equation]
deﬁne F : D K,M −→ R as:
[Equation]
and let x (cid:63) :
= arg max x ∈ D K,M F ( x ) .
If ηK 1+ η ≥ ( M − 1) , then the number of x ∈ D K,M which uphold
F ( x ) ≥ s · F ( x (cid:63) ) is bounded from above by:
[Equation]
Proof.
By lemma 3, x (cid:63) (cid:39) (cid:16) ηKM (1+ η ) , ηKM (1+ η ) , . . .
, ηKM (1+ η ) ,
K M (1+ η ) , . .
K M (1+ η ) (cid:17) .
By lemma 7, the number of n s between 0 ≤ n ≤ K such that (cid:0) Kn (cid:1)
η n ≥ s · (cid:0) K ηK 1+ η (cid:1) η ηK 1+ η is upper bounded by
K √ (2ln(
[Equation] [Equation]
[Equation] [Equation]
ln
[Equation]
) , by lemma 6 the number of non-negative integer M -tuples a 1 , . .
. , a M such that a 1 + . . .
+ a M = ηK 1+ η and (cid:0)
[Equation]
(cid:1) is bounded from above by
[Equation]
, and the number of non-negative integer M -tuples b 1 , . .
, b M such that b 1 + . . .
+ b M = K 1+ η and (cid:0)
[Equation] [Equation]
(cid:1) is bounded from above by
[Equation] [Equation]
M
[Equation]
In total, without taking into consideration the interactions between the three multiplicands (so our bound is not tight)
[Equation]
the number of x ∈ D K
[Equation]
M which uphold
F ( x ) ≥ s · F ( x (cid:63) ) is upper bounded by:
[Equation]
multinomial and since ηK ( M − 1)(1+ η ) ≥ 1 and s ≤ e − 1 .
5 (and hence 2ln ( s − 1 )
− 1 2(ln( s − 1 ) − 1) ≤ 2 ), this can be further bounded by:
[Equation]
Lemma 9.
Let K, M be two ﬁxed natural numbers, η ∈ (0 , 1] , and s ∈ (cid:0) 0 , e − 1 .
5 (cid:3) a constant sensitivity parameter.
Denote:
[Equation]
deﬁne F : D K,M −→ R as:
[Equation]
and let x (cid:63) :
= arg max x ∈ D K,M F ( x ) .
If K 1+ η ≥ M 2 , then the number of x ∈ D K,M which uphold F ( x ) ≥ s · F ( x (cid:63) ) is bounded from below by:
[Equation]
Proof.
By lemma 6, the number of non-negative integer M -tuples b 1 , . . .
, b M such that b 1 + . . .
+ b M =
K 1+ η and (cid:0)
K 1+
[Equation]
b 1 ,...,b M
[Equation] [Equation]
K
[Equation]
M
[Equation]
1+
[Equation] [Equation]
K M
[Equation]
(cid:1) is bounded from below by
[Equation]
Since these b s are only a subset of the elements of D K,M which F takes into consideration, this is also a (quite loose) lower bound on the number of x ∈ D K,M which uphold F ( x ) ≥ s · F ( x (cid:63) ) .
Now, we know that:
[Equation]
and since s ≤ e − 1 . 5 , it holds that ln (cid:0) s − 1 (cid:1) > 1 , we get:
[Equation]
and therefore:
[Equation]
C L OWER BOUNDS ON THE ε - SEPARATION RANK C. P RELIMINARIES C.1. T ENSORS AND THEIR MATRICIZATION
We begin by laying out basic concepts in tensor theory required for the upcoming analysis.
The core concept of a tensor may be thought of as a multi-dimensional array.
The order of a tensor is deﬁned to be the number of indexing entries in the array, referred to as modes .
The dimension of a tensor in a particular mode is deﬁned as the number of values taken by the index in that mode.
If A is a tensor of order N and dimension M i in each mode i ∈ [ N ] , its entries are denoted A d 1 ...d N , where the index in each mode takes values d i ∈
[ M i ] .
We will make use of the concept of the matricization of A w.r.t.
the balanced partition ( P, Q ) , denoted (cid:74) A (cid:75) P,Q ∈ R M N/ 2 × M N/ 2 , which is essentially the arrangement of the tensor elements as a matrix whose rows correspond to P and columns to Q .
Suppose A ∈ R M ×···× M is a tensor of order N , and let ( P, Q ) be a balanced partition of [ N ] , i.e. P and Q are disjoint size N / 2 subsets of [ N ] whose union gives [ N ] .
The matricization of A w.r.t.
the partition ( P, Q ) , denoted (cid:74) A (cid:75) P,Q , is the M N / 2 -by- M N / 2 matrix holding the entries of A such that A d 1 ...
d N is placed in row index 1 + (cid:80) N / 2 t
=1 ( d p t − 1) M N / 2 − t and column index 1 + (cid:80) N / 2 t
=1 ( d q t − 1) M N / 2 − t .
We now present the concept of grid tensors, which are a form of function discretization (Hackbusch, 2012).
Essentially, the function is evaluated for a set of points on an exponentially large grid in the input space and the outcomes are stored in a tensor.
Formally, ﬁxing a set of template vectors x (1) , . .
, x ( Z ) ∈
[ V ] , the points on the grid are the set { ( x ( d 1 ) , . . .
, x ( d N ) ) } Zd 1 ,...,d N =1 .
Given a function y ( x 1 , . . .
, x N ) , the set of its values on the grid arranged in the form of a tensor are called the grid tensor induced by y , denoted A ( y ) d 1 ,...,d N ≡ y ( x 1 = x ( d 1 ) , . . .
, x N = x ( d N ) ) .
C.1. ε - RANK
We will make use of the concept of ε -rank
Alon et al.
(2013) of a matrix A deﬁned for any ε > 0 as the minimum rank over matrices that approximate every entry of A to within an additive ε .
We will prove lower bounds on the ε s for which the ε -rank a matrix remain high by the following lemma:
Lemma 10.
Let M ∈ R n × n be symmetric matrix and ε > 0 , then:
[Equation]
Proof.
Let E ∈ (cid:2) − ε 2 n , ε 2 n (cid:3) n × n , we need to prove that rank ( M + E ) ≥ k .
Since M is symmetric, M is diagonalizable with eigenvalues λ 1 ≥ λ 2 ≥ · · · ≥ λ n .
Denote by v 1 , v 2 , . .
. , v n the eigenvectors that are normalized according to the l 1 norm, then for any i ≤ k we have that:
[Equation]
In particular for any i ≤ k we have that ( M + E ) v i (cid:54) =
and since v 1 , v 2 , . .
. , v k are linearly independent we conclude that rank ( M + E ) ≥ k .
Finally, we will use the following lemma for lower bounding the amount of small eigenvalues of symmetric matrices:
Lemma 11.
Let M ∈ R n × n be a symmetric matrix with diagonal entries that equal to 1 , and denote its eigenvalues by λ 1 ≥ λ 2 , . . .
, λ n .
Then:
[Equation] [Equation]
Proof.
Since the trace of a matrix equals to both the sum of its eigenvalues and its diagonal entries, we get:
[Equation]
Eq.
28 follows from the fact that:
[Equation]
C.1. H IGH -D IMENSIONAL S PHERES
We will use the Lebesgue measure on the sphere for taking expectations on S d .
For any measurable subset A ⊆ S d this measure is deﬁned as the d + 1 dimensional volume of the "wedge" in the ball
B d +1 :
[Equation]
Where λ d +1 denotes the Lebesgue measure on R d +1 .
We will also use an unnormalized version of this measure:
[Equation]
Smith & Vamanamurthy (1989) showed that µ unnormalized (cid:0) S d (cid:1) is monotonically decreasing for d > 5 .
The following lemma bounds the rate of this decrease:
Lemma 12.
For any d > 5 the following holds:
[Equation]
Proof.
Since d > 5 we know that µ unnormalized (cid:0) S d (cid:1) >
µ unnormalized (cid:0) S d +1 (cid:1) and therefore:
[Equation]
Finally, Gipple (2014) showed that µ unnormalized ( S d − 1 ) µ unnormalized ( S d +1 ) =
d +12 π , completing the proof.
Finally, we will use a well known fact regarding the variation of the sphere volume for different radii (see for example Smith & Vamanamurthy (1989)):
Fact 1.
For any d ∈ N , R > 0 the following holds:
[Equation] [Equation]
C. P ROOF OF THE LOWER BOUND
In this subsection, we prove theorem 2 of the main text.
We will follow the proofs of Levine et al.
; Wies et al.
(2021), with important adjustments to the ε -sequential-separation rank deﬁnition.
We begin by showing that high ε -rank
Alon et al.
(2013) of the grid tensor matricization implies high ε -sequential-separation rank (see section 2.2 of the main text) of the function.
Essentially, we apply claim 1 from Levine et al.
to ε -approximations obtained from the ε -separation-rank deﬁnition.
This relation, which holds for all functions, is formulated below for functions realized by the analyzed Transformer network:
Lemma 13.
For y ( p,i
[Equation]
,L,
[Equation]
context as deﬁned in theorem 1 of the main text.
Let ε - seq
[Equation]
sep
[Equation]
denote its ε -sequential-separation rank.
Then, for any integer Z , any ε > 0 , any set of template vectors x (1) , . . .
, x ( Z ) ∈ R
d x and any sub-matrix M of
[Equation]
holds that:
[Equation]
(36) where A ( Z y ( p,i ) ,L,dx in-context ) is the grid tensor of Z y ( p,i ) ,L,dx in-context with respect to the above template vectors.
Proof.
If ε - seq - sep (cid:16) y ( p,i ) ,L,d x in-context (cid:17) =
∞ then the inequality is trivially satisﬁed.
Otherwise, assume that ε - seq - sep (cid:16) y ( p,i ) ,L,d x in-context (cid:17) =
K ∈ N , and let ˜ y be an ε -approximation for Z y ( p,i ) ,L,dx in-context with seq - sep
(˜ y )
=
K .
By claim 1 of Levine et al.
we have that rank ( (cid:74)
A (˜ y ) (cid:75) a , b ) ≤ K .
Denote by ˜ M the sub-matrix of (cid:74)
A (˜ y ) (cid:75) a , b that corresponds to the rows and columns in M .
Now, since ˜ y is an ε -approximation for Z y ( p,i ) ,L,dx we have that (cid:13)(cid:13)(cid:13) M − ˜ M
(cid:13)(cid:13)(cid:13) ∞ ≤ ε .
Finally, by deﬁnition
[Equation]
Relying on lemma 13, we will bound the ε -sequential-separation rank from below via the ε -rank of sub-matrices of (cid:74) A ( Z y ( p,i ) ,L,dx in-context ) (cid:75) a , b .
Denote d := ( d x − H ) / 2 , λ :=
3 L − 2 , lemmas 10, 11 assure us that for n := (cid:16)(cid:16) dλ (cid:17)(cid:17)
=
Ω (cid:0) 2 L · d x (cid:1) it is enough to prove that there exists an assignment to the network’s weights, as well as choice of templates vectors, for which there exists sub-matrix
M ∈ R n × n of
[Equation]
) (cid:75) a , b that is symmetric with diagonal entries that equals to 1 and with
[Equation]
M (cid:107) F ≤ (cid:16)(cid:112) ( d + 1) (cid:17) n 34 , in order to show that ε - seq - sep (cid:16) y ( p,i ) ,L,d x in-context (cid:17) ≥ n 1 / 4 2 √ d +1 for ε ≤ 12 n 2 .
Now we will use a corollary that is direct results of the proof in Levine et al.
This corollary shows that if Z y ( p,i ) ,L,dx in-context is able to produce vectors that do not change the analysis in Levine et al.
(2020), then for any matrix B ∈ R n × d with rows that are l 2 normalized, there exists an assignment to the networks weights, as well as choice of templates vectors, for which there exists a sub-matrix of the grid tensor matricization that is equal to 3 M = (cid:0) BB T (cid:1) (cid:12) λ .
Importantly M is symmetric.
In addition, since the rows of B are l 2 normalized, the diagonal entries of M equals to 1 .
Therefore, M upholds the assumptions of lemmas 10, 11, 13 and it is enough to ﬁnd B for which (cid:107) M (cid:107) F ≤ (cid:16)(cid:112) ( d + 1) (cid:17) n 34 .
Corollary 2.
Let d, λ > 0 , assume
that for any matrix A ∈ R (( d 3 L − 2 )) × d with rows that are lnormalized, there exists a choice of template vectors x (1) , . . .
, x ( Z ) , an assignment to the embedding layer and the ﬁrst self-attention layer key and query weights, such that for any j 1 , j 2 ∈ (cid:104)(cid:16)(cid:16) d 3 L − 2 (cid:17)(cid:17)(cid:105) the output of the ﬁrst self-attention layer on (cid:16) j 1 , j 2 + (cid:16)(cid:16)
d 3 L − 2 (cid:17)(cid:17)(cid:17) is:
[Equation]
(cid:33) u for
[Equation]
Otherwise where φ ( j ) ≡ (cid:98) j − 1 / d a (cid:99) ·
( d a − 1) + ( j − 1 mod d a ) + 1 .
Then for any matrix B ∈ R n × d with rows that are l 2 normalized, there exists an assignment to the networks weights, as well as choice of templates vectors for which there exists sub-matrix of the grid tensor matricization that equal to M = (cid:0) BB T (cid:1) (cid:12) λ .
Now we will shows that indeed Z y ( p,i ) ,L,dx in-context is able to produce vectors that do not change the analysis in Levine et al.
(2020) and the assumptions of corollary 2 holds.
Lemma 14.
Let A ∈ R
(( d 3 L − 2 ))
× d with rows that are l 2 normalized, then there exists a choice of template vectors x (1) , . . .
, x ( Z ) , an assignment to the embedding layer and the ﬁrst self-attention layer key and query weights, such that for any j 1 , j 2 ∈ (cid:104)(cid:16)(cid:16) d 3 L − 2 (cid:17)(cid:17)(cid:105) the output of the ﬁrst self-attention layer on (cid:16) j 1 , j 2 + (cid:16)(cid:16)
d 3 L − 2 (cid:17)(cid:17)(cid:17) is:
[Equation]
1 ,h (cid:33) u for ∀ α ∈ [ d x ]
[Equation]
d 2 N ( α − 1) mod d a = d a −
Otherwise where φ ( j ) ≡ (cid:98) j − 1 / d a (cid:99) ·
( d a − 1) + ( j − 1 mod d a ) + 1 .
Proof.
We will ignore Z y ( p,i ) ,L,dx in-context ’s element-wise multiplication with vocabulary embedding matrix by choosing ∀ i, j M V i,j = 1 (by the terms of corollary 2 it sufﬁces to ﬁnd any assignment of the learned weights).
For any i ∈ (cid:104) 2 (cid:16)(cid:16) d 3 L − 2 (cid:17)(cid:17)
+ 1 (cid:105) our templates vectors will be:
[Equation]
N
[Equation]
Otherwise
We will implement summation of the inputs embedding in the ﬁrst self-attention layer, we will follow
Levine et al.
and set the ﬁrst layer self-attention key and query weights to:
[Equation]
W
[Equation]
This assignment implements summation of the inputs embedding in the ﬁrst self-attention layer since:
[Equation]
where (1) is because W Q, 1 ,h = W K, 1 ,h are matrices that are zero everywhere except for entry (1 , d a ) and that all the entries in the vocabulary embedding matrix equals to 1 , and (2) because of linearity.
Therefore, for any j 1 , j 2 ∈ (cid:104)(cid:16)(cid:16) d 3 L − 2 (cid:17)(cid:17)(cid:105) the output of the ﬁrst self-attention layer on j 1 , j 2 is:
[Equation]
[Equation]
Finally, we need to show that indeed for any j 1 , j 2 ∈ (cid:104)(cid:16)(cid:16)
[Equation]
[Equation]
eq 44 give the desired u :
[Equation]
Otherwise
The third and forth cases are clear from x (cid:48) s deﬁnition, so it remain to prove the ﬁrst and second cases.
For this we will examine d 1 , d 2 .
d 1 = j 1 ≤ (cid:16)(cid:16)
d 3 L − 2 (cid:17)(cid:17) and therefore:
[Equation]
2 (cid:17)(cid:17)
[Equation]
[Equation]
d 3 L − 2 (cid:17)(cid:17)(cid:105) and therefore:
[Equation]
So it clear that also the ﬁrst and second cases upholds.
Returning to ﬁnding B for which (cid:107) M (cid:107) F ≤ (cid:16)(cid:112) ( d + 1) (cid:17) n 34 , we will use the probabilistic method for proving the existence of such B , i.e. we will show that for random B the expectation of (cid:107) M (cid:107) F ≤ (cid:16)(cid:112) ( d + 1) (cid:17) n 3 4 and therefore in particular there exists such B .
Lemma 15.
For any d, λ ∈ N such that λ ≥ d the following holds:
[Equation]
Proof.
We start by bounding the expectation of the squared norm:
[Equation]
Now, E u,v ∼ S d
[Equation]
u, v (cid:105) 2 λ
[Equation]
≤ ( d + 1) (cid:16)(cid:16) dλ (cid:17)(cid:17) − 12 by lemma 16 and therefore we got that:
[Equation]
Finally
[Equation]
by Jensen inequality we have that:
[Equation]
C. T ECHNICAL LEMMAS
Lemma 16.
For any d, λ ∈ N such that λ ≥ d the following holds:
[Equation]
Proof.
We will use conditional expectation to make reduction for simpler expectation.
From rotational invariance of the uniform measure on S d we know that for every rotation matrix R ∈ SO ( d ) and unit vector v ∈ S
d we have that:
[Equation] [Equation]
where the ﬁrst equality holds because R is orthogonal.
Therefore, by choosing R such that Rv = ewe will get that:
[Equation]
Now we can calculate the last expectation directly:
[Equation]
Now
[Equation]
by lemma 12 and fact 1 we have that:
[Equation]
Therefore we have that:
[Equation]
Finally
[Equation]
by lemma 17 each term in the integral is upper bounded by (cid:16)(cid:16) dλ (cid:17)(cid:17) − 12 and thus:
[Equation]
Lemma 17.
For any x ∈ [0 , 1] and d, λ ∈ N such that λ ≥ d the following holds:
[Equation]
Proof.
Note that since x 2 λ (cid:0) 1 − x 2 (cid:1) d 2 = 0 in the boundaries ( x ∈ { 0 , 1 } ) , it is enough to prove the inequality for critical points
[Equation]
Therefore
[Equation]
x 2 = 2 λ 2 λ + d is the only critical point and:
[Equation]
where the last inequality follow from the fact that:
[Equation]
D E XPERIMENTAL D ETAILS D. K NN-TAPT
We conducted the network training described in section 3.1 of the main text with AdamW optimizer (with the parameters suggested in the original RoBERTa paper: β 1 = 0 .
9 , β 2 = 0 .
98 , ε = 10 − 6 and weight decay of 0.01), with batch sizes of 128 or 256 (depending on model size) and sequences oftokens each.
We started with pretrained RoBERTA-base weights from the HuggingFace
Transformers repository 4 , and continued training them on the MLM task with masking probability of 15%, where each masked token had a probability of 80% of being replaced with the special [ MASK ] token, 10% of being replaced with a random token and 10% of being kept the same.
The data used for this phase of training was created using the four different procedures described in section 3.1.
After the training was ﬁnished, we evaluated the models’ performance using the SentEval kit.
D. K NN-P RETRAINING
We conducted the network training described in section 3.2 of the main text with AdamW optimizer (with the parameters suggested in the original GPT-2 paper:
β 1 = 0 .
9 , β 2 = 0 .
95 , ε = 10 − 8 and weight decay of 0.1), with batch size of 512 and sequences of 256 tokens each.
We pretrained a
HuggingFace Transformers implementation of GPT-2 from scratch on Wikipedia with the standard
LM objective, and switched to a mixture of the standard data and our generated kNN data in two different points during training.
After the training was ﬁnished, we evaluated the models’ performance on the Natural Questions benchmark.
E K NN-P RETRAINING AT DIFFERENT CHECKPOINTS
The following table includes F1 evaluation scores of zero shoe closed book Natural Questions examples for different model sizes at different training checkpoints.
Overall, further pretraining seems to improve the effectiveness of kNN-Pretraining.
[Table]
F K NN-P RETRAINING ON ADDITIONAL BENCHMARKS
The main text describes experiments on the Natural Questions dataset.
We test how kNN-Pretraining affects other NLU tasks, by examining several tasks from the GLUE benchmark (Wang et al., 2018) – Multi-Genre Natural Language Inference (MNLI)
(Williams et al., 2017), Recognizing
Textual Entailment (RTE) (Dagan et al.
(2010) and others), and the The Winograd Schema Challenge (WNLI)
(Levesque et al., 2012).
As in the case of Natural Questions, we evaluate the zero-shot performance of our models since it is a direct probe to the abilities of the model straight after the process of pretraining.
In contrast to Natural Questions, the GLUE tasks we examined are classiﬁcation tasks and not generation tasks, so assessing zero shot performance on them is not straightforward.
We therefore follow the template-based method of Gao et al.
(2021a)
for converting the tasks’ data into a format processable by unidirectional language models.
Notably, the examined GLUE classiﬁcation tasks are not easy for the examined unidirectional models in zero shot.
Table 3 includes the zero-shot scores of the 345 M parameter model that trained regularly for 200 K steps and then continued training for 20 K steps of kNN-Pretraining, versus the average of 3 baselines that trained regularly for the same number of overall steps (the same models used in ﬁgure 1).
Similarly to the results on Natural Questions (ﬁgure 1), all examined models score only slightly better than random guess on the examined GLUE tasks.
However (and again similarly to the case of Natural Questions), we get a clear signal that kNN-Pretraining signiﬁcantly moves the needle when applied for just 10% of the regular pretraining time.
We conjecture that when using stronger models (that train for longer and over more data), the positive effect of kNN-Pretraining will be enhanced, since as the model improves, it can better understand and utilize the various in-context hints that kNN-Pretraining provides.
[Table]