PaperlayerView exact revision on arXiv

Source-linked AI summary

A Mechanism for Sample-Efficient In-Context Learning for Sparse Retrieval Tasks

Jacob Abernethy, Alekh Agarwal, Teodor V. Marinov, Manfred K. Warmuth

arXiv:2305.17040v1cs.LGcs.CL

TL;DR

In-context learning allows pretrained language models to adapt from a handful of examples without parameter updates, yet its sample efficiency remains difficult to explain. This paper gives transformer mechanisms for segmenting prompts and learning sparse linear hypotheses, and establishes implementability and sample-complexity guarantees. Its empirical findings also validate delimiter dependence and show correspondence between attention maps and the proposed hypothesis-learning steps.

  • Problem

    How pretrained language models achieve sample-efficient in-context learning from a handful of examples, despite the difficulty of explaining its versatility, remains an open mechanistic question.

    Source: §1 · p. 1

  • Method

    The paper proposes transformer mechanisms that segment delimiter-formatted examples, infer a sparse linear regressor, and apply it to a test input.

    Source: §1 · p. 2, §2.2 · p. 4, §3.2 · p. 5

  • Results

    The proposed procedure is implementable with transformers and has sample-complexity guarantees, while experiments validate delimiter dependence and correspondence between attention outputs and hypothesis-learning steps.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The work provides an end-to-end mechanistic account of segmentation and hypothesis learning for the studied in-context learning tasks.

    Source: §7 · p. 9

    The prompt encoding assumes distinct example and label delimiters that do not occur in task inputs or outputs.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

We study the phenomenon of \textit{in-context learning} (ICL) exhibited by large language models, where they can adapt to a new learning task, given a handful of labeled examples, without any explicit parameter optimization. Our goal is to explain how a pre-trained transformer model is able to perform ICL under reasonable assumptions on the pre-training process and the downstream tasks. We posit a mechanism whereby a transformer can achieve the following: (a) receive an i.i.d. sequence of examples which have been converted into a prompt using potentially-ambiguous delimiters, (b) correctly segment the prompt into examples and labels, (c) infer from the data a \textit{sparse linear regressor} hypothesis, and finally (d) apply this hypothesis on the given test example and return a predicted label. We establish that this entire procedure is implementable using the transformer mechanism, and we give sample complexity guarantees for this learning framework. Our empirical findings validate the challenge of segmentation, and we show a correspondence between our posited mechanisms and observed attention maps for step (c).

1 Introduction

In-context learning lets pretrained language models adapt from a handful of labeled examples without updating parameters, but its sample efficiency and versatility remain difficult to explain. This work proposes and studies transformer mechanisms for segmenting prompts, learning hypotheses, and applying them to new inputs.

1 Introduction

  • The paper abstracts pretraining as a fixed prior over sequences and studies an end-to-end account of in-context learning.This contrasts with prior approaches that require pretraining on the downstream task or a mixture of downstream tasks.Source: §1 · p. 2
  • The proposed segmentation mechanism selects example and label delimiters by maximizing the segmentation likelihood under the pretraining prior.Its sample complexity depends logarithmically on the number of candidate delimiters and inversely on the gap between correct and incorrect segmentation likelihoods.Source: §1 · p. 2
  • For sparse retrieval tasks, attention can identify a hypothesis consistent with the prompt examples and carry it forward to answer later queries.The task family includes outputs that are an input token or a sum of a subset of input tokens.Source: §1 · p. 2
  • Empirical validation shows dependence on identifiable delimiters and correspondence between attention outputs and the proposed hypothesis-learning steps.The authors also train transformers to solve the studied sparse retrieval tasks.Source: §1 · p. 2

2 Problem Setting and Notation

The paper frames in-context learning as using a pretrained language model to infer a task from prompted examples and predict a test output. It specifies a decoder-only transformer and an encoding that separates examples and labels with distinct delimiters.

2.1 The Transformer Architecture

  • A decoder-only transformer embeds token sequences, processes them through layers, and produces next-token predictions.The described layer uses attention heads followed by a residual update and GeLU; layers compose to form the model.Source: §2.1 · p. 3

2 Problem Setting and Notation

  • A supervised task is a function paired with an input distribution, and its examples are i.i.d. labeled samples from that task.The learning objective is to estimate a function with low expected loss on a fresh input.Source: §2.2 · p. 3

2.2 In-Context Learning

  • In-context learning uses the prompted examples and a test input as the language-model input, with the desired output equal to the task function’s value on that test input.The encoding separates examples with one delimiter and separates each input from its label with another; the delimiters must be distinct and absent from task examples.Source: §2.2 · pp. 3–4

