The Inductive Bias of In-Context Learning: Rethinking Pretraining Example Design
Authors: Yoav Levine
Noam Wies
Daniel Jannai
Dan Navon
Yedid Hoshen
Amnon Shashua
Venue: ICLR 2022 Spotlight (2022)
Keywords: Language Modeling
Pretraining
Self-attention
Transformers
Expressivity
Separation Rank
Sentence Embeddings

Abstract: 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 fine-tuning stages, which do not necessarily appear related at first glance. Second, our result clearly indicates further improvements to be made in NLM pretraining for the benefit of Natural Language Understanding tasks. As an example, we propose ``kNN-Pretraining": 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. 

Published as a conference paper at ICLR 2022
THE INDUCTIVE BIAS OF IN-CONTEXT LEARNING: RETHINKING PRETRAINING EXAMPLE DESIGN
Yoav Levine, Noam Wies, Daniel Jannai, Dan Navon, Yedid Hoshen & Amnon Shashua The Hebrew University of Jerusalem (yoav.levine) noam, wies, Samiel.jannee Ban.nay @mail www..huiiii..... i1
ABSTRACT
Pretraining Neural Language Models (NLMs) over a large corpus involves chunk- ing the text into training examples, which are contiguous text segments of sizes processable by the neural architecture. We highlight 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 twofold role. First, formalizes the motivation behind broad line of recent successful NLM training heuristies, proposed for the pretraining and fine-tuning stages, which do not necessarily appear related at first glance. Second, our result clearly indicates further improvements to be made in NLM pretraining for the benefit of Natural Language Understanding tasks, As an example, we propose . we show that including semantically related non-neighboring sentences in the same pretraining example yields improved sentence representa- tions and open domain question answering abilities. This theoretically motivated degree of freedom for pretraining example design indicates new training schemes for self-improving representations.
INTRODUCTION
Beyond excelling in their core task of pure language modeling, modern Neural Language Models (NLMs) show impressive zero- and few-shot abilities in general Natural Language Understand- ing (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.
Specifically, 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 implies that NLU capabilities that involve integrating information from different examples across the corpus (see, e.g., figure I under-favored by design in the current setting. Thus, if one sentence in the corpus can elucidate the meaning of another sentence (e.g., defines a hard concept or provides auxiliary information), our result implies that model that saw them in different training examples will enjoy this elucidation less than model that saw them in the same training example.
While standard approximation results examine the expressivity of an architecture over single input, our theoretical approach pertains to the entire training process, and examines the expressive capacity of the resultant NL. 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

---

Published as a conference paper at ICLR 2022
the network involves very small learning-rate term 7) our analysis formalizes and quantifies an "expressivity toll" that the model pays when making use of such harder-to- access stored information.
We employ the tool of function's separation rank with respect to subsets of its variables, which quantifies its ability to model input dependencies between these subsets. The separation rank was employed for analyzing the dependencies modeled by convolutional (Cohen & Shashua (017)) recurrent (Levine et al. 2018), and self-attention (Levine et al.) 2020) networks with respect to single input example. In order to analyze NLM's ability to model dependencies between different training examples, we refine the usage of this measure in two manners: (1) we introduce the e-separation rank, which measures the effective ability of function model dependencies in finite 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.
Specifically, we upper bound the log of the separation rank of a depth L width dy self-attention based NLM, with respect to two sentences that are shown in its input, by Ld) 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 Ó(d
Given common learning-rate values of E [10-6, : 10-4), this implies a guaranteed "depth deficit" 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 deficit is more significant, and may behave like a fraction of L. We leave attempts to tighten the depth deficit estimates to future work.
0.5 (1ga(11
1.1 THE 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 significantly improve both task-specific training and pretraining of NLMs. Gao et al. (2020) advance the frontier in k-shot learning via finetuning. They show that by concatenating several related training examples per input, instead of using standard fine-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. [2020); Thakur et al. (2020): when training for sentence similarity tasks, including both sentences in the same input leads to performance gain of around 10 points relative to separately encoding each sentence. In the challenging setting of open-domain question answering, [zacard & 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 a body of work aimed at reducing the quadratic dependence the Trans- former 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 s ability to integrate text across different documents in the corpus. The following approaches take further step and enable direct cross corssc connections during pretraining. Lewis et al.] (2020) attend to related documents when maximizing the likelihood of 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 Knowledge-Retriev 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. (221)) train Transformer based protein-LMP LM that receives multiple related protein sequences in-context. Their protein-I. L.M surpasses previous methods which process one sequence per input by a wide margin, with significant parameter efficiency.
model that has access
1.2 LEVERAGING 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 unified, and importantly, new methods for tilting the pretraining bias towards NLU tasks are indicated. Following the presentation of our theoretical results in section [ we detail in section 3 two controlled setting exemplifications of new methods that directly leverage the in-context bias.
Our first 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

---

Published as a conference paper at ICLR 2022
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 significant gains of the KNN-TAPT over regular TAPT on SentEval sentence similarity tasks. A dedicated ablation study shows the significance of adding the general corpus neighbors in-context, versus in separate training examples, during KNN-TAPT.
Our second experiment introduces task-independent pretraining phase, dubbed KNNN-Prttainiigg As in KNN-TAPT, we group together sen- tences with similar sentence repre- sentations in the same training ex- ample, but in KNN-Pretraining we
460 po ge about pound, cunces...
(..pens can weigh roughly 20-50 grams...
"Which is heavier, pen book?'
agproximataly
Zero-shot Closed Book
Pretra iinn
Natural Questions:
Baseline mean
0.01
Baseline std
Fl -sccre
use only sentences from the gen-
eral pretraining corpus. This can be
0.005
viewed as sentence- focused varia-
Beginning of regular pretraining
Beginning of Khil-pretraining
tion of the above surveyed pretrain-
ing schemes in Lewis et al. (2020)
and Guu et al. (2020), who oper-
0.0 000000
5000
10000
15000
20000
ate on full documents (up 512
KNN- NNN-Pretraining Steps
each), and is very similar to RETRO by [Borgeaud et al. (202T) (Deep- Mind), who show the benefits of this approach given much larger resources. Figure SaareeTshone shows that after regular pretraining for 200K steps on Wikipedia, the zero-shot closed book performance of 3 different randomly initialized GPT2- medium models (345M parameters) on open domain questions from Wikipedia (Kwiatkowski et al.. 2019) is very low (correct less than 50 questions out of 20,000 in the evaluation set). Adding ENN-Pretraining for 20K steps raises performance significantly (correct on roughly 250 questions in the evaluation set), reflecting the enhanced ability to integrate knowledge from related sentences, acquired via the in-context bias.
Figure A 10% addition of KNN--NNNNaaaaiaa Pretraierainn boosts zero-shot closed book QA score by~ 5X (evaluation set size is 20,000).
In summary, 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 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.
2 THEORETICAL ANALYSIS: THE IN-CONTEXT BIAS OF SELF-ATTENTION
In this section, we consider the entire NLM training procedure as functional that receives an unlabeled training corpus and outputs a trained NLM. Our analysis focuses on the corpus segmentation into training examples
syper-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 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 exist in broad range of architectures, but focus on self-attention since almost all modern NLMs are based on the Transformer architecture of Vaswani et al. (2017). Our theoretical framework is based on that Levine al. (2020); Wies) et al. (2021), who analyze a simplified, 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 's ability to correlate two subsets of its variable set, referred to as the separation rank. The analyzed framework captures the connectivity of attenttten but omits its softmax and ReI.U non-lineesssies (see eq. [ beloow.. We refer the reader to Levine et al. (2020): Wies et al. (202T) for a discussion on the impact of these relaxations. Essentially, they are shown weaken the overall network power but still allow meaningful comparison of the self-attention integration abilities.
3

---

Published as a conference paper at ICLR 2022
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.
2.1 THE ANALYZED IN-CONTEXT AND SEQUENTIAL REPRESENTATIONS
For an input sequence of N embedding vectors {xl E Rd. Jf- denote the function realized by the analyzed H heeaded depth-L width-d, Transformer architecture at output location E [N] by: L,,dz
(x) ...., ...,..,"))) E Idd,, where W stands for learned parameters, recursively defined:
Sw
1
Bw BBW ,. BW
1.1+ 1..da
N.ld))
Ewolh
a WiyWW gi J.d.d
(1)
ang = (Wo.. itd, w.K.k. j..dz) BW
i.0.d, Bw
Sw
where W is composed of Key, Query, Value and Output matrices: 4A W.K.I.,, WQAA WVA.A, (WO.LA) E Rdax xdz, where we assume the standard choice d, = d
[4], E [H]
For word from vocabulary of size V, w E [V], the translation into the Transformer dimension is done via a mapping EMy: [V] + Rdo;
Eus (w) = (M²)
where MV E Rd is the learned vocabulary embedding matrix, and (M²), is its wth column, also referred to as the learned word embedding for w. Overall, the function of the analyzed Transformer over sequence of N words {w> E [V]],,
1 can be written by composing eqs. I and
L.L.d, Y.M (w" ...., w gw LLLdd (ENr (w*) ,..,, EM (w
(3)
For simplicity of presentation, we examine two sentences S and S2 of equal length N: Si = {w]) and S. = [w]] j-1 The in in-context representation simply concatenates both in the input: y Vin-context ($,, 82)= LLL.d, W.M" (wi...., w/ w ..., w w.,"")
(4)
For the sequential approach, we consider a setup in which sentence S is inserted into the network training step and sentence S2 is inserted into the network at training step + 1. The output of the network at training step is therefore: y I.L.d (Si), where Wr, MM² stand for all the learned WW.My
weights before training step t. Focusing on autoregressive NLMs for simplicity of presentation (the analysis holds for bidirectional NLMs as well), the log- likelihoood loss is given by (())
- -, log (
(softmar { ((M))
and the gradient update for any learned y. (8i))) w W.
weight E (W., M is: 0.41 (S); 10-1..060 -.(5))0))))/ where 77 is the learning rate. Accordingly, the analyzed sequential representation is the network output after training step L + 1:
I.I.d y sequeential (S1. S2) (S.))( (Simm) (S2).
(5)
In practice, two relevant non-neighboring sentences are not necessarily shown in consecutive pre- training steps. In comparison to the realistic scenario of S and Sy appearing at any training step, this simplifications tilts the representation in favor of modeling high correlations between S and S2. Thus, by upper bounding the ability to correlate S₁ and S2 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 network's '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.2 A MEASURE FOR MODELING IN-CONTEXT AND SEQUENTIAL DEPENDENCIES
In this section, we refine the separation rank, used in prior work in order to analyze the dependencies between two sentences appended in-context. In section T.2.T we present the separation rank and introduce a finite precision refinement 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 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 both the in-context and sequential cases.
4

---

Published as a conference paper at ICLR 2022
2.2.1 THE EFFECTIVE SEPARATION RANK
The separation rank has been established as measure of dependencies modeled by deep convolutional, recurrent, and self-attention networks (Cohen & Shashua, [2017 Levine et al. 2018 2020). For function y(A, B) over variables A {a' E زز and B (b' E &) 10 j==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.f. (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 > RR w.r.t. (A, B) is defined as follows:
(A (3) == min
R E N: .91...99............ .R. 91. 9R XM - R.t. y (A,
Σ gr (A) (B)
(6)
If the separation rank of a function w.r.t. (A, B) is 1, is multiplicatively separable w.r.t (A, B). meaning it cannot take into account consistency between A and B. The higher (A, (v) is, the
farther 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:
(A,) (y) = min| Rt' < (A (V) ENE: Bÿ: 2 s.t.
() (0) = R [Vx E X2M (0(z) - y [
In words, if function has high separation rank, but it can be approximated up error e by function with a low separation rank, then it has a low e-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 definition 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 stored in the network weights in 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 e-separation rank. Our techniques for bounding the e-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. 2.2.2
THE SEQUENTIAL SEPARATION RANK
Levine al. 2020). who were the first to apply the separation rank to functions realized by Trans- former architectures, studied classical architecture expressivity questions which apply only to the in-context representation. Accordingly, they analyzed only the separation rank of g g....d,, defined in eq. [ and the input variables considered for calculating the separation rank were the word embed- ding vectors. A fundamental difficulty 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₂ 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 S1.
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, be E Rdd,, which explicitly associates each employed vocabulary embedding vector with the sentence index $ E (1, 2} that invoked its usage:
2 (EM (w)) = Far @
if 8 = 1
(8)
E (w!) b
if s == 2
where @ denotes dlement-wise multiplication. We define the sentence association operation over the
analyzed representations, denoted Zy (a, with y E (yinfoo (S1, S2) y sequeniil (S1, (eqs. 4
o 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 marked with the identity of the sentence that invoked them. Finally, we define the following specialization of the separation rank measure to our setting, referred to as the sequential separation L,d, rank of E {y in-comtext (S1, S2). y sequeettal (S1, S2)):
seq-sep(y) := sep(a.b) (Zy)
E-soq-sop(y) == E-seP(a. (Zy)
5

---

Published as a conference paper at ICLR 2022
Clearly, when the introduced variables are vectors of 1, the auxiliary layer in eq. 8 is the identity operation and so (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 2, 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. (020). Thus, the properties of the existing framework are unchanged under the new definition. At the same time, for the sequential case brought forth in this paper, the sequential separation rank considers both the effect of S on the gradient-updatee word embedding and the introduction of S2 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.3 THE EXPRESSIVE ADVANTAGE OF IN-CONTEXT LEARNING
We show below that the function computed by a self-attention based NLM when inserting sentences Si and Sy together in its input (the in-context representation) can model more elaborate dependencies between S₁ and S2 than the function attained when showing S1 in the input, modifying the network's s weights according to its loss, and then showing S2 in a subsequent input (the sequential representation). We begin by stating the following corollary, following from theorem 2 in Levine et al. (2020) and proposition 1 in appendix [A] which upper bounds the sequential separation rank of the in-context representation:
(,,s).L.dz be the p E [dz] entry of the analyzed in-context representation defined in Corollary 1. Let Vm conteat
eq. 4
Assume that L > loga dz. Then (Ó notation omits log terms: log d log L, log H):
} = (L de)
log seq-sep Vin-contexs
(10)
However, the e-separation rank of the sequential representation is upper bounded by a lower term:
Theorem 1. (See proof in appendix B Let Usequentiall be the p E [dg) entry of the analyzed
sequential representation defined in eq. [ Assume that all learned parameters and all gradients are
bounded: v0 E (W, MM : Amin S 101,1/255)/000 S Amax P N < dzs and that L > loga da. Then, Ve 0:
log s-so
( Vioadoafaa
)} - Ō (LL + 0.5 dz)
(11)
Therefore, gap between upper bounds on the ability to model dependencies between S and S2 is indicated. Since the learning rate is 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:
(See proof in appendix C.. For sirssneess as defined in corollary there exists an assignment of the network weights for which the following holds:
Theorem 2.
)] π (L-da)
log s-seq-sep
(Min-comext
(12)
where 8 = (3' 3
((
Notably, corollary 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, defined eq. [ is a degree 3 polynomial over its 2N dz inputs, rendering the whole network a degree 3"d polynomial. We write this polynomial as a sum over many monomials, and by definition, the separation rank of any monomial composing the polynomial is 1. Since the separation rank of 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 variables affect the
'To see this, note that for s the operation in eq. 8 includes both M³ (a) and variables from b. "The upper boundedness assumption resembles practices of gradicnt clipping and weight decay, and the lower boundedness assumption resembles finite precision.

---

Published as a conference paper at ICLR 2022
computation only via the gradient, so their impact is expected to be limited. However, considering that S2 first encounters gradient updated vocabulary matrix entries (m ) (m²)) 170C(5))/((m")) appears that both S and S2 variables enter the self-attention stack via its input, similarly to the in-context case. So the integration between S₁ and S2 occurs right from the start. and indeed show that the separation rank of both representations is similar. However, since any function of S₁ is accompanied by the learning-rate 17, the monomials for which there are many S1 variables will be multiplied by high powers of 1j. This causes many monomials to be negligibly small, and accordingly not contribute the e-separation rank. By combinatorial considerations show that the number of monomials that are not attenuated by (have sufficiently 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. Specifically, 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 (77" layers that has seen them in the same context. Common learning-rate values are on the order of E [10-6, 10-"], implying a deficit of ~ 6 layers in the sequential case. As shown in in Levine et al. (2020); Tay et al. (202T). ) 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 defined "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.
3 KNN BASED PRETRAINING EXAMPLE DESIGN
Our theoretical analysis quantifies the relation between the small magnitude of the learning rate, and the deficiency 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 specific 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 identified via kNN search in their RoBERTa-large sentence representation space (Reimers & Gurevych 2019), performed using the FAISS library (Johnson et al. 2019).
3.1 KNN TASK ADAPTIVE PRETRAINING
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 general corpus. This implics 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 100M Wikipedia sentences and appended in-context neighbors that have embeddings with over 0.8 cosine similarity to the SentEval
7

---

Published as a conference paper at ICLR 2022
STS12
STS13
STS14
STS15
STS16
STS-B
SICK-R
Avg.
Basline Roberta Model
32.1
56.3
45.2
61.3
62.0
55.4
62.0
53.5
TAPT
43.0
62.2
51.6
70.6
64.9
63.0
63.5
59.8
ANN-TAPT (random, in-batch)
40.2
62.7
51.9
64.9
62.1
61.5
65.4
58.4
KNN-TAPT (neighbors. in-butch)
40.8
62.4
53.1
66.1
63.0
61.3
65.2
58.8
KNN-TAPT (random, in-context)
44.62
62.64
51.4
65.28
64.93
64.31
66.96
60.0
KNN-TAPT (neighbors, in-context)
44.9
63.4
52.1
66.2
65.3
66.5
68.3
61.0
Table 1: KNN-TAPT, which augments the Task Adaptive Pre'Training (TAPT) setting of Gururanaan
et al. 2020), harnesses the in-context bias and improves SentEval sentence similarity scores.
example embedding, with special token inserted between different sentences. We continued until finding no more neighbors or reaching a maximum of 256 tokens in the RoBERTa vocabulary (Liu et al. 2019). This search yielded 170K examples, over which we continued training pretrained RoBERTa-base model for 5 epochs. using the first epoch for learning-rate warmup and examining peak learning rates of (1,3,5, 5, 7 10 5. See appendix D for implementation details.
Table I shows zero-shot SentEval sentence similarity scores, attained by using the average word embedding of an inserted sentence 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 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" flag relates SentEval example to its actual neighbors from the kNN search, while the "random" " flag relates it to random Wikipedia sentences from the overall neighbors pool attained in the search. The "in batch" flag 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" flag 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 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 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.2 KNN PRETRAINING
We extended the above to more general related non- non-neighboring sentences given only the general pretraining corpus. KNN-Pretraining is also motivated by the KNN-LM results of Khandelwal wall et al. (2019), who show significant benefits 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.
designing pretraining examples with
Specifically, 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 kNN- Pretraining, we pretrained GPT-base and GPT-medium (110M and 345M parameters) architectures from scratch over Wikipedia in the regular pretraining scheme, and switched to KNN-Pretraining at two different points during pretraining (200K and 400K). The training examples were of maximal size 256, and the batch size was 128 for the GPT-medinn models and 256 for the GPT--ase 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
8

---

Published as a conference paper at ICLR 2022
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. (2019); ; Asai et al.
(2019); Roberts et al. (2020). 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 figure ] demonstrates, 3 baseline models, pretrained with the regular scheme, achieve very low F1< 10 scores on this task. In contrast, ANN-Pretraining shows a low-scoring but significant improvement.
To increase the credibility of the signal, we evaluated our models on the first 20K 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 100s 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 NQ scores of models of different sizes when starting ENN-Pretraining at different checkpoints, and in appendix F zero-shot scores on several GLUE tasks, which demonstrate clear gains of KNN- Pretrainnn over the baselines.
4 DISCUSSION
Modern NLM pretraining schemes have tremendously advanced the natural language landscape, since they allowed powerful models to train on huge amounts of unlabeled text. But NLMs are now challenged with tasks which require deeper and more nuanced understanding of text, and means of improving the basic pretraining process should be considered. For a given architecture, pretraining can be improved by adding more data or finding more sophisticated training objectives to apply over existing data. In this paper we highlight a parallel path for improvement, which employs the same data and objective, but redistributes the available strength of the Transformer architecture such that important connections within the pretraining corpus are learned more effectively. Specifically, highlight the bias of the trained NLM towards modeling dependencies between chunks of text that appeared within the same training examples. In current pretraining schemes, this means that dependencies between non-neighb chunks of text under-favored. If such dependencies matter for the task at hand, we suggest rearranging the data into corresponding training examples.
We formalize the above notion. Our theoretical setup asks expressivity questions that pertain the training set rather than to a single example. We thus tie the construction of the training example with the available expressivity of the architecture: we prove that the connections that can be modeled between different training examples are bounded by the connections that can learned by shallower and weaker architecture, if these examples were inserted within the same input.
The advantage in including related text in the input of the NLM is noticed and leveraged in the empirical landscape. With that, it clear that showing the model related data is meaningful even if is in different training examples, and many leading methods elect to do just that. Our quantification of this trade-off is intended to aid informed decisions and highlight the expressivity advantage to be gained by smarter training example designs. We follow up on these recommendations and demonstrate the immediately available gains to be achieved by designing training examples that include nearest neighbors in embedding space. This method can be enhanced, and other more explicit biases can be introduced. For example, multiple mentions of the same entity, event, or concept can be concatenated within the same training example.
The gains achieved by using similarity in representation space indicate a path for self-improving representations, left for future work. After first cycle of kNN-Pretraining, the representation is refined and applying a new kNN search over it can lead to more informative next round of KNN- Pretraining. This way, deeper insight can be elicited from a given pretraining corpus.
Lastly, while this paper focused on leveraging the identified in-context bias for pretraining, it can also be tied to recent successes of in-context inference methods. From the in-context few-shot prompts Brown et al. (2020). to in context augmentations such as in Gao et al. (2020): SShhick & Schütze (2020) and many others, the benefits of biasing the prediction by appending text in-context are now widely established. The tools brought forth here can assist in clarifying the theoretical advantages of such practices. Overall, our work aims to provide timely theoretical interpretations, to help guide the rapid empirical advances of our field.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
We thank Or Sharir, Kevin Leyton-Brown, and Ori Ram for useful discussions. This research was supported by the ERC (European Research Council) and the ISF (Israel Science Foundation). Experiments were performed with Cloud TPUs and supported by Google's 's TensorFlow Research Cloud (TFRC). Yoav Levine supported by the Israel Academy of Sciences Adams fellowship.
Noga Alon, Troy Lee, Adi Shraibman, and Santosh Vempala. The approximate rank of matrix and its algorithmic applications: approximate rank. In Proceedings of the forty-fifth annual ACM symposium Theory of computing, pp. 675-684, 2013.
Akari Asai, Kazuma Hashimoto, Hannaneh Hajishirzi, Richard Socher, and Caiming Xiong. Learn ing to retrieve reasoning paths over wikipedia graph for question answering. arXiv preprint arXiv:1911.10470, 2019.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, bean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. arXiv preprint arXi: 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. arXiv preprint arXiv:2005.1411 2020.
Nadav Cohen and Amnon Shashua. Inductive bias of deep convolutional networks through pooling geometry. In 5th International Conference on Learning Representations (ICLR), 2017.
Alexis Conneau and Douwe Kiela. Senteval: An evaluation toolkit for universal sentence representa- tions. arXiv preprini 1003v::880300555 2018
Ido Dagan, Bill Dolan, Bernardo Magnini, and Dan Roth. Recognizing textual entailment: Rational, evaluation and -erraccceeeerrrtum...mm Natural Language Engineering, 16(1):105-1M) 105, 2010.
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Gold- ing, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, September 2021a. htts::ts::// /oii org. wttp://../o....////... 5wt///2/1///....o 53 71628
Tianyu Gao, Adam Fisch, and Dangi Chen. Making pre-trained language models better few-shot learners. arXiv preprint 1xXi:20022.152255 2020.
Tianyu Gao, Xingcheng Yao, and Danqi Chen. Simese: Simple contrastive learning of sentence embeddings. arXiv preprint arXiv:2104. 08821, 2021b.
Jake Gipple. The volume of n-balls. Rose-Hulman Undergraduate Mathematics Journal, 15(1):14, 2014.
Suchin Gururangan, Ana Marasović, Swabha Swayamdipta, Kyle Lo, Iz Beltagy, Doug Downey, and Noah A Smith. Don't stop pretraining: Adapt language models to domains and tasks. arXiv preprint arXiv:2004.10964 2020.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. Realm: Retrieval- augmented language model pre-training arXiv preprint 2xXi::2002.08909, 2020.
Wolfgang Hackbusch. Tensor spaces and numerical tensor calculus, volume 42. Springer Science & Business Media, 2012.
Samuel Humeau, Kurt Shuster, Marie-Anne Lachaux, and Jason Weston. Poly-encoders: Archi- tectures and pre-training strategies for fast and accurate multi-sentence scoring. In International Conference on Learning Representations, 2020. httliippss/opeeeeeeeee.eeee
id=SkxgnnNEvä
Gautier Izacard and Edouard Grave. Leveraging passage retrieval with generative models for open domain question answering. arXiv preprint orx1v:2007.01282, 2020.
10

---

Published as a conference paper at ICLR 2022
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 2019.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv: 191100172. 2019.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
Kenton Lee, Ming-Wei Wei Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. arXiv preprint arXiv: 2019.
Hector Levesque, Ernest Davis, and Leora Morgenstern. The winograd schema challenge. In Thirteenth International Conference on the Principles of Knowledge Representation and Reasoning, 2012.
Yoav Levine, Or Sharir, Alon Ziv, and Amnon Shashua. Benefits of depth for long recurrent networks. (ICLR 2018) International Conference on Learning Representations workshop, 2018.
long-term term memory of
Yoav Levine, Noan Wies, Or Sharir, Hofit Bata, and Amnon Shashua The depth-to-width interplay in self-attention. In Advances in Neural Information Pro- cessing Systems, 2020. URL. (00of//5000430000552282211111113
nips. cc/ wttppppccpppppp/////// /22pp///20///illll httpp//r..iippscc///pp//////////
Mike Lewis, Marjan Ghazvininejad, Gargi Ghosh, Armen Aghajanyan, Sida Wang, and Luke Zettlemoyer. Pre-training via paraphrasing. arXiv preprint 2rXiv:2006.50222 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv 2019.
Roshan Rao, Jason Liu, Robert Verkuil, Joshua Meier, John F Canny, Pieter Abbeel, Tom Sercu, and Alexander Rives. Msa transformer. bioRxiv, 2021.
Nils Reimers and Iryna Gurevych. Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv:1908. 10084 2019.
Adam Roberts, Colin Raffel, and Noam Shazeer. How much knowledge can you pack into the parameters of a language model? arXiv preprint arXiv:2002.089110, 2020.
Timo Schick and Hinrich Schütze. It's not just size that matters: Small language models are also few-shot learners. arXiv preprint arXiv:2009.07118 2020.
David J. Smith and Mavina Vamanamurthy, How small unit ball? Mathematics Magazine, 62 12):101-107, 1989. doi: 10. 10.1080000255500X1199999 1080/00255705 989. 11977419
URL htttps://doi.oo
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint arXiv:2011.000 1.04006, 2020.
Yi Tay, Mostafa Dehghani, Jinfeng Rao, William Fedus, Samira Abnar, Hyung Won Chung, Sharan Narang, Dani Yogatama, Ashish Vaswani, and Donald Metzler. Scale efficiently: Insights from pre-training and fine-tuning transformers. arXiv preprint arXiv:2109. /0686, 2021.
Nandan Thakur, Nils Reimers, Johannes Daxenberger, and Iryna Gurevych. Augmented shert: Data augmentation method for improving bi-encoders for pairwise sentence scoring tasks, arXiv preprint arXiv:2010.082400 2020.
11

---

Published as a conference paper at ICLR 2022
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances neural information processing systems, pp. 5998-6008, 2017.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1804. 2018.
Noam Wies, Yoav Levine, Daniel Jannai, and Amnon Shashua. Which transformer architecture fits my data? a vocabulary bottleneck in self-attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro-
ceedings of Machine Learning Research, pp. 11170-11181. PMLR, 18-24 Jul 2021. URL
https::/ppoceeciie
mlr pressppessss.... http//////////eeeeeeee wttp:/weelaa......m wwtsela..tsmmllll.aaaaa....m.
Adina Williams, Nikita Nangia, and Samuel R Bowman. A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint 1704.5:70. 2017.
A PROOF OF COROLLARY
Proposition Let Si = ( { و
j=l and S2 = [u {
1M j=1 be two sentences, gw the Transformer
I.L.d,
operation of y in-cooneext (S1, S2) and M be a vocabulary embedding matrix. Then:
sep-seq
L.L.d innconteea
(21, S (EPP((N),)NNNN))) Bw
J.L.d.
iLid2 Proof. Assume that sep (M)((N))))))) (Rd:)* N R and gí. ......999 I' (Rd))
gy
= R, then by definition there exist 9l,...,9R
IR such that for any { (x' J==1
2N
d. (x2,...,,?3) gw
X2N)
D gr (x.,...,xX)
x²) go N-1 (x"*)...,,22
x2N)
Now, given a, b E Rd, write:
y d. (2) S2) (a, b) Σ gr (Eary (wt) aa.... Easy (wf*) o a)
g (E (wb) @ b...., Eny (w @ b)
Clearly, this form of presenting Z y has R summands, conclude that:
V S3)) (a, b) is separable with respect to (a, b), and since
sep-seq (y Lo
(2), sep(a,b)
(yyt. (S), R
Corollary now follows from an upper bound on SeP((N)),2NN(M) (2020).
1.L,d2 BW
given Levine et al.
B UPPER BOUND FOR THE E-SEPARATION RANK
Definition 1. For an expression that can be represented as a sum of some terms,
an
12

---

Published as a conference paper at ICLR 2022
denote by f the corresponding sum, but with each term replaced by its absolute value, that is:
N
ann
and note that by the triangle inequality holds that:
11
L...). be be the E [da] entry of the analyzed sequential represen- Assume that all learned parameters and all gradients are bounded:
Theorem 3. Let Isequential
tation defined in eq.
5
40 (W, MV Amin S 10 102(5:)/al) S Amax for some 0 Amin S Amax N < dgg
n E (0, 2(14ridz < 8h, 21 11) de 3", In addition, assume that there exists M 2 ffr which
holds that zt
M on its domain. Then:
(1, S2)
log sep
(prii)LLdam Vscqucental
= Ô ([L 0.5 logs (n)l dz)
Proof. Denote:
Z (S2)
(b) = 2 p.s.H. L. (a, b)
The proof outline is as follows:
We start by finding a representation of (2Ss)
(4, as a sum of terms, where each term separable with respect to (a, b). We then turn to finding a subset of these terms, denoted G, such that the sum of all terms in an e-approximation of 2 (m. Lastly, since it follows from the definition of the E-separation rank and the construction of G that (a (z H. L.da (S (Si.Ss)) is upper bounded by the cardinality of G (which is the number of summands in the approximation), we find an upper bound to |G]. which therefore an upper bound to (a,e)(a.b)
Z L. (S as well, Negammal
which by definition is equal to e-seq-sep
H.S.H.L. H.d.. Vesquentaal
STEP A SEPARABLE REPRESENTATION OF Z H. (
(SrS8s)
Following Levine et al.
2020
Wies et al.
2021
(8a)
(b)
can be written as:
(1,1(1,313))
z (82)
(4,1 (b)
N
da
Σ
(0,))
Jaundees =1 he(M))(()) P
(13)
(C(L++1
(CL)
II
(((ch)
wie)
II
@rets wie)
where wt = g (b)² + na (a)" ② (this form follows from g (b)² = Exy (w 3 the entry- wise product of b with the embedding of wy prior the /th training step, and f (a) 8c(S. a) ( >,g
the gradient update performed to wijs embedding at time f), the p(c.A) and ((e.h) terms are sums of products of the networks inner (i.e., non-embedding) weights which were also updated with respect
13

---

Published as a conference paper at ICLR 2022
to C (Si;a), and for convenience we denote jC(L)+1 = i and P(C(L)+IM) = ploM)
E
Σ
Q
KE[H) ((u) POi
C(L)+1
II
plah) 9 (b)² + nf (a) o 0b)
To
C(L)
II
(o (ah) 1 (b)² n/ (a)a ob)
Separating to vocab- geaab-gradient terms and vocab terms:
N Σ
da
(10.h)
Jumocaa = (()) -0
involving
II plahh) nf (a) bb) (a)))b))
II eElq
(olet) (f) оb)
\eETp
involving
(plah) 1 (b)²)) )
II [E]C(L)+1]Jp
II CE(C(L))VO
(olat) (b)a))
Terms invalving just
Opening indices:
da
Σ
Σ
Σ
Σ
C(L))) Janndow Io((C(L)
Σ
II eETP
plan) nf (a)dr, bao
II
((a))) P of (a)az. bg.
eElQ
=1
plc,) Te g (b)
II oE(CC())N
((e.h)
0 (b)
cE(C(L)+I]MI
Separating to weights and variables:
N Σ
Σ
Boun
Co C GCC++1 =1 (oCL)
PpC 11 Joudd 10 1o5(C(L)
II keEp
nf (a) bas
II nf (a) b
neTa
II (E|C(E)+I]Ve
g (b):
II
g (b)/
oe(o(L)iVo
14

---

Published as a conference paper at ICLR 2022
where:
de
E
((,,))
PrP
kE(H) In Tou
cCL+1
C(L)
plac)
Il
Q (aak)
Tosaa
1
Compressing summation to count variable powers:
20(L)+1
Σ
E
NA
Pi + tPd NA ny nd, 20
Total power of (a)
many are equal to each of (d.)
Σ
Σ
fo m. mi 20([)+1-
OSPL.I Pdda N NA Yoeld, [4.] 2 Pa.; Pei MJE[N] Ede Po =z
VEENNN)) 2, m,,m mod? #1tm,El mod 2
How many indices each JE[N]
How distribute powers ((a)
Σ
AN..P.
US 1o Tidy S2C(L)+I-N YaEld,) EN 7o [JE[N] di Pao, =m,
How to distribute the powers of g(b)
de
de II == a=1
IIII
nf (a)² ba Pass
(9 (b):
aai
where:
AN..P
(oCIC(L) Mrlt
de
Tan...BBrL)))
15

---

Published as a conference paper at ICLR 2022
Pushing in summations on N, only the parity matters:
20(L)+1
Σ N4=0
Σ
Σ (0.1)) NV
P Pd. Rad -CC(L)+1-N
=
>n ey mod 2 b 6y Smii (N
AN..P.R.
Network's is weights, function of
N de
Σ
II (1(a)))
Pais
d NSN NA 21 t
f=la=1
VaE(dal WN P V) dazz Pa. z5 VjElN] Ee mod
A.p.per function of
d2
d2
II (ba)
0S Rdg -- j=1 (b)..
E
me m my (0())
YaE(d=] E Aa.=maa
VjE[N] D =a;;=;;
WEE[M me mod my (1-6)
e.
20(L)+1
Σ N
Σ
=
NA=0
- ndy -2((L)+1-N
Σ
ANaDM OADe" T Bane
X eE(0,1
N mod 2
MN Smin (M)
where each summand is separable with respect to (a, b).
STEP 2 AN APPPOOXIMATION OF H...d (
VERQURNTAL
Now that we have a representation of Z L. Vasucutia
L.d. ( as sum of (a, b)-seabreble terms, we
turn to finding subsum that can approximate S. I. (( up to an e-precision.
First, let define the set of all legal indices in the last sum:
02NA520(L)+1
Id= P.=NA de n. -C(L)) VJE[N], e E mod Ey (
D D2c(1)+1dd.. (NN.p.ee)eNPtttt
16

---

This will ensure that Za is an e-approximation of Published as a conference paper at ICLR 2022
And for G C D, denote:
2g b),
Σ
AM QADe
B.n.
(N..D..e)EEG
which is the sum of all terms with indices in G. Clearly, summing over all possible indices gives us the original expression:
DD (a. b) Vocaaniad (S (a. b)
Given O, wish to find subset of the indices, G C D, such that the sum of all terms whose indices are in is an e-approximation of 2 Viqquettia "(S. (a, b). That is, we are looking for
G C D such that for all a,
Ep (a, b) - Za (a, b)l = Zpic Se
Note that since we assume that 2 (a. b) =
< M, we can get:
Z (a, b)l
DR. (a, b)
E (a, b)l
2 () <
M
2 (a, b)
2 (a. b)
and it follows that is enough for us to show that:
2 (a, b) 2 (a. b)
1 6
M
which is equivalent showing that:
ZDiG (a, b) 2g (a, b) < M
(14)
under the assumption:
OC(S))
YB E +
0,
if (a)) ba. 9 (b)² [Amins Amax]
(15)
00
which we make going forward. p.H. R. Rapperaa (S) aa. b).
Now, assume that v0 E O :9:00
80 [Amin Amax] and by Levine al. 2020 the Ps and Qs in eq. are products of up to L matrices, each of their coordinates is bounded in [A A and we assume without loss of generality that Amin S 1 < Amax (otherwise we could have picked a smaller Amiu and larger Amax). Then for each (Na. P. n, e) E D the following inequalities hold:
8C(S))
C(L)+1
(CL))
Σ
Amax
II
Amax
II Cmm
Amax
hE(m Cum caLl -1
C=
dd((D)d,A20(L)+ 2202))+
ANA.P.
Σ
(OL) 2
di ACCL)) (CLL))+ max
1PCIC(L)+1 (OCIC(L [
VkEld, MoE Hitp://e....e.iitoooooeooo../
(2C (20(L)+1) +1 -20 NA
( (
NA
)
20 (20 (L) +1 1- NA NA) ((L) (1)),,A20())) 20(L)+2
Pd.
M1, ... ndd
17

---

Published as a conference paper at ICLR 2022
N dz
Q.p.p $ Σ
Pa. max
ISPLI N S.N
Jim a=1
21t1488 NA
VoE(d,) Σ Paij Pa
[JE[N] Eos Pa. MjE[N] F3 me
E
A max
a Pnm N S.N. #1 =
VaE(da] Σ Po.s=pa
VjE(N] Eas Po.s =2j VjE[N] Sey mod
d
da
V V.P.pncc S
A Po max
Σ
Ra
max
05 H4 (2C(L)+1 j=l
mi mi+.+my -
VaEld,) > Ba.j
[JE[N] Ede Bay (e mod m;=(1-8,) mod
Σ
A
2C(L)-1
OSAI (20(L)+1-N, my mutt+mm 1C(
VaEld,] En= =R
mmy Vje(N] Ede mmy YjE[N] my=ey m==((1e)) mod
and therefore:
ANaPO OApe B.p.n
Σ
Σ
1. N SN
OSP1 SEPII.....PA Pde (20())+ mi - H1
20 IN NA
VaEld, EN- Po. Bo
Vogld, Σy Bay mn
[jE[N] Ede Pai mzz VjE[N] =y=e mod
YjE[N] Ede Ray==m
ViE (JE(N] mey mod
m1=(1-e,) mood
20 (2)(L)+1) (1))+1) NA
(p.,...ppp
NA
20 (L)
d d (2C10)-)))))0000))111 max
d
m1,-.. Ad.
18

---

Published as a conference paper at ICLR 2022
Relaxing the parity constraints inside the brackets and recalling that Amax bound:
gives upper
na
(2C (L) (L)+1) +1
NA
2C (L) + NA
NA
Pi,...,PP
Pd.
.1,...., Ad,
which further bound using lemmas and in Levine al. (2020) until are left with:
e (2d, + +2C(L) 1
2d.N
N
2C (2) + +1)
NA
2C (T)) +1)+1- NA)
Pl..... PI,..., Pda
NA
.1,...., ,dd
On the other hand:
NN...n @A..p B.p.n.
20 (L) (L) +1
(
NA
2C (L) [ +1 -
(I....,pd.
NA OLL do ((2(())2))
.p....,Po..
B1...., 1dd2
Σ
A min
OSPL Pul N S.N. 214m04g = NA
YoE(d,) Σ Pa
VJE[N] Ed=; Pa.. =2 (JE[N] =e mod?
Σ
A 20(L)+1 min
USB1 by.. Rdg N 1-
+...-mx 2 (L(+))) Vor(de) N
mna
AJE[N] Ede =m, ()E(M) my mey mod
(4=(1-e;) mod 2 NA
20 (L) +1
(2 (2 (L) +1 +1- No
NA
Ph....,Po Pa.
mI,..., ndd,
20 (L) +1
NA
2C (L) + +1- NA Ri....,Pdd Add
(o(()) do(
( (2+2) min
NA
PI....,Pa
Combining the upper and lower bound for ANw.p.n ΦA..e V. we get that for G C D:
Epia b).
Za (a, b)
e N2d + 2C (L)
2d.N
Amax Amin
(L+2)(201L)+2)
d,N
(20(L+1)) ENNnpmeeEDD 1f (2
No
(20(L)+1-NA)
Pa
EINADOO No (20()++1 (((CC)))
Nr
(2C(L)+I-Na
pp...ppa
Bi....idd
19

---

Published as a conference paper at ICLR 2022
so in order to show that 14 holds, it suffices to show that:
LNNnpneeDDIG N (2 Manaa
(20(L)+1-N))
Einnpmeeee 11 NA (20(1++1
Na
(2C(L)+I-Na)
Pi...aa.
N e(2d, 20 (L)++1)
2d,N
Amia
(2+2)22(1)+22
Amax
M
Now, we can limit ourselves to subsets G C D of the form:
G (T) = (NN.....))D (20 (n NA (2)
and in this case we get that:
NA Emmupeeeeeaam 1 (MCCC))
(C(L)+I-NA)
Pdo
(NN 11 (2C(L++1)
(
Na Pa....pa
(2(L)++-NA)
N₂
E(Map..)EPp)am T
(eNeEi))) T DI (T)
G (T))
Let us define:
(EN,520(L)41
D Dro(L)+1,, (NN.
Sode De NiA
Ede no 20([))-1-111))
(T) = (NND..)ED N (app20))4)) Moyss m
(2CCL)41-NN) NA)-T
and note that:
(TT) s [G(T) |(T))
D S DI < 2 |
and also:
2C (L) + .)) |
(2c 2d
(2dz + +C' (L) + 1)
2ds
d,
where the inequality is due to lemma 3 in Levine et al. (2020).
Hence:
2N (e(2d, (CC(L))))
24
-lā(T)
D (T) IG (T)
d.
G (T)
And after rearranging we get that for each T 2 such that:
N (e (2d, + 2C' (L) + 1)) 2d
(G(n) 2
(16)
1+ e(2d, (2C
N.N
2d.
(Amia
(L+2)
ddd
Anmax
M
In order for Zorry (a, b) to be an e-approximation of M,n., ((S. (a,b). b).
20

---

Published as a conference paper at ICLR 2022
(z (Z, L. (S,
STEP 3 AN UPPER BOUND (-SEP(a.b)
In the last step we have found condition on subsets of indices, such that summing over any subset who meets this condition will yield an E-approximation Z 2.p......d...S.... (S,, We will find specific subset who meets this condition, and it in order bound Besesnnn p. (S.. from above.
Regania
(a))
We will focus our attention on T's of the form:
m2 147
2C (L) +1
(()) =
1(2C(L)+1)
I+.
x2C(())+1
20(L)+1
1+7
1+:
((2C(L)+1) (1+1)da
20((()) (1+7)d
2011)+1 (1+7)da
20(L)+1
(1+m)da
for E (0, e-1.] which we'll determine later.
By lemma 9 we have that:
((((())
me (2C (L) + +1)
2d (1 + 1)) (In (In(s-))) "- "3
2 de d=
so we can choose:
s==exp (11 (TEH, (T (TC
22N (e(3d +2000) 11)--
2dz N Amii (1-2)00(1++)))
# 22 (ma)20)))))))
and since this upholds 1-5 we get that 16 indeed holds for G (T (s*)), and therefore that (x)) is e-approximation of Z p.i.H. d.d ( So)
Now, from lemma we get:
(TT())) 2 2 (2)
2T d,
(dz 1)²
(dz 1)
22N
(2d, 2C (L) 1))
N.N oldd (2CCLLHH)
2d,NN
(Amia Amax
(1+2)(20(1)+2)
2
de
and therefore:
s-seq-sep
Uhequennial
-Eseep(a.b)
U L. (8)
(T
52N (T (s*))
2 222 (L) +1)
(00.(225.) /T ds-1
(dz 11²
(d, 1)
23N (e (2d, 20 (T.) 11) ddz
2d,N
(L+2)(20(L)+2)
(L
2
1+
d.N
Amia Amac
Rd.
21

---

Published as a conference paper at ICLR 2022
B.1 LEMMAS FOR ESTIMATING THE NUMBER OF COFFFICIENTS
Remark 1. For brevity and clarity we will use expressions of the form (g....)). regardless of whether K is divisible by M not. For the latter case, this expression should actually be:
K
(1
(K
mod M) times
expressions of the form ig
K Brim 2o +1)
should be read as:
K
(K mod M)+ rimes
and expressions of the form
(
K w..
-11
should be read
x
iK mod M = 0
Jos vance
K
otherwise
(K mad MF staee
Lemma 1. For all K, .M E N, the maximal value of asx) holds that (زj Ojzl S 1.
maw)) is achieved when Vj1,22 E [M]
Proof. Let a1...., am be a sequence of non-negative integers such that +... + as K and (a.....as is maximal. Assume towards contradiction there exist J1.22 E [M] such that
K
aj Ujz
#1
< 1, then:
d3s
K
K!
01....., aM
-== az!
K!
Ajz +
(IhEMM)Mn))a)) an!) (aj - 1)! (aja + 1)!
aj
K!
[IEEMIMAS asl (aj, - 1)! (ajz 1)!
K
a1 ai....., -.....,a, + 1, ..... amd
in contrary to the maximality of K Ragg) Therefore, Vji.12 E [M] jag, azzl 1
Lemma 2. Let K, be two fixed natural numbers, 17 E (0, 1] and denote
(n) =
K
17
7
n1
71
K-n M
nn
M
Then for all n E [K]
[ -11 II
(n(K (j()- as
(nn)= =
(1+1)M
1=0
(17)
mod M
1/ (K m M)
K
[ + 1 M
K
M : .....
M

---

Published as a conference paper at ICLR 2022
Proof. We will prove by induction.
Base case: n
n
K
(
17
0
0
K-0
K-0
MI 1...
M
M
[ -1
1) 7 jM)
as
II
(j 1
J=0
0 mod M
"
(K- L M) MM
K
(L
+ M
M
M
where the second equality is due to the fact that:
L
[ )
M
πI
=1
j=0
an empty product, and:
n (K 0 MM)
0
mod M
1 M
Therefore,
17 is true for = O.
Induction step: Let n 2 0 such that 17 holds for n.
S (n
( K
17 nt
n +
+ .))
K - n -
11 + 1)
K-n K-n
M
17 (K 12)
S (n)
(3r + 1) M
n(K - n)
n "K (K- (K jjM 16
II
n2 + 1) M
(j+1) 1) M
j=0
" mod M
(K
M M)
K
(L
+ 1) M
M
M7
--1
(n (n jM)
1J+1)M
j=0
(n+1) mod
n7 (K
x+1
M)
K
M
n+1
+ 1)
M
K
K
Thus. 17) holds for n + 1, and the proof of the induction step is complete. Hence, by induction. [7 is correct for all n E [K] UK]UK (0).
Lemma 3. Let K, M be two fixed natural numbers, and E (0, 1). Then the maximum of:
KK'
K n
n"
(1,,...,M
011
(1y...,b .ba
23

---

Published as a conference paper at ICLR 2022
is achieved when:
n = (1 n - M
'/K 1+
( (1+ +1)
M =
and all E [M], la S 1 and (b. M
S 1.
Proof. From lemma we know that a multinumial coefficient reaches its maximum when the sum evenly distributed between all indices. Since the a,s and can 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 [a;
ad I and b;
7
K K-n
M S
Define S as in lemma 2
Note that the function:
K (x)
T
K-I
M
M
Gaussian-shaped and therefore unimodal and has unique maximum. 17" for E (0, is monotonically decreasing, and therefore their product also unimodal.
Since we are only interested in solutions for n E N, we can reduce our problem to finding the maximal such that:
(n) (n 1)
Sn ()
and from lemma 2 we have:
S (m)
17 (K
n-1 M
2
S (n 1))
n-1 + 1) M
n
R M
M
<
M
1 1+
So we get that 8 (n) is monotonically increasing long as integer for which this condition holds is:
M
nK-M and the largest
I-
(L n - M
(1+ (
+1 M
n
1+
Lemma 4. Denote by N (B²) the number of integer lattice points in By (the d-dimensional zero- centered ball of radius R). Then:
ne
2R
-1) d
ne
12
S (BA)
V ad
V nd
Vd
Proof.
Let I (B% be the set of all integer lattice points in B For x E I (B define:
C = xx=
+ = F
ViE (d
2 S y こ
Let y E UnE Ca, so there exists E I (B such that y E Ca, and from the triangle inequality we get:
1 V
1 = R +
Vd 2
|
| + |/
æl| S R+
2
2
and therefore y E Bd R4
Since y was chosen arbitrarily, we get that:
v7
U
Co C
Bd R-
va
xEI(BTg

---

Published as a conference paper at ICLR 2022
Note that Vol (Ca) 1 for all a E I (B% and that for x' E I (B%) such that a + x', Ca CCd
a of measure zero, hence:
N (B%) I (BR)
Vol
U
C
#EI(BA)
Vol
13dd
R+ اهمد
R+ Vd
r (2 + 1)
Assume for convenience that d mod 2 = 0, so I' yields:
(5 +11 = - !, and Stirling's approximation
n
Va
d
rc 2
2R
d
R+
1
VTd.
2
V
Vd
2c
On the other hand, note that Bd
S
(1:) Ca. and therefore:
R-
N (B%) Vol
U Co
+EI(BA)
2 Vol (dd
R- vis
R Vd
r ($ +1)
12R
d
1
VTd
Vd
Lemma 5. Let K, M be two fixed natural numbers, (0, 1) a constant sensitivity parameter, and let"
TR.M = { MaENM
(,,""
K
K
(ar,...,aak aM
2
K M
Then:
K
2
K
a M a T K,
a
In
(s
-1
M
M
C TK.M
K
C NM
a T
K,
I
4K In (5-1)
M
Proof. By Stirling's approximation, holds that:
a! 2
and slightly more accurate version is:
1
a
z! ~
2m
x +
6
25

---

Published as a conference paper at ICLR 2022
and by plugging this approximation to the definition of a multinomial coefficient we get after some rearranging:
K
K
+
M
(2T))
,1, ,0M
ii=1
a
exp
(E a In (x)))
Note that using this approximation and some rearrangements, we get that the following are equivalent
K
K
01 01,...,aM
AM
2
K M
K +
M
LM
6
(2T))
i==1 V
+
6
1
exp Ei= aj In aass)
1-M
K
+
6
M
2 (2T)
2
(18)
i==1
M
K
exp
Ei=
K M
In
R 18
M
M
M
K
II
exp
Σ
a
In
S 5-1
(19)
K
M 6
so we can characterize the set:
TR.M= = {a E NM a i=K (ITi vactd)
exp a Ia (a.))aa)ssa (15+
as T K.M TK,M-
We'll start by finding condition that will assure us that (a.,..., ar) E Tr.M (ie. we will characterize subset of T T..M).
First, note that by the AM-GM inequality,
Σ i= (as
+
K
+
M
K M
II
a
M
M
and therefore:
A
M
K
1
II
04 +
6
II
a +
6
M
+
6
Since we're interested in a subset of TK,M, we can show that the last inequality holds when we
replace the first term in the left-hand side (I v Va: + f with ( + +)
# (if the new inequality
holds, (99 must hold as well), and we' we're left with:
K
1
*
M
M
(
K
4
+
exp
Σ
az
In
K (1)
S st
+
M
6
M
6
M
M
In
< In (5-1)
(20)
K
1=1
Now, for i E [M], let hy ==
M
1, so:
K
M
M
M
Σ
0j In
K a
M Σ
(1 + hn (1 (1+hi)
and the last inequality becomes:
K M
(1 + h:) In (1 h2) 5 In (s
M
i=1
M
M
(1+h))((1+h)>
In (1-
K
26

---

Published as a conference paper at ICLR 2022
Observe that for all x 2 -1, it holds that:
In (1++)) Sx
and therefore:
2) (1+ x) S x² + x
so it suffices (again, we're only interested in a subset TK.M) to show that:
M
Σ
(he + h.)
M -
In (s (
(21)
K
1=
Now, observe that:
M
M
Σ h4 Σ
K ag
1=1
1=1
=k
M
MM
K 1=1
0
so becomes:
M
h2 < In (5-1)
K
K
2
K
In (8-2)
M
and we get that:
}
K M
K
AEN" a
a
TI
K,
a
i
In
(s (5
c TK.
M
Let us now turn finding condition that will assure that (a1,..., as) & TK.M (i.e., we will characterize subset of the complement of TK.M). Note that:
M
II
Σ a; In
a.
exp
K
K
M
M
M
1
6Ma;; M 6K + M )
exp
Σ
aj
In
+
In
K
2
so if (a (a1...,,as) aMM & TK.M, we must have that:
K
M
1
6Maa; M
6K + M )
exp
Σ
ag In
0g
+
in
K
2
i=1
+
M
6
M
.
6Maa; + M
)]
0; In
+
In
> In (s-1)
K
2
6K + M
27

---

Published as a conference paper at ICLR 2022
and using the same definition of h; as before, we get:
M
M a2
1
6Maa; + M
0g In
+ 2
In
6K +
+M M
M
+ 1 In
66K (1 + h4) + M
(1 + h1s In (1 + h.)
M
2
6K + M
K
(1 + h:) In (1+
M
[6K (1 + h1+
K
3K
2 In
hi
6K +M
M
6K +
M
where the last equality is due to the fact that
M i==
h4
Observing the first order Taylor polynomial of the function:
K (x) M
(1+ (1 + x) + In 6K (1 (1+x)++
x) M
K
3K
6K + M
M
6K M
at x with the remainder the Lagrange form, we get:
(
K
3K
h2
K
3K
(h,)
h4 + l" (5:)
h4
M
6K
+
M
2
M
6K
M
h
f (E:)
2
K
18K²
h2
(22)
M &1++)
(6K (1) + M)²
2
for some &, between 0 and h,
Note that f (E.) is monotonically decreasing with & for & E (-1, M 1]. and using this fact and the fact that K 2 1, (22) is lower bounded by:
A2
4M2
So can limit ourselves to looking at the cases where:
M
K
h >
In 1s
4M2
1=1
K
4K 4K (5
0j
M
1=1
and we get that:
AEEN a
aTi== K,
K
_
a
>
M
5 5 a-I=='' \TK.
c (a NM a-i K,
I I 2
+ TK,M
a
<
M
Combining the two results together we get:
A a I
K,
a
K
<
K M
In 11
M
TKKM
K
2
C a N NM
a Ti=
K,
a
i
<
4K
(8
M
28

---

Published as a conference paper at ICLR 2022
Lemma 6. Let K, M. be two fixed natural numbers, and E (0, 1) constant sensitivity parameter. Then the number of multinomial coefficients, (a a which uphold:
K
K
K
2
a1:..., RM
M
upper bounded by:
111
M-1
Z
In (8-1)
M -1 +1)
(M
1) V
and lower bounded by:
M-1
M-1
nC
K In (s11)
2
2
-
M
V
M
Proof. Let (aj....oss) be multinomial coefficient for which it holds that:
K
K
2
K
a1...., 0M
and denote:
K
K
TKMM
AENM
E NM
2
K
a1..... aM
M
K
K
T
a NM
a 1 K,
a
i
In (s
M
M
K
T
a E C NM
a T K,
a
<
4K In (s-1)
M
then by lemma
TKMM TR.M S STK.
and therefore:
Tk.M & (T.Ml & |T
so in order to bound the cardinality of TK.M (which is the quantity we are interested in), we can find an upper bound on the cardinality of I K.M and a lower bound on the cardinality of T K.M
Let B E {L, U} and TB K.M- and denote:
V (s
=f L ln
RB)
if U
For E [M], denote T; 0, So the problem has changed finding the number of integer
M-tuples .1,..., such that ΣM 0 and [<]] < R RBB), which is the number of integer lattice points æ in the zeno-centered M dimeensiooa ball of radius (B) that uphold E. 3; 0.
Note that the intersection of d-dimensional ball of radius R with the hyperplane H yy E IRd y. I) }} is a (d - -immeenioonal ball, so we can assume that the number of in- teger lattice points in the zero-centered d-dimensional ball of radius R whose coordinates add-up to zero is J/2 / 11) where V/a is the cosine of the angle between H and one of the axis-aligned
(d - -)-dimensional hyperplanes in Rdd In our case, R R RBB) and d = M, so by lemma 4 the
29

---

Published as a conference paper at ICLR 2022
cardinality of TKi is upper bounded by:
M
M-1
N (BM-L ViKin( 1j
2 2 4K In 1)
+1
VM
M
a (M
VM-1
M-1
me (M
Kin(s-1)
(V
M- +1
1) VT
and the cardinality of TK K.M is lower bounded by:
M--1
2 K In (s
1
[BM-1
W
VM
V 17
In(=-1)
v VM
h (M 1)
VM -1
M-1
ze
(VKin(s1)
2
M -
M
VT
Lemma 7. Let K be fixed natural number, 1) (0, 1]. and E (0, 1) constant sensitivity
parameter. Then number of integer us such that (H) n" 2 s n 179 is upper bounded by:
ak
K V (1 21n (s)) (1+1) n) 4n
211 + 11) (In (s-1) -11 1)
Proof. Denote n nn 1+1
+ x. S0:
K
K
nK 1-
n" 2
y
n
n
I+m
K
17th
n
nK
1+1
Recall that by Stirling's approximation we know that:
K
K
KK
2nn (K - n) n (K n) n
So the number æs which uphold:
K
K
2T
It7
+
xz)
1+.
K
-x)
nK
+ x
nk +x
(
K
K
z
1tn
1+n
K
KK
8
2T
I++
K
nk
ITw
K
k
1+1
I+n
ak 1=
to
I4 +1
nK
K
In
8
1 1+
1 1+
(23)
+a+
njK
ln
+ r
17
I+
1 1
approximates the number we trying to quantify.
30

---

Published as a conference paper at ICLR 2022
After some rearranging, one can observe that
23 is equivalent to:
nk
(1 + 11)
In (s-
2
+
+
in
1
1 1+11
2
n
(24)
K
1
(1 + 17) x:
+
x
+ 2
In
+11
K
and since for all f > -1 the number of integer for which it holds that:
141 I In (1 + t), the number of æs which uphold 24 is upper bounded by
In (s (s- 2 (1-1p) Kr - 2(1 1)) x²
2n + 2 - 172) -2 =7)²² x²
- 2( + 1)²² (In (1 -1)=²
(25)
+ (1 - 21n 1)) (1 - -1²² Kx- 21n (8-1) nk²
n 00
Recall that for inequalities of the form ax² bæ S 0 where 0, the set of all values of which satisfy this inequality is
-b- 4mc
4ac
= bb v And
and the number of integer values of (the interval's length).
2a
which satisfy this condition is approximately yha 4ac In our case,
25 gives
= 2(1 + =)²² (In (s (s-1)-11 - b = (1 21n (s-1)) (1 (1- K
ln (s 1) nk²
and therefore:
K KV (2 (s- 1)² (1+ 17) 4n
12 4ac
a
2(1+ 17) (s- - 1)
Lemma Let K, MM be two fixed natural numbers, n (0, 1]. and E (..| constant sensitivity parameter. Denote:
DK. = (n, N2M+1
OSnSK bottttbv-=-nn
define F: DK.M
IR. as:
(K)
K (K- (
F (n,a,, b)
17" a1...., aa bi,..., bys (b.....bbs))
and let x* arg maxF (x). If ςk x
147 2 (M 1), then the number of x DK.Ms which uphold
F F(x) 2 s. Fx*) is bounded from above by:
M-1
2K
25TTeK In 1s-1) Vi
(M - 1)²
2(M (1 +1)
Proof. By lemma
x*
nK
n
nKK Mf(1+) M(14n)
K
K M(1+))
By lemma 7
M(1+n)
M(1 +)))+)
the number of ns between
< < K such that (*) 2
14n
is upper bounded by
K 12n (1)n))))) by lemma 6 the number of non-negative integer M-tuples a1,.. a1,..., AM
21 ([+=))(((=))-1)
n and JK
such that 01 +... aM I+7
2 8
ITT
is bounded from above by
OK M0
31

---

Published as a conference paper at ICLR 2022
M-1
(M- ( (4V
nk In(s 11
+1)
and the number of non-negative integer M-tuples b1...., bbr is bounded from above by
(1(4))
such that bj +.... bw
K and
1+ bi.....ba
2
I+
M(1-9)
M-1
In total, without taking into consideration the interactions between the three multiplicands (so our bound is not tight), the number of DR,M which uphold F (x) 2 s F (x*) is upper bounded by:
(M-1)== V (M- (X-)))4)) +1 1)
In(e-1)
Ky (2
1)² (1 n)² 4n 2 21 + 7) (In (s (s-)s) 1)
binornial Tj expoment
nC M-1
M-1
K In (s-1)
" (M - () +1
(M 1) VT
a- muliinomial
M-1
M-1
xc
T
K Kin(s-1) + (M (1
4
1) V
h-mutiiomiil
nK and since (M-1)(++)
2 and S 115 (and hence 21n( )-1
2(nnss ) 22), this can further bounded
by:
2K (M - 1)² ा
25TeK In (s-1) vi
2 2 (1++))
Lemma Let K, M be two fixed natural numbers, 7 E (0, 1], and s E (0,e a constant sensitivity parameter. Denote:
DK.M = { a,, N2M+1
OSaSK taa. bettttber---
define F DK.M R as:
(k)
K
K bl.....
F an,
a1....,aM
and let x* == arg maxF (x). If 1+4 M2, then the number of E DK.M which uphold F (x) 2 xEDx.M
from 3. F (x") is bounded from below by:
M-1
TeK In (s-1) 2M (1 +11)
Proof. By lemma
6 the number of non-negative integer M-tuples bb. such that by is bounded from below by
+ bu I+n
and
b.....bx
M
In(s-!)
M-1
2
1+n
Since these bs are only subset of the elements of DK.
M //
which F' takes into consideration, this also a (quite loose) lower bound on the number of x E DK.M which uphold F (x) F (x*).
Now, we know that:
K
> M
+
32

---

Published as a conference paper at ICLR 2022
and since 8 S <e-1,, it holds that In (s-1) > 1, we get:
K
K
(s
ス M²
1 ++1
1 1+
In 1)-
I+1
1>
M
and therefore:
M-1
M
K
In 2)
-1
M V
M
M-1
M-1
K In (8-1)
In (8-1)
2
I+7
V
It4
M /
M
M
TeK In (s-1)
M V
2M2 (1 + 17)
C LOWER BOUNDS ON THE S-SEPARATION RANK
C.I PRELIMINARIES
C.1.1 TENSORS AND THEIR MATRICIZATION
We begin by laying out basic concepts in tensor theory required for the upcoming analysis. The core concept of tensor may be thought of as a multi-dimensional array. The order of a tensor is defined be the number of indexing entries in the array, referred to as modes. The dimension of tensor in particular mode is defined as the number of values taken by the index in that mode. If A is a tensor of order N and dimension M, in each mode E [NV, its entries are denoted Ad,...xx where the index in each mode takes values d; E [M&].
We will make use of the concept of the matricization of A w.r.t. the balanced partition (P, Q). denoted [[]r. E M/3 M/2 which is essentially the arrangement of the tensor elements as matrix whose rows correspond to P and columns to Q. Suppose A E RM M is a tensor of order N, and let (P, Q) be a balanced partition of [N], f.e. P and Q are disjoint size N/2 subsets of [[V] whose union gives [N]. The matricization of A w.r.t. the partition (P, Q). denoted [A]],; the M M 8/2 matrix holding the entries of A such that Ad...Ay is placed in row index 2/2 (d, - N/2 and column index + (dd 1/2-
1 1
form of function discretization (Hackbuschl 2012). Essentially, the function is evaluated for a set points on an exponentially large grid in the input space and the outcomes are stored in a tensor. Formally, fixing a of template vectors x)..... (2) (2) function y(x "....,,'),, the set its values on the grid arranged in the form tensor are called
We now present the concept of grid tensors, which
[V]. the points on the grid are the set ((((4)))...((())) )
xx(xx))x)
"d.....d== Given
the grid tensor induced by y, denoted Aly)d,.....d = yxx² x(11) x x(dw)).
1.2 e-RANK
We will make use of the concept of E-rank Alon al. (2013) of matrix A defined for any > 0 as the minimum rank over matrices that approximate every entry of A to within an additive e. We will prove lower bounds on the es for which the s-rank a matrix remain high by the following lemma:
Lemma 10. Let M E x be symmetric matrix and
0, then:
ê -ank (M) 2
Wk S n d (M) 2 =
(26)
2n
33

---

Published as a conference paper at ICLR 2022
Proof. Let
2n 2n. nxx we need to prove that rank (M E) 2 k. Since M is symmetric, .
diagonalizable with eigenvalues A 2 Ag 2 ... 2 An Denote by U1, U2...., Un the cigenvectors
that are normalized according to the 1 norm, then for any i S & we have that:
(I E) lell 2 || |Eu X Eu,l, A
0
(27)
2n
In particular for any < we have that (M + E) 0s # and since 01 ....... U/ are linearly independent conclude that rank (M + E) 2 k.
Finally, we will use the following lemma for lower bounding the amount of small eigenvalues of symmetric matrices:
Lemma 11. Let M E guxa be symmetric matrix with diagonal entries that equal to 1, and denote
its eigenvalues by . An. Then:
max
k
dx
(28)
T2
M
Proof. Since the trace of matrix equals to both the sum of its eigenvalues and its diagonal entries, we get:
n
n Σ [ = trace (M)
Σ
A SrA + n-T
rA 1r +1
(29)
n
Eq. 28 follows from the fact that:
| V Xp + . A1
(30)
C.1.3 DIMENSIIENSIIN SPHERES
We will use the Lehesgue measure on the sphere for taking expectations on Sd. For any measurable subset A C gd this measure is defined the d + dimensional volume of the "wedge" in the ball gd+1,
d++1 (ftx) x E A,t [0, 1]})
(A)
(31)
Ad+1 (Bd+1)
Where Xd+1 denotes the Lebesgue measure on 1Rd+1 We will also use an unnormalized version of this measure:
Munnormalized (A) = 1 (A) d++ (Bd+1)
(32)
Smith & Vamanamurthy (1989) showed that Munnormaiieed
(5d) is monotonically decreasing for > 5.
The following lemma bounds the rate of this decrease:
Lemma 12. For any d > 5 the following holds:
Munnormalized
(33)
Punnormalized (Sd)
2x
Proof. Since d > 5 we know that Munnomaliizdd
($d)
Munnemmaizzed (Su+1)
and therefore:
Punnormalized (dd
(3d-1)
Munnormaizzdd
(34)
Munormalized
Mumamalized (Sd+1)
Finally, Gipple (2044) 4) showed that
Parnernctedd
d+1
completing the proof.
Matoomaann
Ba41
2x
Finally, will use 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> the following holds:
Punnormalized (R
.d++1
(35)
Hunnormalized (Sd)
34

---

Published as a conference paper at ICLR 2022
C.2 PROOF OF THE LOWER BOUND
In this subsection, we prove theorem 2 of the main text. We will follow the proof' of Levine et al. (2020); Wies et al. (222I) with important adjustments to the e-sequentiall-sepaaaie rank definition.
We begin by showing that high E-rank Alon et al. (013)) of the grid tensor matricization implies high e-ssegtential-sepaaaaeee separation rank (see section 2.2 of the main text) of the function. Essentially, apply claim from Levine et al. (2020) to E-approximations obtained from the e-separation- definition. This relation, which holds for all functions, is formulated below for functions realized by the analyzed Transformer network:
Lemma 13. For Vin-context as defined in theorem 1 of the main text. Let E-seq-sep (1 (p.t),L,d, L, denote its E-sequential-separede
(Hin-contett rank. Then, for any integer Z, any O, any set of template vectors x()) ((),...,(())) E I& and any sub-matrix M of [A(2
llab it holds that:
Pmoniss
e-seq-sep
e-rank (M),
(36)
where A(Z 1, is the grid tensor of Z L.d, with respect to the above template vectors. Vincames
Proof. If e-seq-sep yin-contest (p.3), (o.d),,dd2 L.d. sume that e-seq-sep Vin-contest
oo then the inequality trivially satisfied. Otherwise, as-
K E N, and let j be an e-approximation for 2 (r.),a122 with Nusoom seq-sep (7) K. By claim of Levine al. (2020) we have that rank ([A(T)la,b) S K. De- note by M the sub-matrix of [A (i) that corresponds to the rows and columns in M. Now,
since ÿ an e-approximation for 2, 2y0. we have that |M x/|l__
S.. Finally, by definition
e-rank (M) < rank (M) S rank ([A()))))) K.
Relying on lemma 13 we will bound the e-sequential--eparaie rank from below via the e-rank
of sub-matrices of T 2 L..s). (.n...)))aa ]m,bb. Denote d == (d= H = 3 lemmas 10 11 assure
us that for n = the network's 's weights, as well as choice of templates vectors, for which there exists sub-matrix M E R of [[A[2 (o.9). ) a that is symmetric with diagonal entries that equals to and with
Ω (22-dd) it is enough to prove that there exists an assignment to
|| (V(d+1) nē, in order show that e-seq-sep
(p (..)).Ldd. Vin-context
for S 2n2
2.0+I
Now we will use a corollary that is direct results of the proof in Levine et al. (2020) This corollary shows that if 2 (x.)).ddd is able to produce vectors that do not change the analysis in Levine al.
(2020), then for any matrix B E rxd with rows that are 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 is equal to3 M (BBT) @A Importantly Ad is symmetric. addition, since the rows of B are 12 normalized, The diagonal entries of M equals to 1. Therefore, M upholds the assumptions of lemmas 10
13 and is enough to find B for which ||M|] S (V(d+ ((dd+)) na.
Let d, X > O, assume that for any matrix E R(a with rows that are l2 an assignment to the embedding
Corollary 2.
normalized, there exists a choice of template vectors x(1) x!.,...,, (2),
[[ 31-2
layer and the first self-attention layer key and query weights, such that for any J1, E
the output of the first self-atten attention layer on (jr. + ((x=2)))ur is:
H
WO...
'ee ignored the constant that appear in eq 28 of Levine et al. 2020 since we can get rid of this constant by dividing the last layer output matrices. Importantly, this constant is larger than 1 and therefore the network weights boundedness assumption is not violated
35

---

Published as a conference paper at ICLR 2022
for
Aji,((a)
(a- 1) mod de < 1=1 A ((a) S d
?
A 12.0(a- )ard)
d--1 < (a 1) mod da < da 1 A Φ(a d--1 Sd
Ya E [dz]
ua
2
2.V
(a- 1) mod da da -
Otherwise
where o(j) =
() mod da) +
Then for any matrix B E Raxd with rows that are l2 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 (BB*)* SA
Now we will shows that indeed L., is able to produce vectors that do not change the analysis Misoomen
in Levine et al. 2020) and the assumptions of corollary 2 holds.
Lemma 14. Let E TR ( ) xd with rows that are l2 normalized, then there exists a choice of x(2), an assignment the embedding layer and the first self-attention the output of the first self-attention
template vectors x
layer key and query weights, such that for any J1, [( ) layer (ji,82
( 31 )) is:
x(l))
for
A Ajjo(a)
(a 1) mod do < - A x(()) <
A 12.00a- (a-a))
d, S (a 1) mod da < du 1A p(a do-1 <
Ya E [dm]
un
2.V
(a- 1) mod da da -1
Otherwise
where o(j) / [i-1/an] (do - + mod da)
Proof. We will ignore Z a,.,,... element-wise multiplication with vocabulary embedding matrix Visamtel
by choosing Vi,j M,, = 1 (by the terms of corollary 2 suffices to find any assignment of the leamned weights).
For any i E (2 ((1==)) )) + +1 our templates vectors will be:
(
Ai,00)
A mod du x da sola)sd
(siin jiss ( -
A { --1
^ da- s(j-1) mod du <dd, IA6la 1oo--da-11 do=
N
N
(j - 1) mod do = d,
0
Otherwise
We will implement summation of the inputs embedding in the first self-attention layer, we will follow Levine et al. (020) and set the first layer self-attention key and query weights to:
W W.K.I.b
K.1A 1if
W.Q.1 1.
1, 1Aj
36

---

Published as a conference paper at ICLR 2022
This assignment implements summation of the inputs embedding in the first self-attention layer since:
N
(xx,))
)
ΣΣ
W.LL (Mvut o x((h)) ..K.I. (Mvut @ x((d))
(37)
W.J.A.
(Mvwy @ xx(1))
(38)
N
H
WRilk (Mvug o x(())) w.K.M. (Mvwe @ (x(e))))
Am1
(39)
(Mvwg W.s.A.V.....
@ x1))))
(40)
(di) W W.l.hwihhhhh
x du
(41)
da
(da)
Xd
(d2) Wo..MppVai
(42)
h=1
H
W w... N (xld)) + x((s))
(43)
where (1) because WQQLIA W. are matrices that are zero everywhere except for entry (1, d.) and that all the entries in the vocabulary embedding matrix equals to 1, and (2) because of linearity.
((24") ))
Therefore, for any J1.J2 E
the output of the first
layer on j1, J2 is:
H
(1x))
Σ
W.O.M.W....
(x(()) + x((2)
(44)
Finally, we need show that indeed for any 21,32 E [( ((.x=)) eq 44
give the desired u.
Ajjo(a)
(a 1) mod do d-=1 )((a) S d
A j2.0(0- da-1)
d,-1 2 s (a 1) mod da < du 1 p(a d sd
Va E [d,]
un
2N
(a- 1) mod da d,
0
Otherwise
The third and forth cases clear from x's definition, so it remain prove the first and second cases. For this we will examine d1, d2. d1 = ji < (1
(sL² d and therefore:
(di) = (0- {A
(a - 1) mod do < d=ي1 A M()) S
2
d, < (a-1) mod d, < da - 1 ^ pla 111a-11111 d
d2 = j2 + (3/ )) E 1+ ( (sr )) 2 (g==a) )} and therefore:
(a 1) mod d, < dayi á((a) S d
x((2)
d_
2 < (a - 1) mod d, < <d, - A δ(a 1 <d
12.0
So it clear that also the first and second cases upholds.
Returning to finding B for which || S (v(++1) ni, we will use the probabilistic method for proving the existence of such B, Le. we will show that for random B the expectation of || S (vd++)) 3 and therefore in particular there exists such B.
37

---

Published as a conference paper at ICLR 2022
Lemma 15. For any d, X E N such that X 2 d the following holds:
(BBT) ox || - +(
E Bsyy
(45)
Proof. We start by bounding the expectation of the squared norm:
|| (BB") e6
(BB") 622
E (84ay
R (sdy"
(46)
F
E u.ond () e))a = n² u.usu (uu 2)22]
(47)
Now, E8 Euwwgg (u, 0) S (d+ 1) (1)
by lemma 16 therefore we got that:
(By (BB') 11pp
S (d++1) (1))
(48)
Finally, by Jensen inequality we have that:
(s-r || (BB*) | s E BBBi8y
(BBT)ON
S (d+1)
(49)
C.3 TECHNICAL LEMMAS
Lemma 16. For any d, X E N such that 2 d the following holds:
Eu,uns (uu 1)22]] S (d 1)
d
(50)
Proof. We will use conditional expectation make reduction for simpler expectation. From rotational
invariance of the uniform measure on gd we know that for every rotation matrix R E SO (d) and unit vector gd we have that:
E Buwgd (u, w) [u] Enrage (Ru, Ru) 22 [/]
Enngge ((u, Ru)22
(51)
where the first equality holds because R is orthogonal. Therefore, by choosing R such that Rv e1 we will get that:
E Bug (uu,))2) u) u)
Eg Ennaso (uu 8) 10/2
(52)
Eyaga Eungd U1 21|1 Eungg
[uPA]]
(53)
Now we can calculate the last expectation directly:
u 22 (u) - Sr
Munoornalized ((VI
x² gd-1)
Eunga [u
22 dr
(54)
Punnormalized 3d
Now, by lemma 12 and fact
we have that:
Munnormalized ((VI - x2 Sd-1)
d+1
1
x²)
(55)
Punmormalized (gd)
2x
Therefore we have that:
Ey (u, u) 2
d+1 S"
(1-22) dx (d (d+1)
L' (1 (1=
2x dx
(56)
2
Finally, by lemma 17 each term in the integral is upper bounded by (1)) and thus:
Eu (1u.))2) < (d+1)) 1)
(57)
38

---

Published as a conference paper at ICLR 2022
Lemma 17. For any x E [0, 1] and d, X E N such that X 2 d the following holds:
-
22 (1 - x² v
d
(58)
Proof. Note that since 2²²hA (1 x²)²
0 in the boundaries (xr {0, 1J), it is enough to prove the
inequality for critical points.
2x (1 - x²2
x²)$) $)
= (atx (1 x2)
2d
2 x
2 (1- - 11=22) x2)8-1
(59)
21
- 0 22 (1 - x²²) dx² x² =
(60)
2X + d
Therefore, x² = 22+d
22 is the only critical point and:
2A
2A
222 (1 22)
(61)
2A + d
2A + d
d
d
(62)
2A + d
2A +
d
(63)
3Λ
1.5 (A + d)
(V
(Ve) -
E
(64)
1.5
eea d)
d
$
-
(65)
e e(A + d)
where the last inequality follow from the fact that:
(
(A+d-1))
A + d
(
(cAA+d-1)
ee
(e(A+a)) d
<
(66)
d
d
d
D EXPERIMENTAL DETAILS
D.1 KNN-TAPT
We conducted the network training described in section 3.T of the main text with Adam W optimizer
(with the parameters suggested in the original RoBERTa paper: 81 0.9, 322 0.98, 10-6 and weight decay of 0.01), with batch sizes of 128 or 256 (depending on model size) and sequences of 256 tokens each. We started with pretrained RoBERTA- repository
weights from the HuggingFace Transformers [ and continued training them on the MLM task with masking probability of 15%, where each masked token had 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 finished, we evaluated the models' performance using the SentEval kit.
D.2 NNN-PRETRAINING
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: 31 = 0.9, Br = 0.95, = 10 and weight decay of 0.1), with batch size of 512 and sequences of 256 tokens each. We pretrained HuggingFace Transformers implementation of GPT-2 from scratch on Wikipedia with the standard LM objective, and switched mixture of the standard data and our generated kNN data in two different points during training. After the training was finished, we evaluated the models' performance on the Natural Questions benchmark.
https: httts:://buggiingfcc.oo/ http:///ww...g..tto..onooooooaae
39

---

Published as a conference paper at ICLR 2022
E KNN-PRETRAINING AT DIFFERENT CHECKPOINTS
The following table includes 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 KrN--rNPeriiit PrN--rretaining..
Model Size
Reg.+KNN Steps
NQ
F1
110M
200 400++0
< 10-3
110M
200+40
6.2 10-3
110M
400+40
7.9-
345M
200 / 400+0
< 10-3
345M
200+10
9.6 10-3
345M
400+10
Table 2: Impact of model size and regular pretraining steps on the Natural Questions F1 score of ANN-Pretraining.
F KNN-PRETRAINING ON ADDITIONAL BENCHMARKS
The main text describes experiments the Natural Questions dataset. We test how KNN-Pretrriniin affects other NLU tasks, by examining several tasks from the GI.E 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 2012). As in the case of Natural Questions, we evaluate the zero-shot performance of our models since it is direct probe to the abilities of the model straight after the process of pretraining. In contrast to Natural Questions, the GLUE tasks examined are classification tasks and not generation tasks, assessing zero shot performance straightforward. We therefore follow the template-basee method Gf et al. a (2021a) for converting the tasks' data into format processable by unidirectional language models.
them is not
Notably, the examined GLUE classification tasks are not easy for the examined unidirectional models in zero shot. Table 3 Jincludes the zero-shot scores of the 345M parameter model that trained regularly for 200K steps and then continued training for 20K steps of ONN-Pretraining, versus the average of 3 baselines that trained regularly for the same number of overall steps (the same models used in figure [.. Similarly to the results on Natural Questions (figure) D. all examined models score only slightly better than random guess the examined GL.E tasks. However (and again similarly to the case of Natural Questions), we get a clear signal that ENN-Pretraining significantly 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, can better understand and utilize the various in-context hints that INN-Pretraining provides.
GLUE Task
MNLI
RTE
WNLI
Random guess
33.3
50.0
50.0
3 baselines Average score
35.1
52.0
51.1
baselines Max score
35.3
52.3
54.9
KNN-Pretraining
35.5
53.0
56.3
Table 3: Zero-shot accuracy scores on several GLUE tasks of KNN-Pretrained model versus baselines that trained regularly the same data.
40


Reviews:
[{"review_id": "sUPa1qI4yF", "reviewer": "Reviewer_yRBu", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n1.\tThe authors analyze the in-context bias of the self-attention model, which could inspire some research works on designing training examples.\n2.\tThe authors propose to include related texts retrieved by the kNN method in a single training sample, which is proved effective in solving sentence similarity tasks.\n3.\tA theoretical analysis of the in-context bias.\n\nWeakness:\n1.\tThe introduction of the motivation (the concept of in-context bias) is not easy to understand at the very beginning.  The paper said: “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.”  Acutally it seems quite natural for me and I did not realize it is a problem until I saw more explanations in section 1.1.\n2.\tThe theory is a bit complicated and not easy to follow. \n3.\tThe experiments are limited. The authors only conduct the evaluation on sentence similarity tasks and open domain QA tasks. However, there are many other tasks that involve sentence pairs. For example, sentence inference tasks such as MNLI and RTE are common tasks in NLP field. The authors should conduct experiments on more types of sentence pair tasks.\n", "summary_of_the_paper": "The paper addresses the bias induced by the inconsistency between the pretraining examples and downstream examples. Specifically, the text is continuous during pretraining while could be non-neighboring in downstream tasks. The authors proposed two methods: kNN-TAPT and kNN-pre-training in Task Adaptive Pre-training (TAPT) step or general Pre-training step respectively, where each training sample is composed of a sentence and its semantically closed neighbors (by kNN search). The method is effective in sentence similarity tasks (in kNN-TAMP scenario) and closed book open domain QA tasks (in kNN-Pre-training scenario).  The paper also gives a theoretical analysis of the so-called in-context bias,  by quantifying 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).", "summary_of_the_review": "The paper addresses the inductive bias of in-context learning. The studied problem is meaningful and could be a potential breakthrough for pretrained language model. However, the written of the paper is a bit hard to follow and the experiments is somewhat limited.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VluDZlu8XBI", "reviewer": "Reviewer_ADar", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main results of this paper are of broad interest. Questions about how the nature of the input to LMs affects learning are important and poorly understood. The claim that the within- and between-example is crucial is novel and thought provoking.\n\nThe arguments are highly technical and require a lot of expertise, (including familiarity in particular with Levine, 2020, on which the proof builds substantially) to understand. This is not necessarily a strength or a weakness---there is great value to developing techniques for making this kind of formal argument. However, I find it difficult to follow even the thread of the argument at an intuitive level, which is a shame because it's interesting to me. I suspect I'm not the only potential reader who will have this problem. The paper could reach a much wider audience (which I think would be an improvement) if it provided more intuitive paraphrases of the main results in Section 2.3. \n\nThere are a few unaddressed issues with the authors' main suggestion: that we can improve LMs’ abilities to make cross-example connections by constructing examples out of interleaved texts. Assuming a fixed example length, this means portions of text from the same source that otherwise would have been in the same example in the standard chunking scheme would now be separated. However, it seems that, in general, dependencies within a text are more important than dependencies between random semantically related sentences from different texts. Another problem is that a language model trained on this kind of interleaved text would presumably also generate less coherent texts.\n\nI also have some doubts about the empirical results. The results about kNN-task adaptive pretraining seem to have a confound: The method for retrieving data for this intermediate training step involves finding sentences with high similarity, but the task on which this method is shown to lead to improvements is itself a semantic similarity task. The claim is that the method should be helpful for many kinds of tasks, but my intuition is that this \"coincidental\" alignment of the data collection and the evaluation task makes the result less generalizable.\n\nThe empirical results about zero-shot QA don't help much because the task is clearly too difficult. The improvement on the models trained with the new method is from F1 of ~0.001 to ~0.01. Do larger LMs succeed at this task? If so, then it seems that models tested in this paper are just too small or ineffective for the task. Is there an easier task where this method leads to strong performance but ordinary pretraining does not?\n\n", "summary_of_the_paper": "This paper provides an argument for reconsidering how we construct a single training example during LM (pre)training. It gives a formal proof that dependencies between two texts are vanishingly weak when the texts are presented in separate training examples (but not when they are in the same training example). Intuitively, in order to make cross-textual inferences between texts in different examples, a model must be able to access information stored in its weights, but commonly used learning rates and model architectures ensure that the updates to the weights are too small to recover this information.\n\nThe paper suggests remedying this problem by changing the way in which we construct training examples. Rather than obtaining examples by separating texts into contiguous chunks, we can construct examples from multiple semantically related chunks found by computing the similarity of text embeddings in the training corpus. Empirical results are presented which show that task-adaptive (i.e. post-pretraining) training on such data leads to greatly improved performance on zero-shot semantic similarity. Further results show that doing pretraining with these kinds of examples improves zero-shot QA performance by a small margin (from an extremely low baseline).", "summary_of_the_review": "This paper makes an intriguing claim about the effect of the input on LM learning, which is backed up by a formal argument. These kinds of arguments are valuable. However, the argument is difficult to follow for a non-specialist audience. Also, the empirical results the authors report have some confounds due to questionable task selection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "S5GhA_Db7Fv", "reviewer": "Reviewer_cz9E", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I thought that this paper was very thought-provoking, and I appreciated the attempts to better understand what is going on with pre-trained language models, why they work well, and what might we be able to improve from theses insights.\n\nStrengths:\n- Thorough theoretical analysis that reveals the connection between (practically-necessary) small learning rates and inability to use dependencies across text chunks\n- Useful framing and discussion of the \"in-context bias\", where models are more likely to learn dependencies within text chunks seen during pre-training.\n- reasonable initial experimental results demonstrating some ways to help models better use cross-text-chunk dependencies (put them into a contiguous text chunk), providing some hope that these results could make models better.\n\nWeaknesses:\n- I felt like the empirical validation could have been stronger. The only two tasks examined are sentence similarity tasks (which seem a bit more like a sanity check), and NaturalQuestions. It's particularly surprising to me that this works so well on NQ, and I wish the authors had dug a bit deeper into this, but I also recognize that page limits exist.\n- Where else do you think the in-context bias could be useful? It wasn't intuitive for me that it'd be useful for NQ.\n- Do you have any initial experiments on the \"self-improving\" aspect of this technique? This is mentioned several times, but there are no initial results or anything suggesting that it might be a promising direction to pursue.", "summary_of_the_paper": "This paper formalizes the notion of an in-context bias---the pre-trained neural language models are better at modeling dependencies that appear within one contiguous pre-training text chunk, rather than those that appear in different pre-training text chunks. The authors perform a theoretical analysis of this phenomenon and link the inability to model dependencies between different examples with the models' low-pretraining learning rate. The authors then propose two methods for improving pre-training example design by adding sentences that are similar (KNN), but from different documents. They empirically show that this works better than reasonable baselines on sentence similarity tasks (e.g., STS and SICK), as well as NaturalQuestions.", "summary_of_the_review": "I thought this paper was interesting and brings a new perspective to practical design decisions used in pre-training language models. I think this paper is relevant to the ICLR audience, and that said audience would be excited to know about it.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "O_F_Z7yCpXv", "reviewer": "Reviewer_twvb", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n * Clearly written paper\n * The theoretical contributions will have a high impact on training transformer-based models\n * The theoretical analysis is supplemented by experimental analysis\n\n Weaknesses:\n  * No major weaknesses\n\n  \nSome minor aspects:\n * \"correct on less than 50 questions out of 20, 000 in the evaluation set\" – Separating thousands by commas (American system) is ok. Separating them by space is also ok. Doing both is not.\n * \"correct on roughly than 250 questions in the evaluation set\" – \"than\" shouldn't be there\n * Equation 4 would have been easier to read if the right-hand side was expanded to S1 and S2's `{w_2^j}` notation from the previous line.", "summary_of_the_paper": "The paper presents a theoretical analysis of the strengths of including similar input examples S1 and S2 into the same actual transformer input. The analysis contrasts using examples of concatenated S1;S2 with running the two in separate inputs. The paper then proves empirically that including related examples in the same input (via a method the article calls KNN-pretraining) allows transformers to learn much faster in cases when cross-document dependency is relevant.", "summary_of_the_review": "Despite its theoretical focus, the paper is easy to read and to follow. The theoretical analysis is supplemented by empirical results.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "921tlMcxv4l", "reviewer": "Reviewer_Z6SX", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: 1. Some experimental results are interesting. Simply adding related sentences in the pre training input context helps end performance.\n\nWeaknesses: 1. I think the presentation of the paper needs to be improved. For the theoretical analysis, it was not clear to me what is the contribution of the current analysis compared to the Levine 2020 paper. I felt similar conclusions can be drawn from the results of that paper as well. It will be good to rewrite highlighting the contributions. I also noticed a lot of repeated text in section 2 from Levine 2020 paper; will be good to modify.\n\n2. The empirical part of the paper shows improved performance of adding similar sentences to the context of LM training. I felt this was quite separate from the theoretical analysis. It does not follow from the theoretical results that adding similar sentences will be a good thing. As a result, having these in the same paper looked incoherent to me.\n\n3. The experimental results are weak. The gains are not super high. I will be more convinced if evaluation is done on a wider range of tasks.    ", "summary_of_the_paper": "The authors show empirically and theoretically that transformer based language models cannot effectively exploit dependencies between sentences in different inputs. For the theory part, they introduce a new measure called \\eps-sep rank. A function is said to have low \\eps-sep rank if it can be approximated up to error \\eps by a function with low separation rank. Bounding the measure shows that if learning rate is low, the information from different training instances cannot be well integrated. Empirically, they present experiments where the task specific pre training was altered to include non-task specific similar sentences in the context of the LM. They show that augmenting context of pre training inputs with related similar sentences from Wikipedia helped improve end performance on similarity and zero shot open domain QA tasks. ", "summary_of_the_review": "I was not clear about the impact of the theoretical analysis presented, compared to past work. It will be good to rewrite the section clarifying the contributions and novel conclusions from the analysis. The empirical section sounded disparate from the theoretical part of the paper, and the results are weak. I suggest evaluating on a wider variety of tasks. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a novel framing of what's at stake when selecting/segmenting text for use in language model pretraining. Four reviewers with experience working with these models agreed that the conceptual and theoretical work here is insightful and worth sharing. The empirical work is fairly small-scale and does not yet support broad conclusions, but reviewers did not see such conclusions as necessary for the paper to be valuable.
Decision: Accept (Spotlight)