Source-linked AI summary
Theoretical Understanding of In-Context Learning in Shallow Transformers with Unstructured Data
Yue Xing, Xiaofeng Lin, Chenheng Xu, Namjoon Suh, Qifan Song, Guang Cheng
TL;DR
Prior theoretical ICL analyses often assume structured examples, while text prompts distribute words across tokens without explicit structural or causal links. This paper studies shallow transformers on unstructured prompts and finds that two attention layers with a look-ahead mask support ICL, while positional encoding helps connect inputs and outputs.
Problem
Prior theoretical studies leave attention masks and positional encoding out of analyses of single-layer transformers, despite real text prompts separating example inputs and outputs across tokens.
Source: §1 · pp. 2–3Method
Results
With infinite prompts, unstructured-data ICL requires at least two attention layers and a look-ahead mask; positional encoding can connect x_i and y_i, and higher embedding dimension improves performance.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
The paper's account favors structured data, but studies unstructured prompts because real text is not curated with direct structural or causal links between words.
Source: §1 · pp. 2–3Without input embedding, positional encoding lacks flexibility when D is much larger than p = d + 1.
Source: §5.2 · p. 8
Abstract
from arXiv · showhide
Large language models (LLMs) are powerful models that can learn concepts at the inference stage via in-context learning (ICL). While theoretical studies, e.g., \cite{zhang2023trained}, attempt to explain the mechanism of ICL, they assume the input $x_i$ and the output $y_i$ of each demonstration example are in the same token (i.e., structured data). However, in real practice, the examples are usually text input, and all words, regardless of their logic relationship, are stored in different tokens (i.e., unstructured data \cite{wibisono2023role}). To understand how LLMs learn from the unstructured data in ICL, this paper studies the role of each component in the transformer architecture and provides a theoretical understanding to explain the success of the architecture. In particular, we consider a simple transformer with one/two attention layers and linear regression tasks for the ICL prediction. We observe that (1) a transformer with two layers of (self-)attentions with a look-ahead attention mask can learn from the prompt in the unstructured data, and (2) positional encoding can match the $x_i$ and $y_i$ tokens to achieve a better ICL performance.
1 Introduction
Prior theoretical ICL studies largely use structured examples, while real text prompts distribute words across tokens without explicit structural or causal links. This paper analyzes shallow transformers on unstructured prompts, finding that two attention layers with a look-ahead mask support ICL and positional encoding helps match inputs to outputs.
1 Introduction
- Prior analyses of single-layer transformers omit attention masks and positional encoding, while ordinary text prompts spread each example across multiple word tokens.The paper contrasts these prompts with structured formats that place an example's input and output in one column.Source: §1 · pp. 1–2
- For unstructured data E2, one layer fails, whereas two attention layers with a look-ahead mask can perform ICL.The paper identifies both the two-layer structure and mask as crucial.Source: §1 · p. 2
- Positional encoding can improve performance by better matching x_i and y_i before the second layer performs ICL on structured data.A larger input embedding dimension and multiple attention layers can also help matching.Source: §1 · p. 2
- The paper does not claim unstructured data is superior: its account prefers structured data, though unstructured text remains unavoidable in practice.Text prompts are often uncurated and lack direct structural or causal information between words.Source: §1 · pp. 2–3
2 Related Works
Related theoretical work studies ICL in simple transformers, often using structured data, and examines transformer components including attention masks, layers, and positional encoding.
2 Related Works
- Theoretical ICL research studies simple one-attention-layer transformers on structured data, including expressive power, convergence, and generalization.The cited literature also studies linear attention.Source: §2 · p. 3
- Related work discusses attention masks across vision and multimodality, and analyzes internal layer mechanisms in BERT.The cited mask studies report empirical successes in those fields.Source: §2 · p. 3
- Positional encoding supplies word-position information and has modality-specific variants, while Bai et al. [2023] prove it enables more complex ICL procedures involving algorithm selection.The paper notes limited theoretical understanding of positional encoding.Source: §2 · p. 3
3 Notations and Architecture
The paper specifies a linear-regression data distribution and a shallow-transformer setup for studying ICL. Its two-layer architecture uses masked softmax attention followed by linear attention, with controlled experiments to assess component effects.
3 Notations and Architecture
- Examples and the query use independent standard Gaussian inputs, with noiseless linear responses sharing a prompt-specific parameter θ.The parameter is drawn from N(0, I_d/d), and the study omits more complex data relationships to focus on positional encoding.Source: §3 · p. 3
- Training assumes infinitely many prompts and minimizes prediction loss for example outputs and the query output.The supplied text identifies these predictions as the quantities in the loss.Source: §3 · pp. 3–4
- The first layer applies masked softmax self-attention to embedded inputs plus positional encodings; a second linear-attention layer produces the query prediction.The prediction is a linear function of the final column of the second-layer output.Source: §3 · p. 4
- Experiments use input dimension d = 5, D = 30 examples per prompt, and 50k training steps, with ten-repetition results deferred to the appendix.Without a mask, positional encoding can exploit later y_i tokens during training, yielding near-zero loss that does not help inference.Source: §3 · p. 4
4 Two-Layer Transformer + Attention Mask
For unstructured E2 prompts, the paper shows that two attention layers with an attention mask support ICL, while a single layer fails; it then explains these outcomes theoretically.
4.1 Empirical Observation
- With E2 and no positional encoding, prediction loss drops as D increases only for two-layer transformers with an attention mask; omitting either component results in ICL failure.Source: §4.1 · p. 5, Figure 2 · p. 5
4.2 Why One-Layer Transformer Fails
- A single attention layer cannot pass input–target similarity information to the prediction: x_i tokens carry zero-valued output entries, while y_i tokens lack the needed similarity information.Theorem 1 formalizes the failure under its assumptions: as D grows, the optimal output vector norm tends to zero and the minimal prediction error remains in Θ(1).Source: §4.2 · p. 5
4.3 Why Two-Layer Transformer Succeeds
- Theorem 2 shows that two layers with an attention mask can perform ICL on E2 even without positional encoding, under specified first-layer parameter assumptions.The second layer can then predict y_q; the proof also uses a cross-term that becomes negligible when D is large enough.Source: §4.3 · pp. 5–6
- Simulations show only a slight performance difference when the first layer’s W_KQ,1 is fixed to zero, consistent with that layer mixing E2 columns asymmetrically for second-layer ICL.Source: §4.3 · p. 6
5 Positional Encoding
Positional encoding improves ICL performance and changes first-layer attention to better match input tokens with their corresponding labels. The section analyzes this matching and its limits across embedding dimensions and attention-mask settings.
5 Positional Encoding
- With infinite training prompts, positional encoding improves ICL performance.This section also studies fully trainable positional encoding.Source: §5 · p. 6
5.1 Intuition: Why Positional Encoding Helps
- Without positional encoding, first-layer attention is nearly uniform across columns, while second-layer high weights often fall on odd positions containing y_i information.With positional encoding, first-layer attention instead favors x_i positions and gives near-zero weight to odd positions.Source: Figure 4 · p. 6, Figure 5 · p. 6, §5.1 · pp. 6–7
5.2 Positional Encoding without Input Embedding
- With the constructed positional encoding and D < d, the first layer combines each x_i with its corresponding y_i with high probability as d → ∞.Without input embedding, positional encoding lacks flexibility when D ≫ d, and the first-layer output can retain information from earlier x_j; a small simulation found the construction close to fully trainable PE and better than no PE.Source: §5.2 · pp. 7–8
5.3 Benefit of Positional Encoding with Large Input Embedding Dimension
- With D < p and p ≫ d under stated conditions, positional encoding matches x_i with y_i in the first layer, supporting a later linear-attention predictor with E(by_q − y_q)^2 = O(d/D).Linear attention in the second layer avoids positional-encoding bias in attention scores; when D increases, small p performs worse than large p with smaller D.Source: §5.3 · pp. 8–9
5.4 With PE but without Attention Mask
- With positional encoding but no mask, label leakage can make training loss nearly zero while prediction loss remains Θ(1); simulations find low ICL loss only with two layers and a mask.The first-layer label information leaks from token 2i to the preceding token, but this does not yield successful inference-phase prediction.Source: Figure 6 · p. 9, §5.4 · p. 9
6 Conclusions
The conclusion identifies at least two attention layers with a look-ahead mask as necessary for learning from separated-token prompts; positional encoding and higher embedding dimensions can further improve ICL performance.
6 Conclusions
- At least two attention layers with a look-ahead mask are needed for unstructured prompts; positional encoding connects x_i and y_i, and higher embedding dimensions improve performance.The paper also says positional encoding balances the mean and variance in prediction loss.Source: §6 · p. 10
- The authors identify softmax attention in the second layer, deeper attention stacks, positional encoding’s generalization gap, and chain-of-thought as future directions.The paper studies positional encoding with almost infinite examples and does not consider the generalization gap.Source: §6 · p. 10
- The analysis uses a simple transformer and may not capture the full characteristics of transformers used in practice.The authors suggest extending the analysis to large-scale transformers.Source: §6 · p. 10
A One-Layer Transformer without PE
- The one-layer proof pairs θ and −θ, making the loss-minimizing prediction approximately zero.This is the proof’s stated key idea for the one-layer case.Source: A One-Layer Transformer without PE · p. 13
B Two-Layer Transformer without PE
- In the two-layer analysis, A1 is closely related to x_q, while A2 is related to the average of x_i values and is negligible.The passages identify A2 as negligible after comparing it with A1.Source: B Two-Layer Transformer without PE · p. 16
C.1 Theorem 3
The theorem’s analysis describes how positional encoding shapes first-layer attention scores, concentrating attention on paired input and output representations.
C.1 Theorem 3
- For the stated regime, positional encoding is the dominant component in determining the attention-score matrix.The supplied passage gives the regime as “d log(d).”Source: §C.1 · p. 17
- The positional-encoding score exp(v_x c^2) dominates the sum of other scores in the same column.The comparison is stated for the attention scores associated with columns O_{:,2i}.Source: §C.1 · p. 18
- The two paired representations each receive attention score 1/2 + o(1), while all other scores are negligible.The paired representations are described as (p x_i1 + x_i, 0) and (p x_i2, y_i).Source: §C.1 · p. 18
C.2 Theorem 4
Theorem 4 gives a condition under which one score term dominates another, while Theorem 5 reports concentration of ξ1 on the x_i term.
C.2 Theorem 4
- When ∥p̃x_jk∥ ≫ 1, v_x ≫ 1, and D ≫ d^2, ζ1 dominates ζ2.Source: §C.2 · p. 19
C.3 Theorem 5
- Theorem 5 assumes all singular values of W_in are equal to κ.Source: §C.3 · p. 19
- For d log p, ξ1 concentrates on the xi term; the supplied passage introduces an “i.e.” clarification but does not include its contents.Source: §C.3 · p. 21
C.4 Theorem 6
Theorem 6 states that positional encoding introduces no bias in the second-layer ICL prediction and gives an error bound; the appendix also identifies negligible terms and describes configuration controls.
C.4 Theorem 6
- In the Theorem 6 analysis, C2 and C4 are negligible relative to C1, C3, and C5.Source: §C.4 · p. 22
- The second-layer prediction has no positional-encoding bias, and E(ŷ_q − y_q)^2 = O(d/D).Source: §C.4 · p. 23
D Configuration Details
- Table 2 describes instructions for controlling the transformer components in the experiments.The paper identifies attention mask, multiple attention layers, and positional encoding as components.Source: D Configuration Details · p. 24, Table 2 · p. 24
E Other Additional Experiment Results
The additional experiments show attention-score visualizations for prompts with and without positional encoding, and report ICL performance across ten repetitions under different settings.
E Other Additional Experiment Results
- Figure 7 shows one example’s full attention scores for a masked, two-layer, single-head transformer without positional encoding, using 31 x-y pairs across 62 tokens.Each row is the attention of one token.Source: Figure 7 · p. 24
- Figure 8 shows full attention scores for a masked, two-layer, single-head transformer with positional encoding in E2 format.Source: Figure 8 · p. 24
F.1 Ten Repetitions
- Figure 9 reports ICL performance across ten repetitions under settings that vary transformer depth, positional encoding, and attention masking.Source: §F.1 · p. 25, Figure 9 · p. 25
F.2 Additional Simulation Regarding PE
The simulation varies the number of examples with d=5 and p=12 to examine positional encoding at large p and performance as D grows. Its comparisons show better performance for p=64 at D=30, while increasing D to 42 at p=12 does not match the Table 1 result.
F.2 Additional Simulation Regarding PE
- The simulation follows Garg et al. [2022]1, fixes d=5 and p=12, and varies the number of examples D to examine positional encoding and growing-example performance.Source: §F.2 · p. 26
- At d=5, p=64, D=30, test loss is 0.1877; p=64 performs better than p=12 at D=30, while p=12 at D=42 remains worse than Table 1.Source: Table 3 · p. 26, §F.2 · p. 26