Gaussian Match-and-Copy: A Minimalist Benchmark for Studying Transformer Induction

TL;DR

Gaussian Match-and-Copy (GMC) isolates long-range, second-order correlations to study how Transformer architectures develop the $PTH\to IH$ retrieval circuit. The authors show that GMC reproduces core mechanisms observed in larger models, reveals a clear architectural gap between attention-based and non-attention models, and supports a conditional max-margin implicit bias in a minimal two-layer attention model under a high-probability geometric event. Empirically, GMC-trained Transformers consistently exhibit PTH and IH head emergence concurrent with sharp loss drops, and the learned mechanisms transfer to non-Gaussian data (e.g., Omniglot) via lightweight embedding adaptation. The work also provides a theoretical framework showing that, under specific assumptions, gradient descent on MSE converges in direction to the max-margin separator with the weight norm diverging logarithmically, offering a principled lens on induction dynamics. Overall, GMC serves as a cheap, controllable diagnostic for probing induction in sequence models and guides future architectural and theoretical investigations into long-range retrieval and in-context learning.

Abstract

Match-and-copy is a core retrieval primitive used at inference time by large language models to retrieve a matching token from the context then copy its successor. Yet, understanding how this behavior emerges on natural data is challenging because retrieval and memorization are entangled. To disentangle the two, we introduce Gaussian Match-and-Copy (GMC), a minimalist benchmark that isolates long-range retrieval through pure second-order correlation signals. Numerical investigations show that this task retains key qualitative aspects of how Transformers develop match-and-copy circuits in practice, and separates architectures by their retrieval capabilities. We also analyze the optimization dynamics in a simplified attention setting. Although many solutions are a priori possible under a regression objective, including ones that do not implement retrieval, we identify an implicit-bias regime in which gradient descent drives the parameters to diverge while their direction aligns with the max-margin separator, yielding hard match selection. We prove this max-margin alignment for GD trajectories that reach vanishing empirical loss under explicit technical conditions.

Figures (13)

• Figure 1: Gaussian Match-and-Copy task (\ref{['def:mc']}). The query matches a hidden context token via correlation, and the target copies the successor of that token.
• Figure 2: Co-occurrence of Loss Drop and Circuit Emergence. The sudden drop in test loss aligns perfectly with the saturation of PTH and IH attention scores. Here: 2-layer Llama 3 trained on GMC with $T=8$, ${d_{\mathrm{in}}}=16$, $\mathbf C=\frac{1}{(1.2)^2}\mathbf I$, and hyperparameters from \ref{['tab:mc-experiment-specs']} in the appendix. See \ref{['app:robustness-emergence']} for other settings (varying depth, $T$, ${d_{\mathrm{in}}}$, and $\mathbf C$), all showing the same behavior.
• Figure 3: Transfer Learning. A model pretrained on GMC can reuse its attention layers to solve Omniglot few-shot classification via lightweight embedding adaptation.
• Figure 4: Architecture Gap. Transformers solve GMC; alternatives struggle.
• Figure 5: Robustness of Emergence. The alignment between the loss drop (top row) and the saturation of PTH/IH scores (bottom rows) across different architectures (Llama 3, GPT-2) and GMC configurations (varying $c$, ${d_{\mathrm{in}}}$, $T$).

Theorems & Definitions (18)

• Definition 1.1: Gaussian Match-and-Copy
• Theorem 4.1: Max-margin implicit bias for MSE in the minimal model
• Theorem D.1: Implicit bias of $\mathrm{CE}$, adapted from Soudry18ImplicitBiasCE
• Theorem D.2: Implicit Bias of MSE
• Remark D.3: On the scope of the assumptions
• Definition E.1: Correction vector $\tilde{w}$
• Lemma E.3: Restoring force from support-tail of support drift
• proof
• Lemma F.1: Per-sample gradients, feature-space form
• proof