3 An Overview of Results

The overview presents prompt segmentation and consistent-hypothesis learning as the two stages of its in-context learning account. It describes likelihood-based delimiter selection, followed by transformer-based learning for tokenized sparse regression.

3.1 Segmenting an input sequence

  • A transformer can implement maximum-likelihood selection of label and example delimiters under the learned sequence prior.The segmentation objective treats the separated input and output sequences as individually plausible under the base distribution.Source: §3.1 · p. 4
  • Segmentation sample requirements depend on the likelihood gap between correct and incorrect delimiter choices and a minimum-probability parameter.The informal theorem gives a high-probability guarantee for recovering both separators; choosing unambiguous separators makes the gap larger and the sample cost smaller.Source: §3.1 · p. 4

3.2 Learning a consistent hypothesis

  • Tokenized sparse regression represents outputs as sparse linear combinations of inputs under a fixed basis, with coefficients varying across tasks.The paper focuses on the standard basis and motivates the task as an abstraction of input-token extraction and manipulation.Source: §3.2 · p. 5
  • A transformer can implement the consistent-hypothesis estimator, which weights candidate coordinates across examples and selects the highest-weight candidate.Under the stated near-orthogonal embedding condition for sufficiently distinct tokens, n = Ω(s log(m/ϵ)/ϵ) examples guarantee expected prediction error at most 2ϵ.Source: §3.2 · p. 5

4 Segmenting an ICL Instance

The section presents prompt segmentation as selecting the delimiter pair that makes the parsed examples most likely. Under a likelihood-gap assumption, maximum-likelihood segmentation recovers the correct delimiters with high probability.

4 Segmenting an ICL Instance

  • The procedure splits a sequence at candidate example separators, then uses label separators to divide each chunk into an input and label.A candidate pair is infeasible if a chunk contains multiple label separators; chunks without one are treated as inputs only.Source: §4 · pp. 5–6
  • In the Shakespeare example, (/ , ;) is the correct delimiter pair, while (: , -) creates a less natural alternative parsing.The quotes are inputs and play names are labels; other potential delimiters create ambiguity.Source: §4 · pp. 6–7, Figure 1
  • The segmentation objective selects the delimiter pair maximizing the sequence likelihood, including likelihoods for segments that cross delimiter boundaries.Each candidate pair defines a probability model over the parsed sequence.Source: §4 · p. 6
  • Under Assumption 1's likelihood separation and the theorem's probability lower bound, maximum-likelihood segmentation returns the correct delimiters with probability at least 1 − δ when the sample-size condition holds.The guarantee depends on correct parses being more likely on average than alternatives, and on a lower bound ν for alternative-model probabilities.Source: §4 · p. 7

5 Learning a Consistent Hypothesis for Tokenized Sparse Regression

The section constructs attention mechanisms that infer a consistent sparse hypothesis from tokenized examples and apply it to a query. For the s-sparse task, iterative deflation yields a sample-complexity guarantee under an embedding assumption.

5 Learning a Consistent Hypothesis for Tokenized Sparse Regression

  • For s = 1, attention compares each input coordinate with its example's label, identifying coordinates whose values are consistent with that label.The embedding assumption makes inner products large for matching coordinates and small for values separated by more than ϵ.Source: §5 · p. 7
  • A second attention layer aggregates hypotheses across preceding examples, and later layers copy the resulting hypothesis to predict on the next input.The prediction reads the selected coordinate of the next input and outputs its value at the prediction token.Source: §5 · pp. 7–8
  • The tokenized setting can represent tasks where each input is instead supplied directly as a vector in R^m, by reducing that setting to the tokenized one.The reduction is described in Appendix F.Source: §5 · p. 8
  • For s-sparse regression, the construction identifies coordinates one at a time and deflates the target after each identification to avoid inconsistent combinations.The procedure stacks O(s) layers of the 1-sparse construction with a deflation layer after each coordinate.Source: §5 · p. 8
  • For any ϵ > 0, Theorem 6 gives expected prediction error at most 2ϵ with probability 1 − δ when n = Ω(s log(m/ϵ)/ϵ), under its embedding and Assumption 2 conditions.The stated embedding dimension is O(s/ϵ).Source: §5 · p. 8

6 Empirical Results

In the experiments, the model reaches zero loss after one Gaussian example, while Rademacher inputs can require more examples to disambiguate the target coordinate. Its attention patterns broadly correspond to the proposed mechanism.

