PaperlayerView exact revision on arXiv

Source-linked AI summary

The Inductive Bias of In-Context Learning: Rethinking Pretraining Example Design

Yoav Levine, Noam Wies, Daniel Jannai, Dan Navon, Yedid Hoshen, Amnon Shashua

arXiv:2110.04541v3cs.CLcs.LG

TL;DR

Standard corpus chunking may make dependencies between separately trained text segments harder for NLMs to model, despite their importance for broader NLU capabilities. The paper analyzes training-set expressivity with separation-rank bounds and proposes kNN-based pretraining that places semantically related non-neighboring sentences in the same example. kNN-Pretraining improves zero-shot closed-book QA, producing 100s of correct answers after roughly 10% of overall training time versus much less in randomly initialized baselines.

  • Problem

    Standard corpus chunking may make dependencies between separately trained text segments harder for NLMs to model, despite their importance for broader NLU capabilities.

  • Method

    The paper analyzes training-set expressivity with separation-rank bounds and proposes kNN-based pretraining that places semantically related non-neighboring sentences in the same example.

  • Results

    kNN-Pretraining improves zero-shot closed-book QA, producing 100s of correct answers after roughly 10% of overall training time versus much less in randomly initialized baselines.

  • Takeaways & Limitations

    Pretraining example design is a theoretically motivated degree of freedom for improving representations using the same corpus and objective.

  • Takeaways & Limitations

    The proven sequential-case upper bound is not guaranteed to be tight, and tightening the depth-deficit estimate is left for future work.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper formalizes an in-context bias: NLMs model stronger dependencies between sentences shown in the same training example than between sentences separated across examples. It uses this bias to motivate pretraining-example redesign, including kNN-based additions that improve NLU performance.

  • 1 INTRODUCTION: The analysis extends separation-rank expressivity questions from single inputs to dependencies between sentences encountered during training.It introduces ε-separation rank and related refinements to quantify how training-example construction affects accessible cross-corpus information.
  • 1 INTRODUCTION: For common learning rates η ∈[10^-6, 10^-4], separating sentences across examples incurs a guaranteed depth deficit of ∼6 layers.The bound is ˜O(dx[L −0.5 log3(η^-1)]) for unseen sentence pairs versus ˜O(dxL) for pairs shown together; the sequential bound may be loose.
  • 1 INTRODUCTION: A 10% addition of kNN-Pretraining boosts zero-shot closed-book question-answering score by ∼5X on 20,000 evaluation examples.The paper also reports gains over baselines on several NLU tasks, including zero-shot GLUE evaluations.
  • 1 INTRODUCTION: The in-context bias means information within pretraining examples is better represented than information integrated across separate examples.The paper frames this as a network-expressivity property of the training process, not merely of a single input.

2 THEORETICAL ANALYSIS: THE IN-CONTEXT BIAS OF SELF-ATTENTION

The analysis compares how self-attention NLMs integrate sentences shown together in context with sentences encountered in successive training examples. It formalizes an expressivity gap arising because cross-example information reaches later inputs through gradient-updated weights and a small learning rate.

  • 2.1 THE ANALYZED IN-CONTEXT AND SEQUENTIAL REPRESENTATIONS: The analysis treats corpus segmentation as a training-process hyperparameter and compares in-context and sequential representations using separation-rank-based measures.The framework refines separation rank with an effective ε-separation rank and explicitly tracks where each sentence enters the computation.
  • 2.3 THE EXPRESSIVE ADVANTAGE OF IN-CONTEXT LEARNING: In-context representations can model more elaborate dependencies between two sentences than sequential representations formed across successive training steps.The sequential setup inserts S1, updates the network using its loss, and then inserts S2; the in-context setup presents both sentences together.
  • 2.3 THE EXPRESSIVE ADVANTAGE OF IN-CONTEXT LEARNING: The in-context sequential separation rank asymptotically matches the regular separation rank, validating the measure for comparing the two representations.The in-context upper bound is shown to be tight in effective rank, whereas the sequential upper bound is not guaranteed to be tight.
  • 2.3 THE EXPRESSIVE ADVANTAGE OF IN-CONTEXT LEARNING: When sentences are separated across training examples, S1 affects S2 only through gradient-updated parameters, so its contributions are multiplied by powers of the learning rate and many become negligible.The proof sketch attributes the weaker effective rank to the small magnitude of η, even though both sentences ultimately enter the self-attention stack through inputs.
  • 2.3 THE EXPRESSIVE ADVANTAGE OF IN-CONTEXT LEARNING: For common η ∈ [10^-6, 10^-4], dependencies between sentences never shown together incur a guaranteed depth deficit of approximately 6 layers.The sequential upper bound is ˜O(dx[L − 0.5 log3(η^-1)]), versus ˜O(dxL) for sentences shown in the same input.

