PaperlayerView exact revision on arXiv

Source-linked AI summary

In-context Learning of Single-index Targets: Comparing Kernel and Feature Learners

Haotian Gu, Yizhou Xu, Lenka Zdeborová

arXiv:2610.01712v1cs.LGcond-mat.dis-nnstat.ML

TL;DR

Most ICL theory has focused on linear targets, leaving the learning of nonlinear single-index tasks as an open comparison between fixed-kernel and feature-learning approaches. The paper analyzes both architectures on a shared task setup and finds that their relative advantages vary with pretraining, task diversity, and context lengths, alongside different context-length scaling behavior.

  • Problem

    Most theory focuses on linear targets, while nonlinear single-index tasks require learning a nonlinearity whose relevant direction changes between tasks and must be extracted from demonstrations.

    Source: §1 · p. 1

  • Method

    The paper compares fixed-kernel and feature-learning one-layer attention learners on the same task setup, using replica calculations to predict memorization and generalization errors.

    Source: §1 · p. 2

  • Results

    The feature learner typically reaches low error with fewer pretraining prompts, the kernel learner can achieve lower memorization error for small task pools with sufficient pretraining, and their context-length scalings differ.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The relative advantage of kernel and feature learning depends on how each architecture interacts with pretraining data, task diversity, and information available in context.

    Source: §5 · p. 10

    The analysis focuses on one-layer linear attention, spherical-harmonic features, and an idealized two-stage procedure that learns the scalar readout from sufficient additional data.

    Source: §5 · p. 10

Abstract

from arXiv · show

In-context learning (ICL) enables a pretrained model to infer a task from demonstrations without updating its parameters. While much of the existing theory focuses on linear target functions, in this paper we study nonlinear cases by comparing two one-layer attention architectures on the same family of single-index tasks. A kernel learner first maps inputs through a fixed nonlinear feature map and then applies linear attention, whereas a feature learner applies attention to the original input, followed by a learned nonlinear readout. We derive predictions for their memorization and generalization errors using the replica method, retaining the effects of pretraining size, task-pool diversity, and training and inference context lengths. The resulting predictions closely match numerical experiments across a broad range of regimes. Our analysis yields phase diagrams that characterize when each architecture is advantageous as the amount of pretraining data, task diversity, and context lengths vary. We further identify qualitatively different context-length scalings for the two learners. Together, these results clarify how architectural choices interact with the dataset and govern nonlinear in-context learning.

1 Introduction

Existing ICL theory largely studies linear targets, leaving open how nonlinear single-index tasks should be learned when their relevant direction must be inferred from demonstrations. This paper compares fixed-kernel and feature-learning attention architectures on a common task setup and finds their relative advantages depend on pretraining, task diversity, and context lengths.

1 Introduction

  • Most prior ICL theory focuses on linear targets; single-index tasks pose the additional question of how to learn the input nonlinearity when the relevant direction changes across tasks.Prior work also reports that pretrained Transformers can learn single-index targets in context.Source: §1 · pp. 1–2
  • The comparison uses the same task pool, pretraining prompts, and inference prompts, and measures memorization on seen tasks separately from generalization on fresh tasks.The task-pool size varies to examine the memorization–generalization tradeoff.Source: §1 · p. 2
  • Replica-based predictions for both learners accurately predict their performance across a broad range of regimes.The analysis considers how performance depends on input and feature dimensions, pretraining size, task-pool size, and context lengths.Source: §1 · p. 2
  • The feature learner typically reaches low error with fewer pretraining prompts, while the kernel learner can achieve lower memorization error when the task pool is small and pretraining is sufficient.Context-length mismatch can also change their relative advantage.Source: §1 · p. 2
  • For degree-k spherical harmonics, the kernel learner requires context length of order d^k to recover that component, whereas the feature learner transitions at L = Θ(d) and can exploit all teacher Hermite components together.The analysis also characterizes differing training and test context lengths.Source: §1 · p. 2

2 Settings

The setting distinguishes performance on seen task-pool members from performance on fresh tasks, then compares two one-layer attention learners that place nonlinear transformations on opposite sides of attention.

2.2 Kernel learner

  • The kernel learner applies a fixed nonlinear feature map before linear self-attention, with spherical harmonics used as an analytically convenient realization.After embedding, training the attention layer is equivalent to linear regression in the entries of its attention matrix; random-feature learners are left for future work.Source: Figure 1 · p. 3, §2.2 · p. 4