6 Empirical Results

  • The reported loss and attention plots average results over a mini-batch of 64 sequences while fixing the same f⋆ across that batch.The section reports these experiments for the 1-sparse tokenized regression task.Source: §6 · p. 8
  • Zero loss is reached after one Gaussian sample; in the Rademacher setting, loss is higher while several coordinates remain consistent and the model typically learns f⋆ once the correct coordinate is disambiguated.The single-sample Gaussian result is information theoretically optimal.Source: §6 · pp. 8–9, Figure 3 · p. 9
  • For Gaussian inputs, prediction-token attention peaks at f⋆ from the second example; for Rademacher inputs, it becomes more peaked as examples resolve early ambiguity.The Rademacher attention is more variable during the first few examples, when multiple hypotheses can remain consistent.Source: §6 · p. 9, Figure 3 · p. 9
  • Layer 6 attention peaks at previous label tokens, matching the proposed step of aggregating a consistent hypothesis across examples.The authors report that the observed attention maps bear correspondence to their theory.Source: §6 · p. 9, Figure 3 · p. 9

7 Conclusion

The paper presents transformer mechanisms for prompt segmentation and hypothesis learning, with statistical guarantees that relatively few examples can suffice. It situates these contributions against prior ICL work and specifies notation for token sequences, examples, delimiters, and attention.

7 Conclusion

  • Transformers can implement sequence segmentation and hypothesis learning for a family of ICL tasks, with statistical guarantees that fairly few examples can suffice.The conclusion identifies information-theoretically optimal learning in the task types studied here and in prior works [10] [18] [1] as an open direction for further study.Source: §7 · pp. 9–10

A Related Work

  • Prior ICL analyses leave gaps in sample-complexity guarantees, prompt parsing, and explanations of how self-attention implements learning mechanisms.The authors note that existing guarantees rely on specific pre-training distributions or difficult-to-quantify stability, while most prior work assumes prompts are already parsed.Source: A Related Work · p. 12

B Notation

  • The notation represents prompts as token sequences containing i.i.d. task examples and labels, separated by candidate delimiter pairs.It also establishes conventions for token embeddings and transformer attention notation.Source: B Notation · pp. 12–13

C Proofs of the segmentation results

The segmentation analysis compares the prior likelihood of a correctly segmented prompt with alternatives, then implements delimiter selection in a transformer. The construction uses attention and conditional-probability outputs to score candidate segmentations.

C Proofs of the segmentation results

  • The correct delimiter pair is likely to assign greater prior likelihood to the i.i.d. task sequence than any alternative delimiter pair.The proof analyzes the log likelihood ratio and uses concentration plus a union bound to control the probability that any alternative scores higher.Source: C Proofs of the segmentation results · p. 14
  • The proof turns the log-likelihood comparison into a martingale concentration bound, with sample size chosen to control errors across candidate delimiters.A union bound yields probability less than δ that any incorrect delimiter pair has greater likelihood than the correct pair.Source: C Proofs of the segmentation results · p. 14

C.1 Segmenting example and label delimiters

  • For each candidate delimiter pair, attention layers locate separators and check that the proposed example and label positions are consistent.The construction then uses a conditional-probability transformer to score the segmented sequence under the pre-training distribution, and selects the maximizing delimiter pair.Source: §C.1 · pp. 14–16

D Learning the the 1-sparse tokenized regression task

For 1-sparse tokenized regression, the transformer mechanism aggregates evidence across prompt examples to select a consistent coordinate, then applies that hypothesis to a query.

D Learning the the 1-sparse tokenized regression task

D.1 Transformer mechanism for 1-sparse tokenized regression

  • Within each example, label-token attention weights coordinate candidates, and the resulting value vectors accumulate evidence for coordinates consistent with the labels.The mechanism selects the coordinate with the largest cumulative weight across examples.Source: §D.1 · pp. 17–18
  • The first two layers identify an approximately correct hypothesis by aggregating candidate-coordinate evidence across examples and selecting a maximally supported coordinate.The text states that any hypothesis consistent with all examples has small prediction error on future examples in this task.Source: §D.1 · p. 18
  • Inference uses attention to carry the selected coordinate to the next example and produce its prediction.The construction marks coordinates consistent with the learned hypothesis and copies the corresponding query token to the output position.Source: §D.1 · p. 19

D.2 Sample complexity

The sample-complexity analysis bounds the prediction error of the hypothesis returned by ICL for 1-sparse tokenized regression. With a suitable embedding and enough examples, its expected error is at most twice the target tolerance.