3 KNN BASED PRETRAINING EXAMPLE DESIGN

The paper designs kNN-based training examples that place semantically related sentences in context, improving sentence representations and zero-shot question answering. Controlled variants indicate that in-context arrangement, rather than merely sharing a batch, is important.

  • 3.1 KNN TASK ADAPTIVE PRETRAINING: kNN-TAPT concatenates each task example with Wikipedia nearest neighbors selected by embedding similarity, using up to 256 tokens per example.Neighbors exceeded 0.8 cosine similarity, and training continued for five epochs on 170K constructed examples.
  • 3.1 KNN TASK ADAPTIVE PRETRAINING: Over 1 point average improvement over regular TAPT shows that kNN-TAPT benefits SentEval sentence similarity representations.The comparison uses zero-shot SentEval scores from average word embeddings of inserted sentences.
  • 3.1 KNN TASK ADAPTIVE PRETRAINING: The in-context variants differ from in-batch variants by placing related SentEval and Wikipedia sentences in the same training example rather than merely the same batch.All kNN-TAPT variants use the same task and neighbor data; their primary difference is data arrangement.
  • 3.1 KNN TASK ADAPTIVE PRETRAINING: The weakness of neighbor in-batch training supports the theoretically motivated advantage of placing related sentences in context.The paper suggests separate Wikipedia-only examples may dilute the original TAPT signal, whereas random in-context examples avoid that dilution.
  • 3.2 KNN PRETRAINING: kNN-Pretraining extends the approach to general pretraining by adding related non-neighboring sentences while retaining regular pretraining examples.The method was applied to GPT-base and GPT-medium models pretrained from scratch, with half the batch containing regular examples.
  • 3.2 KNN PRETRAINING: On zero-shot closed-book Natural Questions, kNN-Pretraining yields a low-scoring but significant improvement, with hundreds of correct answers after roughly 10% of training.The paper also reports clear gains over baselines on several GLUE tasks.

4 DISCUSSION

The discussion presents training-example design as a parallel route to improving NLM pretraining: redistribute existing data so important corpus connections occur within context. It frames the theory as an explanation for related training and inference methods and as motivation for self-improving representations.

  • 4 DISCUSSION: The paper proposes improving pretraining with the same data and objective by rearranging examples so dependencies between important non-neighboring text chunks are learned more effectively.The central bias favors dependencies between chunks appearing within the same training example.
  • 4 DISCUSSION: The theory links cross-example connectivity to the expressivity of a shallower and weaker architecture than one receiving the same examples in context.This formalizes the connection between training-example construction and available architectural expressivity.
  • 4 DISCUSSION: The quantified trade-off is intended to inform training-example design and motivate concatenating multiple mentions of an entity, event, or concept within one example.The discussion connects this recommendation to empirical methods that already leverage related text in the input.
  • 4 DISCUSSION: The paper extends its proposal toward self-improving representations by repeating kNN search after representation refinement, a direction left for future work.A new representation can define more informative neighbors for a subsequent pretraining cycle.
  • 4 DISCUSSION: The framework is presented as a theoretical interpretation of in-context few-shot prompts and in-context augmentations, alongside its focus on pretraining.The discussion cites Brown et al. (2020), Gao et al. (2020), and Schick & Schütze (2020) as examples of related practices.

A PROOF OF COROLLARY 1

The proof expresses the sequential in-context representation as a sum of terms separable in the two inputs, then selects a subsum that approximates it within ε.

  • Assumptions and bounds: The proof bounds the relevant terms using assumptions on parameter and gradient magnitudes and bounds on products of network matrices.These assumptions provide coordinate-wise controls used in the subsequent estimates.
  • Proof strategy: The proof rewrites the sequential representation as a sum of terms separable with respect to (a, b).The representation is expanded by separating vocabulary-gradient, vocabulary, weight, and variable contributions.
  • Approximation construction: An index subset G is chosen so that the sum of its terms is an ε-approximation of the target representation.The construction reduces the approximation requirement to an upper bound on the discarded terms.
  • Conclusion: The resulting choice of subset, based on a sensitivity parameter and the supporting lemmas, yields the required ε-approximation.The final construction uses s∗ and establishes that ZG(T(s∗)) approximates the sequential representation.

B.1 LEMMAS FOR ESTIMATING THE NUMBER OF COEFFICIENTS

