Attention as Featurizer: How Transformers Learn Nonlinear Regression In-Context

Paper: 2605.05176 Authors: Alexander Hsu, Zhaiming Shen, Wenjing Liao, Rongjie Lai Categories: cs.LG, math.NA

The Gap

Existing ICL theory has mostly tackled linear regression. We know transformers can fit linear functions from examples in the prompt, and we have bounds proving this works. But real-world tasks are nonlinear — polynomial trends, piecewise patterns, smooth curves. The moment you step outside linearity, prior theory goes silent. It doesn’t tell you whether attention can construct nonlinear features, how many examples you need, or what generalization looks like.

The gap: no constructive proof that transformers can realize nonlinear basis functions through attention, and no finite-sample bounds for nonlinear ICL.

Problem: ICL works empirically for nonlinear tasks
   |
   v
Gap: Theory only covers linear regression
   |
   v
Assumption: Attention can act as feature constructor
   |
   v
Method: Explicitly build transformer layers that realize
        polynomial/spline bases via attention mechanism
   |
   v
Evidence: Finite-sample bounds + synthetic validation
   |
   v
Conclusion: Nonlinear ICL is provably learnable with
            explicit sample complexity

The Increment

One sentence: Before this paper, ICL theory stopped at linear models; after, we have explicit constructions showing attention can build nonlinear features and provable generalization bounds for polynomial and spline regression.

Core Mechanism

The paper constructs a transformer architecture in two stages. First, the attention layers act as a feature constructor: given input-output pairs in the prompt, attention weights compute nonlinear transformations (like x², x³, or spline basis functions) of the query point. This happens through carefully designed key-query-value matrices that encode the basis functions. Second, a feedforward layer acts as a least-squares solver: it takes the constructed features and fits them to the prompt examples, producing predictions.

The construction is explicit. For polynomial features up to degree d, the attention layer computes each monomial x^k by setting keys and queries to encode powers of x. For spline features, attention computes local basis functions centered at knot points. The feedforward layer then solves the normal equations to find coefficients that minimize squared error on the prompt examples.

Input prompt: (x1,y1), (x2,y2), ..., (xn,yn), query xq
   |
   v
[Attention Layer 1: Feature Constructor]
   Keys encode basis functions: 1, x, x^2, ..., x^d
   Query xq -> compute phi(xq) = [1, xq, xq^2, ..., xq^d]
   |
   v
Feature vector: phi(xq) in R^d
   |
   v
[Attention Layer 2: Coefficient Solver]
   Compute Gram matrix from prompt examples
   Solve normal equations: beta = (Phi^T Phi)^-1 Phi^T y
   |
   v
Prediction: y_pred = phi(xq)^T beta

Think of it like a custom calculator being assembled on the fly. You hand the transformer a set of example calculations (the prompt). The attention mechanism looks at these examples and figures out what kind of math operations are needed — it’s like selecting which buttons to activate on the calculator (polynomial buttons? spline buttons?). Once the right buttons are identified, the feedforward layer presses them in the right sequence to solve your new problem (the query). The calculator doesn’t come pre-programmed for your specific function; it builds the right program by examining your examples, then executes it.

Key Concepts

• Attention as Featurizer: Normally we think of attention as “looking at relevant parts of the input.” Here, attention does something more structural: it computes basis functions. Given a query point x, attention weights and value projections combine to output φ(x) = [1, x, x², …]. This isn’t retrieval or weighting — it’s function evaluation. The key insight: by choosing key/query/value matrices carefully, you can make attention compute any polynomial or spline basis. It’s like programming attention to be a function generator rather than a content filter.

• Finite-Sample Generalization Bounds: The paper proves that if you train a transformer on T tasks (each task is a function from some family), and each task has n prompt examples, then the test error on a new task with m prompt examples scales as O(1/√m + 1/√T). This tells you exactly how many examples you need in the prompt (m) and how many training tasks (T) to achieve a target error. It’s not asymptotic hand-waving — it’s a concrete recipe: “Want error ε? Use this many examples.”

Framework Shift

Before (linear ICL theory):          After (this paper):

Prompt -> [Black Box] -> Prediction  Prompt -> [Attention: Build Features]
                                              |
Theory: Linear regression only               v
        Implicit feature learning      Feature Vector phi(x)
                                              |
                                              v
                                       [FFN: Solve Least Squares]
                                              |
                                              v
                                       Prediction
                                       
                                       Theory: Explicit construction
                                               Nonlinear bases
                                               Finite-sample bounds

From implicit learning to explicit construction, the core shift is treating attention as a programmable feature generator with provable properties.

Expert Assessment

Problem choice: Real gap. Linear ICL theory was a natural starting point, but the jump to nonlinear is where practical applications live. The authors didn’t manufacture this — it’s the obvious next frontier.

Method maturity: The construction is clever but not deep. It’s essentially showing that transformers can implement classical regression with basis functions. The novelty is in the explicit construction and the finite-sample analysis, not in discovering a new learning principle. That said, explicit constructions are valuable — they turn “transformers probably can do this” into “here’s exactly how.”

Experimental integrity: Synthetic tasks only. The paper validates on polynomial and spline regression with controlled noise. No real-world benchmarks, no comparison to other ICL methods, no ablations on architecture choices. The experiments confirm the theory works as stated, but they don’t stress-test it. For a theory paper, this is acceptable but leaves practical questions open.

Writing quality: The construction section (Section 3) is dense and could use more intuition before diving into matrix notation. The paper front-loads formalism when a simple example would clarify. Section 4 (theory) is solid but assumes familiarity with concentration inequalities. A rewrite of Section 3 with a running example (e.g., “let’s build a transformer for quadratic regression step-by-step”) would make the paper far more accessible.

Verdict: weak accept — Solid theoretical contribution with explicit constructions and finite-sample bounds, but limited experimental scope and dense presentation reduce impact.

Takeaways

For practitioners: If you’re building ICL systems for regression tasks, this paper tells you that attention can explicitly construct nonlinear features — you don’t need to hope the model learns them implicitly. You can initialize attention layers to compute specific basis functions (polynomials, splines, wavelets) and let the model refine from there. This could speed up training or improve sample efficiency.

For theorists: The technique of explicitly constructing transformer layers to realize specific functions is generalizable. If you want to prove transformers can do X, try constructing the weights that make them do X, then analyze the construction. This paper’s framework (feature construction + least-squares solving) could extend to other function classes beyond polynomials and splines.

For skeptics: The bounds depend on the function class being known (you need to know you’re fitting polynomials of degree d). In practice, you don’t know the true function class. The paper doesn’t address model selection or misspecification — what happens when you use polynomial features but the true function is exponential? That’s the gap between this theory and deployment.