2.1 Task

  • Memorization evaluates tasks sampled from the finite training pool, while generalization evaluates fresh tasks sampled uniformly from the sphere; distribution shift is not considered.The teacher assumption requires a nonzero linear component, which places the problem in the O(d^2) learning regime.Source: §2.1 · p. 3

2.3 Feature learner

  • The feature learner applies attention to original inputs and then uses a scalar nonlinear readout, with attention and readout trained in two stages.The readout-training idealization takes additional prompts to infinity, making their share of total prompts vanish.Source: §2.3 · pp. 4–5

3 Main results

The section derives replica-based error predictions for both architectures and compares their pretraining- and context-length behavior. It finds distinct advantages and context-scaling tradeoffs, including a training/inference mismatch for the feature learner.

3 Main results

  • Replica-based formulas predict both learners’ memorization and generalization errors, and the resulting predictions are validated numerically.The kernel and feature results use their respective high-dimensional limits; the analysis uses a Gaussian-equivalence ansatz that is validated numerically.Source: §3 · pp. 5–6
  • Kernel recovery of a degree-k harmonic sector requires L on the order of d^k, while the feature learner uses L = Θ(d) to exploit all Hermite components.For the kernel learner, the sector is fully recovered when L ≫ d^k, partially recovered when L = Θ(d^k), and unrecovered when L ≪ d^k.Source: §3 · p. 7
  • When p > 1, the feature learner approaches its infinite-pretraining limit in the regime d^2 ≪ N ≪ D^2, before the kernel learner does.The comparison follows from the learners’ attention-matrix parameter counts, d^2 and D^2 = Θ(d^2p), and their respective convergence scales.Source: §3 · p. 7
  • Longer inference context always improves the kernel learner, but need not improve the feature learner because its Stage-II readout is fixed after training.When inference length is fixed at eL, both learners reach optimal performance at training length L = eL; increasing training context beyond eL is not beneficial.Source: §3 · p. 7

4 Experiments

The experiments compare the learners on a common teacher and prompts while varying pretraining size, task-pool diversity, and context lengths. Their relative performance shifts across these regimes, and their context-length scaling differs.

4 Experiments

  • The experiments use the same teacher, task pool, and prompts for both learners, and vary pretraining size, task diversity, and training and inference lengths.Unless stated otherwise, d=32, D=559, Δ=0.01, and λ=0.001; the kernel learner uses a complete degree-one and degree-two harmonic basis containing the teacher.Source: §4 · p. 7

4.1 Varying the number of training prompts

  • At N=36864, feature memorization/generalization errors are 0.061/0.421 versus 0.088/0.675 for the kernel learner.The feature learner reaches low error earlier, in the range d^2 ≪ N ≪ D^2.Source: Figure 2 · p. 8, §4.1 · p. 8
  • At N=1200000, kernel memorization error is 0.037 versus 0.046 for feature, but generalization mean error is 0.657 versus 0.354.The kernel learner's generalization is worse at low task diversity.Source: Figure 2 · p. 8, §4.1 · pp. 8–9

4.2 Varying task diversity and context lengths

  • Kernel memorization advantage occurs at large N and small N0, while increasing N0 removes it; its generalization advantage occurs only at low N.With context lengths varied, the kernel learner can be favored by a large training–inference mismatch or when both lengths are large.Source: Figure 3 · p. 9, §4.2 · p. 9

4.3 Learning different harmonic degrees

  • For degree-one and degree-two components, the kernel learner recovers them at separate scales L ≍ d and L ≍ d^2, while the feature learner recovers both together.The feature learner has lower error throughout the displayed range; the separation in sample complexities grows as d increases.Source: §4.3 · pp. 9–10

5 Conclusion and Future Work

The paper compares kernel and feature learners for nonlinear in-context learning, using replica calculations and proving distinct context-scaling mechanisms in the population limit. Its analysis is scoped to particular architectures, features, and training assumptions.

5 Conclusion and Future Work

  • Replica calculations characterize how memorization and generalization errors depend jointly on pretraining size, task diversity, and training and inference context lengths.In the population limit, the paper rigorously proves two context-scaling mechanisms.Source: §5 · p. 10
  • The relative advantage of kernel and feature learning depends on their interaction with pretraining data, task diversity, and information available in context.This is the paper's stated broad conclusion about architectural advantage in ICL.Source: §5 · p. 10
  • The analysis is limited to one-layer linear attention, spherical-harmonic features, and an idealized two-stage procedure with sufficient additional data for learning the scalar readout.The authors identify random features, learned embeddings, jointly trained readouts, softmax attention, deeper Transformers, and broader targets as open directions.Source: §5 · p. 10

