PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions

Satwik Bhattamishra, Arkil Patel, Phil Blunsom, Varun Kanade

arXiv:2310.03016v1cs.LGcs.CL

TL;DR

The paper examines limits of in-context learning, whether attention is necessary, and whether pretrained LLMs learn from examples rather than only recalling pretraining tasks. It tests Boolean-function tasks, teaching sequences, and pretrained models, finding varied Transformer performance, sample-efficient teaching-sequence learning, and LLM performance competitive with nearest-neighbor classification.

  • Problem

    It remains unclear how well Transformers implement learning algorithms across task types, whether attention is essential, and whether pretrained LLMs can learn from in-context examples.

    Source: §5 · p. 7

  • Method

    The paper empirically tests sequence models on sampled Boolean-function learning problems, teaching sequences, and pretrained-model prediction tasks.

    Source: §2 · p. 3, §4 · p. 5, §5 · p. 7

  • Results

    Transformers perform well on simpler Boolean classes but degrade on complex ones, learn more sample-efficiently from teaching sequences, and pretrained LLMs compete with nearest-neighbor classification.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show that some attention-free architectures can approach Transformer performance and that pretrained models can use in-context examples on tasks drawn from a large combinatorial space.

    Source: §1 · p. 2, §7 · p. 9

    The example threshold for in-context learning can depend on the function class, with more complex classes potentially requiring more examples.

    Source: §3 · p. 3

Abstract

from arXiv · show

In order to understand the in-context learning phenomenon, recent works have adopted a stylized experimental framework and demonstrated that Transformers can learn gradient-based learning algorithms for various classes of real-valued functions. However, the limitations of Transformers in implementing learning algorithms, and their ability to learn other forms of algorithms are not well understood. Additionally, the degree to which these capabilities are confined to attention-based models is unclear. Furthermore, it remains to be seen whether the insights derived from these stylized settings can be extrapolated to pretrained Large Language Models (LLMs). In this work, we take a step towards answering these questions by demonstrating the following: (a) On a test-bed with a variety of Boolean function classes, we find that Transformers can nearly match the optimal learning algorithm for 'simpler' tasks, while their performance deteriorates on more 'complex' tasks. Additionally, we find that certain attention-free models perform (almost) identically to Transformers on a range of tasks. (b) When provided a teaching sequence, i.e. a set of examples that uniquely identifies a function in a class, we show that Transformers learn more sample-efficiently. Interestingly, our results show that Transformers can learn to implement two distinct algorithms to solve a single task, and can adaptively select the more sample-efficient algorithm depending on the sequence of in-context examples. (c) Lastly, we show that extant LLMs, e.g. LLaMA-2, GPT-4, can compete with nearest-neighbor baselines on prediction tasks that are guaranteed to not be in their training set.

1 INTRODUCTION

The paper asks how far Transformers can learn in-context algorithms, whether attention is essential, and whether pretrained LLMs can learn from examples. It studies Boolean-function tasks, informative teaching sequences, and pretrained models to investigate these questions.

1 INTRODUCTION

  • The study builds on simplified, well-defined settings because real-world language data and tasks are difficult to define precisely.Prior work used such settings to examine in-context learning, including models trained from scratch on prompts sampled across tasks.Source: §1 · p. 1
  • The paper frames four questions: Transformer limits, the necessity of attention, the value of informative examples, and pretrained LLMs’ ability to implement learning algorithms.Boolean inputs make the framework usable for testing LLMs on discrete tasks.Source: §1 · pp. 1–2
  • Transformers approach optimal learning on simpler Boolean tasks but struggle on more complex classes, including parity, where they remain at chance despite algorithms succeeding with the same data size.Attention-free models nearly match Transformers on most tasks, though gaps remain on some.Source: §1 · p. 2
  • Teaching sequences let Transformers learn more sample-efficiently, and one Transformer can switch between efficient and less efficient algorithms depending on whether the sequence is a teaching sequence.A teaching sequence uniquely identifies its function within the class.Source: §1 · p. 2
  • Pretrained LLMs achieve non-trivial performance, competing with nearest-neighbor classification on tasks sampled from a large combinatorial space.The authors argue these tasks are virtually guaranteed to be absent from pretraining, so the models learn from in-context examples alone.Source: §1 · p. 2