This section develops combinatorial lemmas for estimating how many coefficient configurations satisfy sensitivity conditions around a maximal multinomial configuration.

  • Maximizing configurations: The maximal multinomial coefficient occurs when the summands are distributed as evenly as possible across indices.This property identifies the configuration around which sensitivity neighborhoods are estimated.
  • Analytic bounds: Stirling approximations, Taylor expansions, and unimodality arguments provide tractable bounds for the relevant coefficient counts.The resulting estimates support both upper and lower cardinality bounds.
  • Coefficient counting: Lemma 6 bounds the number of multinomial coefficients satisfying a sensitivity threshold from above and below.These bounds are used to estimate the cardinality of the coefficient sets needed later.
  • Geometric reduction: The coefficient-counting problem is reduced to counting integer tuples and lattice points in a centered multidimensional ball.The reduction imposes a zero-sum constraint on the tuple coordinates.

C.2 PROOF OF THE LOWER BOUND

The lower-bound proof transfers high ε-rank of a grid-tensor matricization into high ε-sequential-separation rank by constructing suitable templates and network weights.

  • Rank transfer: High ε-rank of a grid-tensor matricization implies high ε-sequential-separation rank for the analyzed Transformer function.The argument applies a general rank relation to ε-approximations arising from the separation-rank definition.
  • Matrix construction: The proof constructs a symmetric submatrix with unit diagonal and controls its Frobenius norm to invoke the relevant rank lemmas.The construction uses l2-normalized rows of an auxiliary matrix B.
  • Network realization: Template vectors and first-layer key and query weights are chosen so self-attention sums the input embeddings.The assignment implements this summation through linearity and the specified embedding and attention weights.
  • Existence argument: A probabilistic-method argument establishes the existence of a matrix B meeting the required norm bound.The proof bounds the expectation of the squared norm and concludes that at least one such B exists.

C.3 TECHNICAL LEMMAS

The technical lemmas bound expectations involving spherical coordinates and establish the analytic inequalities used in the lower-bound construction.

  • Spherical reduction: Rotational invariance reduces expectations over the sphere to expectations involving a fixed coordinate direction.A rotation maps an arbitrary unit vector to the first basis vector without changing the relevant measure.
  • Expectation bounds: The remaining expectation is evaluated directly and bounded using preceding lemmas and auxiliary facts.The proof combines these bounds term by term inside the integral.
  • Pointwise inequality: Lemma 17 supplies the pointwise inequality needed for the integral estimate over x ∈ [0, 1].The proof reduces the inequality to critical points because the boundary expression vanishes.
  • Critical-point analysis: The critical point of the relevant expression is x2 = 2λ/(2λ+d), which determines the location used in the bound.This value is identified as the only critical point.

D.1 KNN-TAPT

The KNN-TAPT training used continued RoBERTa pretraining with masked language modeling under a specified AdamW configuration.

  • RoBERTa-base weights were continued on the masked-language-modeling task with a 15% masking probability.
  • Training used AdamW with β1 = 0.9, β2 = 0.98, ε = 10−6, and weight decay 0.01.
  • The setup used batches of 128 or 256 and sequences of 256 tokens, depending on model size.

D.2 KNN-PRETRAINING

The KNN-Pretraining setup trained GPT-2 from scratch on Wikipedia before introducing generated kNN data at selected training points.

  • A HuggingFace Transformers implementation of GPT-2 was pretrained from scratch on Wikipedia with the standard language-modeling objective.
  • Training used AdamW with β1 = 0.9, β2 = 0.95, ε = 10−8, and weight decay 0.1.
  • The setup used batch size 512 and sequences of 256 tokens, then switched between standard and generated kNN data at two training points.

E KNN-PRETRAINING AT DIFFERENT CHECKPOINTS

This section examines KNN-Pretraining across model sizes and regular-pretraining checkpoints using Natural Questions F1 scores. The reported pattern is that further pretraining appears to improve its effectiveness.

  • Further pretraining seems to improve the effectiveness of KNN-Pretraining on Natural Questions.
  • The evaluation uses zero-shot closed-book Natural Questions examples.
  • Table 2 evaluates Natural Questions F1 scores across different model sizes and regular-pretraining steps.

F KNN-PRETRAINING ON ADDITIONAL BENCHMARKS

The paper tests KNN-Pretraining on several zero-shot GLUE classification tasks in addition to Natural Questions. Although scores remain near random, the method produces a clear positive signal after limited additional pretraining.

  • The evaluated tasks are MNLI, RTE, and WNLI, formatted for unidirectional language models and assessed zero-shot.
  • KNN-Pretraining significantly moves the needle on several zero-shot GLUE tasks after only 10% of regular pretraining time.
  • The examined unidirectional models score only slightly better than random guess on these GLUE tasks.
Loading 2110.04541v3…