A Derivation of the kernel learner

The derivation expands the teacher in a spherical-harmonic basis, derives kernel state equations and a limiting spectrum, then relates the kernel learner's pretraining calculation to a linear ICL model using Gaussian equivalence.

A.1 Teacher expansion

  • The teacher is expanded in normalized Gegenbauer polynomials, which converge to Hermite polynomials at each fixed degree; the residual is approximated by an independent Gaussian with matching variance.The paper assumes sufficient regularity for the coefficients and teacher variance to converge.Source: §A.1 · pp. 14–15

A.2 Fixed point equations

  • The kernel learner's state equations determine five order parameters and three conjugate variables.The equations are presented without specializing to a particular task-Gram-matrix spectrum.Source: §A.2 · p. 16, §A.3 · p. 16

A.3 Limiting spectrum

  • When N0/d has a finite positive limit, the degree-one block has a Marchenko–Pastur spectrum, while each degree-k block for k≥2 contributes asymptotically as an identity matrix.If degree one is excluded, the limiting distribution reduces to a point mass at νin.Source: §A.3 · pp. 16–17

A.4 Reduction to Takanami et al.

  • The reduction first establishes an exact algebraic correspondence, then uses Gaussian equivalence; it transfers the pretraining saddle-point calculation, not the inference-error calculation.The approximation characterizes the effective problem through leading first- and second-order statistics while neglecting higher-order correlations.Source: §A.4 · pp. 17, 18, 20

A.5 Derivation of the main results

This section derives the kernel learner’s limiting errors and context-length behavior and develops replica-based fixed-point equations for the feature learner, using Gaussian approximations and symmetry assumptions to reduce the calculations. It also combines expressions for f2 and f3 with stationarity equations to recover Result B.1 under a stated spectral-convergence condition.

A.5 Derivation of the main results

  • Consequently, the memorization and generalization errors coincide; in other words, the passage reports the same error for memorization and generalization, without specifying further conditions in the quoted text.Source: §A.5 · p. 23
  • A degree-k harmonic component is fully recovered when L ≫ d^k, partially recovered when L = Θ(d^k), and unrecovered when L ≪ d^k.For polynomial context scaling L = eL = αd^s, the intermediate case occurs only when s is an integer in K.Source: §A.5 · p. 24

B.1 Fixed point equations

  • The feature learner’s state equations determine five order parameters and their conjugates from (αF, γF, κF, ∆, λ, σ⋆, σ0, SF).The order parameters are m, q, q̄, χQ, and χQ̄.Source: §B.1 · p. 25

B.2.1 The replica trick

  • The replica trick computes positive-integer replica moments, continues them near n = 0, and then takes β→∞; the continuation and limit exchanges are formal steps.The calculation first takes d→∞ at fixed n, β, and (α, γ, κ), then extracts the term linear in n.Source: §B.2.1 · p. 26

B.2.2 Teacher score and attention outputs

  • Under Gaussian equivalence, the teacher score and attention outputs are approximated by jointly Gaussian variables matching their leading first two moments.The assumption is that higher-order correlations are negligible in the high-dimensional limit.Source: §B.2.2 · p. 27

B.2.3 Order parameters and replica symmetry

  • The replica calculation assumes task-pool self-averaging and a stationary solution invariant under permutations of replica indices.These assumptions let task-dependent quadratic forms be replaced by task-pool averages and impose a replica-symmetric form on the order parameters.Source: §B.2.3 · pp. 28–29, §B.2.5 · p. 30

B.2.4 Calculation of f1

  • The f1 calculation reduces the prompt average to a few scalar Gaussian variables and uses scalar minimization in the fixed-point equations.For the identity Stage-I readout, the scalar objective is strictly convex and its minimizer is unique.Source: §B.2.4 · p. 30

B.2.5 Calculation of f2

  • Combining the Q and barred-Q terms gives f2 in Eq. (150).Source: §B.2.5 · p. 30

B.2.6 Calculation of f3

  • An auxiliary Gaussian matrix removes the replica coupling in the f3 integral, making replicas conditionally independent; a bounded determinant contribution vanishes as β → ∞.Combining the resulting f3 expression with Eqs. (156), (157), and (165) gives f1, f2, and f3, and hence f = f1 + f2 + f3.Source: §B.2.6 · pp. 31–32