2 SETUP FOR IN-CONTEXT LEARNING

The setup trains sequence models across sampled learning problems so that they learn an algorithm for predicting labels from in-context examples. It differs from ordinary supervised learning, which learns one target function from a sample.

2 SETUP FOR IN-CONTEXT LEARNING

  • Each training sequence is a learning problem: the model sees labeled examples and predicts labels autoregressively from each prefix.The setup treats learning as learning a learning algorithm, unlike traditional training on one target function.Source: §2 · p. 3
  • Training sequences are generated by sampling inputs from a distribution and a function from a function-class distribution, then labeling each input with that function.Each of the N sequences uses fresh samples of inputs and a function.Source: §2 · p. 3
  • The study primarily uses Boolean functions mapping fixed-length bit vectors to binary labels, while varying the function class across learning problems.The model predicts each next label from the preceding examples and the current query input.Source: §2 · p. 3
  • The experiments focus on Transformers but also compare recurrent, state-space, long-convolution, and hybrid sequence architectures trained from scratch.The models use causal masking and are tuned across depths, widths, and learning rates.Source: §2 · p. 3

3 IN-CONTEXT LEARNING BOOLEAN FUNCTIONS

The experiments test several Boolean-function classes with varying characteristics and compare sequence architectures on their ability to predict labels in context. Transformers do well on some classes but degrade on others, while attention-free models are competitive on most tasks.

3 IN-CONTEXT LEARNING BOOLEAN FUNCTIONS

  • The task suite includes conjunctions and disjunctions, 3-term DNFs and 3-clause CNFs, and parity and sparse-parity functions.The Boolean inputs are sampled mostly uniformly, with modified distributions for conjunctions and 3-term DNF to avoid near-one null accuracy.Source: §3 · pp. 3–4

3.1 RESULTS

  • Transformers attain near-perfect accuracy on some classes, including simple conjunction and disjunction tasks, but performance degrades across other classes and remains at chance on parity tasks.The authors report parity failures despite other algorithms succeeding with the same number of examples, so insufficient in-context data does not explain the failures.Source: §3.1 · p. 4, Table 1 · p. 5
  • Table 1 compares architectures by accuracy on the last input in each prompt, averaged across 1280 test sequences.The reported table values have standard deviation within 0.5.Source: §3.1 · p. 4, Table 1 · p. 5
  • Attention-free architectures perform almost as well as Transformers on most tasks, although RetNet and DSS trail on 0-1 Threshold functions and Hyena differs on nearest neighbor.Hyena is the closest to Transformers overall; recurrent models also struggle to match Transformer performance on linear regression tasks.Source: §3.1 · pp. 4–5

4 IN-CONTEXT LEARNING WITH TEACHING SEQUENCES

Transformers achieve perfect accuracy on subsequent examples after receiving teaching sequences across five Boolean tasks, learning more sample-efficiently than from arbitrary examples. For Conjunctions, a Transformer trained on both sequence types can select between distinct algorithms according to the prompt.

4 IN-CONTEXT LEARNING WITH TEACHING SEQUENCES

  • The experiments place the shortest teaching sequence first, followed by examples sampled from the input distribution, across five Boolean-function tasks.The tasks are Conjunctions, Disjunctions, CNFs, DNFs, and sparse parities; sequence lengths depend on each function’s literals or relevant bits.Source: §4 · pp. 5–6
  • Transformers achieve perfect accuracy on subsequent examples across all five tasks after receiving a teaching sequence.They also succeed on DNFs and sparse parities, where they struggled in the vanilla setting.Source: §4 · p. 6, Figure 2 · p. 6
  • Teaching sequences let Transformers predict with fewer examples, while FFNs require additional examples beyond the sequence to reach high accuracy.A teaching sequence is the smallest sequence required to learn the target function.Source: §4 · p. 6
  • Vanilla-trained Transformers can use teaching-sequence prompts, but need more examples; teaching-sequence-trained models perform poorly without those initial examples.The latter algorithm relies on the first t examples containing the teaching sequence.Source: §4 · p. 6, Figure 3 · p. 7
  • A Transformer trained on mixed Conjunction and Teach Conjunction examples achieves near-optimal performance on both, selecting an algorithm based on the prompt.It uses the sample-efficient approach for teaching sequences and the less sample-efficient approach otherwise.Source: §4 · p. 6, Figure 3 · p. 7