D.2 Sample complexity

  • The returned hypothesis has error at most ϵ on every example in context, which the analysis uses to derive a risk bound.The argument assumes labels are generated by the target hypothesis and applies a finite-hypothesis generalization bound.Source: §D.2 · p. 19
  • n = Ω(log(m/ϵ)/ϵ) examples suffice for E[|fn(x) − f⋆(x)|] ≤ 2ϵ with an embedding in R^O(1/ϵ).Theorem 7 gives this guarantee for any ϵ > 0.Source: §D.2 · p. 20

E s-sparse Tokenized Regression

For the s-sparse tokenized task, the construction iteratively identifies consistent coordinates, deflating already-selected ones to avoid requiring a unique consistent hypothesis. After O(s) layers it obtains a one-to-one approximate match to the true coordinates, then applies the learned hypothesis for inference.

E s-sparse Tokenized Regression

  • Iterative deflation learns one hypothesis consistent across examples, avoiding the direct coordinate-wise approach unless a unique consistent hypothesis exists with high probability.The construction uses O(m) layers with one head per layer.Source: E s-sparse Tokenized Regression · p. 20
  • The coordinate-search construction repeatedly applies the 1-sparse procedure to find a new consistent coordinate after removing those already identified.The resulting consistent coordinate is added to the set of coordinates tracked by the construction.Source: E s-sparse Tokenized Regression · pp. 21–22
  • After L = O(s) layers, the selected set has cardinality s and is in bijection with the true coordinates, with matched values within ϵ on every observed example.The construction also rules out adding a coordinate that is not similar to some true coordinate across the observed examples.Source: E s-sparse Tokenized Regression · pp. 22–23
  • The learned consistent hypothesis is prepared for inference and applied using the final three layers of the index-token task construction.After the coordinate set is converted to a consistent hypothesis representation, those layers perform inference.Source: E s-sparse Tokenized Regression · p. 23

E.1 Proof of Theorem 6

The proof transfers the coordinate-matching guarantee to an inference error bound and gives a sample-complexity bound analogous to the index-token task.

E.1 Proof of Theorem 6

  • The bijection between the learned and true coordinate sets matches each true coordinate to a learned coordinate within ϵ on every example.The lemma states a probability guarantee for the resulting coordinate match.Source: §E.1 · p. 23
  • The approach gives a sample-complexity bound analogous to the one for the index-token task.The passage does not state the bound's explicit expression.Source: §E.1 · p. 23
  • E[|x_j − x_b_n(x_j)|] ≤ 2ϵ for every j ∈ f⋆ under the bijection guaranteed by Lemma 10.The proof says this yields the theorem's result for the chosen n.Source: §E.1 · p. 23

F Vector 1-sparse regression task

For vector 1-sparse regression, the construction recovers a consistent hypothesis bit by bit with attention, then uses the established index-token construction for inference. The same approach has a similar sample-complexity bound and can extend to the s-sparse index-vector task.

F Vector 1-sparse regression task

  • The vector construction learns each bit of a consistent hypothesis sequentially using O(log(m)) attention layers.It builds on the first attention layer's recovery of a consistent hypothesis for an example in the 1-sparse token task.Source: F Vector 1-sparse regression task · pp. 24–25
  • After the bitwise construction, the indicator vector marks exactly the hypotheses consistent with the true hypothesis on the example.The construction replicates the bit-recovery procedure across the hypothesis bits.Source: F Vector 1-sparse regression task · p. 25
  • A consistent hypothesis can then be learned across observed examples and applied for inference using the index-token construction.The sample-complexity bound is similar to Theorem 7, and the construction can also extend to the s-sparse index-vector task.Source: F Vector 1-sparse regression task · p. 25

G.1 1-sparse tokenized regression experiments

The experiments examine 1-sparse tokenized regression under Gaussian and Rademacher inputs, including how pretraining and delimiter choices affect inference. Rademacher-only pretraining transfers to Gaussian inference, whereas Gaussian-only pretraining performs poorly on Rademacher inference.

G.1 1-sparse tokenized regression experiments

  • The setup uses m=5, eight single-head layers, 128-dimensional embeddings, and 8,000 epochs of mini-batch-64 training.Source: §G.1 · p. 26
  • Gaussian inputs permit identifying the hypothesis from one example, whereas Rademacher inputs require Ω(log m) examples for identification with high probability.Source: §G.1 · p. 26
  • Rademacher-only pretraining matches the mixed model on Gaussian inference, while Gaussian-only pretraining performs poorly on Rademacher inference.Source: Figure 6 · p. 27, §G.1 · p. 27, Figure 5 · p. 27
  • In Rademacher inference, attention spreads across consistent hypotheses and peaks on f⋆ at the earliest example where identification is possible.Source: §G.1 · pp. 27–28

G.2 Segmentation

Related papers

Loading 2305.17040v1…