B.2.7 Stationarity

  • The five order-parameter and five conjugate stationarity equations recover Result B.1 when N0/d → κF ∈ (0, ∞) and the empirical spectrum of SF converges almost surely to νκF(ds).The derivation uses the fact that the scalar objective is stationary at v⋆, so differentiation through v⋆ vanishes.Source: §B.2.7 · p. 32

B.3 Derivation of the main results

The derivation obtains attention matrices, Stage-II readouts, and memorization and generalization errors, then analyzes finite-pretraining corrections and context-length limits. For the feature learner, the order-parameter deviation from the infinite-pretraining values is O(d^2/N).

B.3 Derivation of the main results

  • For seen-task and fresh-task test protocols, the derivation uses different overlaps, with fresh tasks averaged over an isotropic direction.These quantities determine the corresponding attention output and inference error.Source: §B.3 · p. 33
  • Equations (168)–(176) determine the learned attention matrix, Stage-II readout, and both inference errors, deriving Result 2.The derivation conditions on the task direction and identifies the optimal readout as a conditional expectation.Source: §B.3 · p. 33
  • For the context-scaling analysis, the derivation first takes the large-pretraining and infinite-task-pool limits, then considers L = eL = αd^s.It sets noise and regularization to zero for the infinite-task-pool limit.Source: §B.3 · pp. 33, 35, 37
  • For the feature learner, order parameters differ from their infinite-pretraining values by O(d^2/N).This yields the feature learner’s finite-pretraining correction.Source: §B.3 · p. 35
  • The derivation represents the Stage-II training channel with Gaussian variables and uses Mehler’s identity to obtain the optimal readout.The resulting attention output is then evaluated at inference context length eL = eαd.Source: §B.3 · p. 36

C.1 Consistency of the Replica method with numerical experiments

The numerical experiments compare replica predictions with order parameters measured from learned attention matrices. Predictions capture the main trends, with visible deviations near the kernel interpolation peak.

C.1 Consistency of the Replica method with numerical experiments

  • Replica predictions capture the main trends of measured order parameters, with visible deviations near the kernel interpolation peak.The comparison uses saved attention matrices; the feature learner and linear attention share the same learned attention.Source: §C.1 · p. 37

C.2 Varying task diversity

Increasing task-pool size narrows the memorization–generalization gap for both learners. The kernel learner trades seen-task accuracy for unseen-task performance, while the feature learner mainly improves on unseen tasks.

C.2 Varying task diversity

  • Increasing task diversity narrows the memorization–generalization gap for both learners; the kernel learner improves on unseen tasks at the cost of seen-task accuracy.For the feature learner, the main effect is improved generalization.Source: §C.2 · p. 38
  • For d = 24, the predicted crossover is N0 ≃ 40, consistent with the crossover between N0 = 24 and 48 in Figure 6.Source: §C.2 · p. 38

C.3 Derivation of Eq. (217)

The derivation of Eq. (217) specializes the analysis to a common teacher and specified limits on pretraining, dimension, task-pool size, and context length. It uses spectral limits to simplify the kernel learner’s expression.

C.3 Derivation of Eq. (217)

  • The derivation takes N to infinity, then d and N0 to infinity at fixed κ = N0/d, and finally α to infinity.It sets noise and regularization to zero and uses σ0(v) = v.Source: §C.3 · p. 38
  • As α tends to infinity, the ratio s/(s + τα) tends to 1 for s > 0 and remains 0 at s = 0, allowing the limit inside the integral.The passage states that this step also applies when κ = 1.Source: §C.3 · p. 39
  • For the kernel learner, the two trace terms combine, and the derivation invokes the limiting spectrum of the relevant matrix.The Marchenko–Pastur spectrum is used in the subsequent limit.Source: §C.3 · pp. 39–40

D A formal statement of Theorem 1

This appendix formalizes the setting and develops supporting derivations for the kernel and feature learners, including polynomial-limit and context-concentration results. It also establishes feature-learner training and inference limits and outlines an inference-error proof using a spherical projection formula and a truncated Gegenbauer expansion.

D A formal statement of Theorem 1

  • In the population setting, task vectors are sampled independently per prompt, and training objectives are population expectations corresponding to γ → ∞ and κ → ∞.The teacher is assumed measurable, independent of dimension and context length, normalized to zero mean and unit variance, with a nonzero first coefficient.Source: D A formal statement of Theorem 1 · pp. 40–41
  • Theorem 2 supposes Assumptions 3 and 4, and defines Aℓ(α, eα) as E_T[Φρα(sα,eαT)Hℓ(T)] for T ∼ N(0, 1).Source: D A formal statement of Theorem 1 · p. 42