5 INVESTIGATIONS WITH PRETRAINED MODELS

The paper tests whether pretrained LLMs can apply learning algorithms to Boolean-function prompts, using frozen GPT-2 and direct evaluation of several LLMs. Directly evaluated models match or exceed nearest-neighbor baselines in dimensions up to seven, while frozen GPT-2 can implement nearest-neighbor prediction but remains below fully trainable Transformers.

5 INVESTIGATIONS WITH PRETRAINED MODELS

  • The study uses Boolean functions to evaluate LLMs on discrete inputs and randomly selected target functions from an exponentially large set, reducing pretraining-memorization confounds.The authors note that this makes it possible to analyze learning algorithms with LLMs, whose inputs and outputs are discrete.Source: §5 · p. 7
  • The frozen-GPT setup trains input and output layers around a frozen GPT-2; direct evaluation uses unmodified LLMs, original token embeddings, and prompts with labeled examples.The direct-evaluation models are GPT-4, GPT-3.5-Turbo, and LLaMA-2-70B.Source: §5 · p. 7

5.1 RESULTS: LLMS PERFORM AT LEAST AS WELL AS NEAREST NEIGHBORS

  • Frozen GPT-2 performs better than nearest neighbors on Conjunctions, but remains well below fully trainable Transformers’ near-perfect accuracy.Its Disjunction performance is described as almost identical to its Conjunction performance.Source: §5.1 · p. 8, Figure 4 · p. 8
  • With frozen weights, GPT-2 achieves near-perfect accuracy on nearest-neighbor-labeled examples, and some attention heads attend to the nearest earlier input.The prompt uses 20 randomly labeled points followed by 80 points labeled by the nearest-neighbor algorithm.Source: §5.1 · p. 8, Figure 4 · p. 8
  • LLMs perform as well as or better than nearest-neighbor baselines on Conjunction and Majority tasks up to seven dimensions; GPT-4 slightly exceeds the baseline at 15 dimensions.At inference, the LLM is not told the task.Source: §5.1 · p. 8

6 RELATED WORK

Prior work has studied in-context learning in stylized settings because real-world data and tasks are difficult to define precisely. This work extends that line of inquiry to sequence architectures and pretrained language models.

6 RELATED WORK

  • Stylized, well-defined settings offer a way to study in-context learning where real-world data and tasks are hard to specify precisely.Earlier work established the phenomenon and investigated practical capabilities and limitations of LLMs.Source: §6 · p. 8
  • Figure 5 reports direct inference-time LLM performance across dimensions for Conjunction and Majority tasks.The top row shows Conjunction results and the bottom row shows Majority results.Source: Figure 5 · p. 9
  • Prior stylized studies showed Transformers can match gradient-based learning algorithms on regression and can represent gradient descent for real-valued tasks.Other work examined sample complexity, mixtures of tasks, and adaptive algorithm selection from in-context examples.Source: §6 · p. 9
  • The paper complements work on sequence-model capabilities and formal-language tasks with a comprehensive study of in-context learning.Prior sequence-model research addressed theoretical expressiveness, empirical capabilities, and limitations.Source: §6 · p. 9

7 CONCLUSION

The conclusion highlights limits in Transformer in-context learning, sample-efficient use of teaching sequences, and pretrained models' capacity to solve stylized tasks. The appendix roadmap and clarifications describe supporting experiments, unresolved mechanisms, and task-design choices.

7 CONCLUSION

  • The conclusion reports that pretrained models learn mechanisms for stylized in-context tasks and that LLMs can match certain baseline learning algorithms in a more practical setting.It identifies future questions about teaching examples, pretraining, parity learning, and the mechanisms used by attention, recurrence, and convolution.Source: §7 · p. 9

A ROADMAP

  • Appendix experiments cover robustness to distribution shifts and noisy labels, sample complexity, task difficulty, parity learning, teaching sequences, LLMs, and implementation details.The roadmap assigns these topics to Sections D through H.Source: A ROADMAP · p. 14

B CLARIFICATIONS

  • The Boolean-input LLM task has over 10^30 possible target functions even at dimension 7, making prior exposure to the same labeled sequences virtually impossible according to the authors.They argue accurate prediction therefore requires using the in-context examples as a learning algorithm.Source: B CLARIFICATIONS · p. 14
  • The authors do not identify the algorithm LLMs use; their results show competitiveness with nearest neighbors, not evidence of nearest-neighbor or gradient-descent implementation.They also note that attention-free architectures match Transformers on most evaluated tasks and are slightly worse on a few.Source: B CLARIFICATIONS · p. 14
  • Task dimensions and example counts were chosen to make learning feasible in light of known algorithms and sample-complexity results.For parity tasks, the experiments distinguish gradient-based learning from approaches such as Gaussian elimination.Source: B CLARIFICATIONS · pp. 14–15

C.1 DESCRIPTIONS OF TASKS

The appendix defines Boolean-function tasks and details input distributions and baselines for evaluating in-context learning. It also summarizes performance curves across function classes and architectures.

C.1 DESCRIPTIONS OF TASKS

  • The test bed includes Conjunctions, Disjunctions, sparse Disjunctions, 3-term DNFs, 3-clause CNFs, and sparse Majority functions.Sparse Disjunctions use at most three variables; the considered DNF/CNF variants have three terms or clauses.Source: §C.1 · p. 15
  • The task set also includes 0-1 Threshold functions, Integer Halfspaces, Parities and sparse Parities, plus a nearest-neighbor prediction task.The nearest-neighbor task labels each later point using its nearest preceding input point.Source: §C.1 · pp. 15–16

C.4 ADDITIONAL PLOTS FOR IN-CONTEXT LEARNING BOOLEAN FUNCTIONS

  • Across in-context example curves, Transformers perform near-optimally on Conjunctions and Disjunctions, but are less accurate on Threshold, Majority, and Integer Halfspace tasks and at chance on Parities.Figure 6 reports mean accuracy by number of examples; Figure 7 supplies curves for five architectures on representative tasks.Source: Figure 6 · p. 16, Figure 7 · p. 17, §C.4 · p. 18

C.1 DESCRIPTIONS OF TASKS

  • Function classes are finite, but Conjunctions alone have 3^n possible functions; at n=30 the class exceeds 100 trillion.For a prompt with 50 examples, the number of input-order permutations is 50! for one input set and target function.Source: §C.1 · p. 16

C.2 DETAILS OF INPUT AND FUNCTION DISTRIBUTION

  • For Conjunctions, Disjunctions, and 3-CNF, inputs are rebalanced because uniform inputs can let a constant classifier exceed 99% accuracy.The modified distribution ensures at least 30% of prompt points have a positive or negative label; function sampling varies by task class.Source: §C.2 · pp. 16–17

C.3 DETAILS OF BASELINES

  • Baselines include a constant Null classifier, an averaging estimator, nearest neighbors, and a tuned one-hidden-layer ReLU feedforward network.The NN baseline predicts from the labels of the nearest prior inputs; the FFN is selected from tested widths, optimizers, and learning rates.Source: §C.3 · pp. 17–18

D.1 ROBUSTNESS OF MODELS

The appendix tests robustness, examines a correlate of task difficulty, and explores sample complexity for Conjunction learning. Results include transfer across function distributions, weak performance under extreme shifts, and near-perfect accuracy with a finite training set.