E.1 Gegenbauer polynomials

  • For fixed degree k, normalized Gegenbauer polynomials converge coefficientwise to normalized Hermite polynomials, while Qk,d converges uniformly to qk.The appendix uses these limits to establish further fixed-degree convergence properties.Source: §E.1 · p. 43

E.2 Kernel learner

  • The kernel-learner derivation obtains an exact finite-dimension test error and proves its asymptotic result when L/eL = O(1).The proof replaces finite-dimensional Gegenbauer quantities with their asymptotic limits.Source: §E.2 · pp. 45–46

E.3 Feature learner Stage-I

E.4 Concentration of the context

  • Lemma 2 gives a context-statistic moment bound for each fixed order r ≥ 2, for sufficiently large d and every context length n ≥ 1.The proof conditions on the task direction and makes its estimates uniform over that direction.Source: §E.4 · p. 47
  • The context-concentration proof separately controls the context mean, norm, radius, and direction using conditional moment bounds and martingale differences.The argument also uses Markov's inequality to control deviations of the context radius and direction.Source: §E.4 · pp. 48, 50

E.5 Feature learner Stage-II

  • Stage II analyzes the population-optimal scalar readout learned with context length L and evaluated at a possibly different length eL, under Assumptions 3 and 4.The analysis defines the conditional predictor from the test response given the feature learner's context statistic.Source: §E.5 · pp. 51–52

E.5.1 Asymptotical parameters

  • The passage establishes uniform convergence of the asymptotic parameter ρd,n to ραn over the full range αn > 0.Assumption 4 keeps r in a fixed compact subset of (0, ∞), proving the last convergence.Source: §E.5.1 · p. 52

E.5.2 A spherical-to-Gaussian approximation

  • The Stage-II technical lemma approximates the randomly rescaled spherical statistic A_d S_d by a Gaussian variable Z for training and inference distributions.The proof combines a high-dimensional spherical Gaussian limit with the Gegenbauer-to-Hermite limit and allows dimension-dependent test functions of uniform polynomial growth.Source: §E.5.2 · pp. 52–53

E.5.3 The Stage-II regression on the training distribution

  • On the training distribution, the regression function is identified through matched bounds using a richer context-conditioned regression and spherical projection formulas.The proof then uses a Gaussian channel and establishes the stated convergence results for the regression and its error.Source: §E.5.3 · pp. 53, 55

E.5.4 Transfer from training to inference

  • The convergence established under the training law is transferred to the statistics encountered at inference, with the argument covering both context lengths L and eL.The transfer controls the good and bad events and applies the resulting convergence to the feature-learner predictor.Source: §E.5.4 · pp. 55–56

E.5.5 Reduction of the inference error to a Gaussian channel

  • The inference-error reduction bounds a residual term using Cauchy–Schwarz and uniform moment bounds.This begins the reduction of the inference error to a Gaussian channel.Source: §E.5.5 · p. 57
  • The cross term in the inference error is treated using the spherical projection formula from Lemma 1.The passage does not include the formula or its resulting value.Source: §E.5.5 · p. 57
  • The derivation then truncates the Gegenbauer expansion and uses Eq. (261), but the supplied passage omits the resulting expression.The fragment also refers to Lemma 1; it does not provide enough context to state the conclusion of this step.Source: §E.5.5 · p. 57

E.5.6 Continuity of the Gaussian-channel error

The appendix establishes continuity of the Gaussian-channel error to justify replacing finite-dimensional parameters with their limits, then completes the feature-learner part of Theorem 2.

E.5.6 Continuity of the Gaussian-channel error

  • The Gaussian-channel error is continuous in its parameters and uniformly continuous when the scale parameter stays in a compact positive interval.This is the continuity property used in the subsequent limiting argument.Source: §E.5.6 · p. 58
  • Bounds place the finite-dimensional and limiting scale parameters in a compact positive interval, enabling uniform continuity to justify the parameter replacement.The argument concludes by proving Eq. (414).Source: §E.5.6 · pp. 57–59

E.5.7 Final proof of Theorem 2

  • The proof defines a Gaussian-channel function and uses the Bayes-optimal predictor with Mehler’s and Parseval’s identities to establish the feature-learner part of Theorem 2.The appendix states that this completes that part of the theorem.Source: §E.5.7 · pp. 59–60
Loading 2610.01712v1…