D.1 ROBUSTNESS OF MODELS

  • Transformers trained on Conjunctions perform effectively when tested on functions with twice the training literal probability, but performance degrades when expected literal counts exceed 80%.Training functions have an expected n/3 literals; the tested shifted distribution has an expected 2n/3 literals.Source: §D.1 · pp. 18–19
  • A Transformer trained on 0-1 Threshold functions achieves over 90% on Conjunctions, outperforming its in-distribution Threshold performance.The authors suggest transfer may reflect Conjunctions being simpler Threshold functions with b dependent on w.Source: §D.1 · p. 19
  • All tested architectures are robust to 10% label flips during training on Conjunctions or Disjunctions when evaluated on clean inputs.Figure 12 presents the noisy-training results for Conjunctions and Disjunctions.Source: §D.1 · p. 19, Figure 12 · p. 20

D.2 WHAT DETERMINES THE DIFFICULTY OF IN-CONTEXT LEARNING?

  • VC dimension does not track model performance consistently: Transformers do well on Conjunctions and sparse Disjunctions but fail on Parities and sparse Parities with similar VC dimensions.The authors frame VC dimension as a measure related to worst-case sample complexity.Source: §D.2 · p. 20
  • Pairwise function correlation tracks performance reasonably well: it is high for classes where models do well and lowest for Parities, whose pairwise correlations are zero.The measure is estimated from sampled functions and uniformly sampled inputs; Table 2 reports estimates across classes.Source: §D.2 · p. 20, Table 2 · p. 20
  • Pairwise correlation is an incomplete difficulty measure because it uses uniform inputs even when some task training distributions are nonuniform.The authors identify a clearer account of in-context-learning difficulty as an open problem.Source: §D.2 · p. 20

D.3 EXPLORING SAMPLE COMPLEXITY

  • With 4,096 training examples, Transformers and LSTMs achieve close to perfect Conjunction accuracy, while increasing Transformer depth has negligible effect.Evaluation accuracy averages across all points in each sequence, so early mistakes keep the mean below 100% even when accuracy later converges near perfectly.Source: §D.3 · p. 21

E CURIOUS CASE OF LEARNING PARITIES

Transformers and other architectures fail to learn Parities and Sparse Parities beyond chance, despite evidence that gradient-trained feedforward networks and known algorithms can perform much better. The authors examine training gradients and discuss possible reasons, but leave the failure mechanism unresolved.

E CURIOUS CASE OF LEARNING PARITIES

  • Transformers and other architectures do not exceed chance-level accuracy on Parities and Sparse Parities, unlike their performance on other function classes.The authors describe this as a failure to find any learning algorithm that performs beyond chance.Source: E CURIOUS CASE OF LEARNING PARITIES · p. 21
  • Gradients in Transformer attention and feedforward parameters approach zero after a few iterations on Parities, unlike on Conjunctions, but the cause remains unclear.Figure 15 depicts parameter-gradient means and loss across training iterations.Source: E CURIOUS CASE OF LEARNING PARITIES · p. 22, Figure 15 · p. 22
  • Feedforward networks trained with gradient descent achieve near-perfect accuracy on PARITY-(10, 2) using the same number of examples as the in-context models.The authors note that the hardness results primarily concern uniform inputs and do not apply to input sets with teaching sequences.Source: E CURIOUS CASE OF LEARNING PARITIES · p. 22, Figure 16 · p. 23
  • For sparse parities, nearest-neighbor classification can perform reasonably well at small k and n; the evaluated LLMs are better than random and comparable to that baseline.For k = 2, the authors explain that relevant bits match with probability 1/4.Source: E CURIOUS CASE OF LEARNING PARITIES · p. 22
  • Gaussian elimination can learn Parities with O(n) examples, but the authors suggest Transformers may not implement it in-context because it is not parallelizable.They present this as a possibility, not an established explanation for the observed failure.Source: E CURIOUS CASE OF LEARNING PARITIES · pp. 22–23

F TEACHING SEQUENCE EXPERIMENTS: ADDITIONAL DETAILS

This appendix details how teaching sequences are constructed for Boolean function classes and describes frozen-GPT experiments and a method for identifying nearest-neighbor heads. It also reports that frozen GPT experiments consistently revealed 8–10 such heads.

F TEACHING SEQUENCE EXPERIMENTS: ADDITIONAL DETAILS

  • A teaching sequence uniquely identifies a function, and the teaching dimension is the largest, over functions, of the size of the smallest such sequence.For Conjunctions and Disjunctions, the described sequence has k + 2 points; the exact function can be identified from those points.Source: F TEACHING SEQUENCE EXPERIMENTS:… · p. 23
  • Conjunction teaching examples use two points setting literals true and other variables to opposite extremes, then one point per literal; Disjunctions use a similar construction.For monotone 3-DNFs, the sequence is assembled from clause-specific examples, and a similar construction is given for 3-CNFs.Source: F TEACHING SEQUENCE EXPERIMENTS:… · p. 23
  • For PARITY-(n, k), a teaching sequence uses k points, each setting one relevant bit to 1 and all other bits to 0; this does not teach the full PARITY-n class.Source: F TEACHING SEQUENCE EXPERIMENTS:… · p. 24

G.1 ADDITIONAL DETAILS OF FROZEN GPT EXPERIMENT

  • Frozen GPT-2 experiments aim to test whether pretrained models contain mechanisms useful for the stylized tasks, while updating only input and output layers.The GPT-2 XL weights remain fixed; a ReLU embedding network and an output feedforward network replace the original input/output layers for Boolean inputs.Source: §G.1 · p. 24

G.2 DETECTING NEAREST NEIGHBOUR HEADS IN GPT-2

  • The nearest-neighbor score measures attention to the label paired with the query’s nearest earlier input, averaged over late sequence positions and examples.Heads scoring above 0.5 are labeled nearest-neighbor heads; frozen-GPT runs consistently found 8–10 such heads.Source: §G.2 · p. 25

G.3 ADDITIONAL DETAILS AND RESULTS OF DIRECT EVALUATION EXPERIMENT

Direct evaluation prompts LLMs with labeled Boolean examples and asks them to predict the label of the next input. GPT-2 fails on both reported tasks even at five dimensions, while LLaMA performance rises with model scale and smaller LLaMAs show nontrivial performance.

G.3 ADDITIONAL DETAILS AND RESULTS OF DIRECT EVALUATION EXPERIMENT

  • The direct-evaluation prompt gives a natural-language instruction and labeled Boolean examples, then asks for the label of the next query point across Conjunction, Majority, and Parity tasks.The experiments use dimensions 5, 7, and 15, with 50 points when d < 10 and 100 when d ≥ 10.Source: §G.3 · p. 25
  • OpenAI-model evaluations were limited to 100 functions because evaluating more was prohibitively expensive.The authors averaged over 2,000 functions for baselines and 1,000 for LLaMA-2 models.Source: §G.3 · p. 25
  • GPT-2 fails on both tasks even at 5 dimensions, while LLaMA performance increases gradually with model scale and smaller LLaMAs achieve nontrivial performance.Figure 19 reports Conjunction and Majority results across dimensions for LLaMA models of different scales and GPT-2.Source: §G.3 · pp. 25–26, Figure 19 · p. 26

H IMPLEMENTATION DETAILS

The implementation uses tuned model architectures and training settings across the experiments, with learning rate and depth among the consequential hyperparameters. The appendix also summarizes prior empirical work on LLM in-context learning and related mechanisms.

H IMPLEMENTATION DETAILS

  • Experiments use PyTorch implementations, including adapted official Hyena and DSS code and a custom RetNet implementation.The experiments were run on 16 NVIDIA Tesla V100 GPUs, with grid search over hyperparameters and Adam learning rates.Source: H IMPLEMENTATION DETAILS · p. 26
  • Learning rate has the most influence on performance across most tasks and architectures, while depth has a nontrivial influence; heads and widths show no nontrivial difference.Source: H IMPLEMENTATION DETAILS · pp. 26–27

I ADDITIONAL RELATED WORK

  • Direct evaluations use greedy decoding with a maximum of two generated tokens; related work discusses prompt selection and order, training-data prevalence, label sensitivity, induction circuits, and scale.Prior work also reports that overriding semantic priors to learn an in-context task emerges with scale.Source: H IMPLEMENTATION DETAILS · p. 27, I ADDITIONAL RELATED WORK · p. 27

Related papers

Loading 2310.03016v1…