On Understanding Attention-Based In-Context Learning for Categorical Data
Authors: Aaron T Wang
William Convertino
Xiang Cheng
Ricardo Henao
Lawrence Carin
Venue: ICML 2025 poster (2025)
Keywords: attention networks
in-context learning
Transformers
Primary area: deep_learning->attention_mechanisms
TL;DR: Examination of how multi-step functional gradient descent may be performed exactly with appropriately designed self-attention and criss attention layers, plus skip connections, with a demonstration on two real datasets.

Abstract: In-context learning based on attention models is examined for data with categorical outcomes, with inference in such models viewed from the perspective of functional gradient descent (GD). We develop a network composed of attention blocks, with each block employing a self-attention layer followed by a cross-attention layer, with associated skip connections. This model can exactly perform multi-step functional GD inference for in-context inference with categorical observations. We perform a theoretical analysis of this setup, generalizing many prior assumptions in this line of work, including the class of attention mechanisms for which it is appropriate. We demonstrate the framework empirically on synthetic data, image classification and language generation.

On Understanding Attention-Based In-Context Learning for Categorical Data
Aaron T. Wang'
William Convertino'
Xiang Cheng' Ricardo Henao
Lawrence Carin'
Abstract
predicts the corresponding output. Early work in this di- rection concentrated simple function classes, like linear regression (von Oswald et al., 2023; Ahn et al., 2023; Zhang et al., 2024; Schlag et al., 2021), with extensions to nonlin- ear, kernel-based regression (Cheng et al., 2024) emerging more recently. These studies have successfully linked the in- ference performed by the Transformer to functional gradient descent (GD) on the latent function.
In-context learning based on attention models is examined for data with categorical outcomes, with inference in such models viewed from the perspective of functional gradient descent (GD). We develop a network composed attention blocks, with each block employing a self-attention layer followed by cross-attention layer, with as- sociated skip connections. This model can exactly perform multi-step functional GD inference for in-context inference with categorical observations. We perform a theoretical analysis of this setup, generalizing many prior assumptions in this line of work, including the class of attention mecha- nisms for which it appropriate. We demonstrate the framework empirically on synthetic data, im- age classification and language generation.
These analyses have primarily focused on functions with real-valued outputs, significant departure from the categor- ical nature of language data, where outcomes are discrete tokens. This paper bridges this gap, extending the functional GD framework to handle categorical observations. This ex- tension is crucial for bringing the analysis closer to the complexities of language models. Furthermore, while prior work has largely relied on simulated data, we demonstrate the applicability of our framework to real-world problems, including in-context image classification and, importantly, language modeling. For the latter, we leverage recent ad- vances in targeted datasets designed to illuminate Trans- former mechanisms (Eldan & Li, 2023), connecting this line of research with the functional data analysis perspective.
1. Introduction
The Transformer architecture (Vaswani al., 2017) has revolutionized natural language processing, achieving markable success in tasks like language generation (Vaswani et al., 2017; Radford et al., 2019; Devlin et al., 2019; Brown et al., 2020; Touvron et al., 2023; DeepSeek, 2025). This success has spurred significant interest in understanding the underlying mechanisms driving Transformer inference, particularly for in-context learning, where the model is pre- sented with "prompt" of input-output examples and tasked with predicting outputs for new inputs (Liu et al., 2022).
1.1. Summary of Contributions
1. We analyze in-context learning with categorical observa- tions, offering new insights into language model behavior. Specifically:
(a) We demonstrate that encoding categorical data (e.g., tokens) as learned embedding vectors arises naturally from a GD perspective on inference (Section 2).
(b) We introduce a new attention- based architecture, em- ploying self-attention and cross-attention blocks, that enables exact implementation of multi-step GD with categorical data (Section 3).
A fruitful line of research has emerged, focusing on sim- plified yet insightful paradigm: predicting outputs for a context-dependent hidden function. In this setup, the prompt consists of onput-output pairs from this function, and the Transformer is queried with a new input. Crucially, the Transformer is not explicitly trained to learn this function; instead, it infers the function from the provided context and
(c) We show that next-word weext-wor-geeerattion generation can be modelled as a in-context learning problem over categorical ob- servations (see Section 3.6).
(d) We provide theoretical analysis showing that, for in- context learning with categorical observations, our GD-based representation is a stationary point of autention-based inference models (Section 5). This theory addresses softmax- -baed attention, widely used in language models.
Ellettrical & Computer Engineering Dept., Duke Univer- sity, Durham, NC, USA. Correspondence to: Lawrence Carin <Iearin @duk..edu>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
2. We empirically validate our framework through experi-

---

On Understanding Attention-Bant Based In-Context Learning for Categorical Data
ments on diverse datasets:
2024). These studies primarily consider real-valued outputs, while our work extends this framework the more chal- lenging case of categorical outputs, while also generalizing the form of attention that may be considered.
(a) We tackle in-context image classification on ImageNet (Russakovsky al., 2014), demonstrating the model's ability to handle high-dimensional data and adapt to unseen image categories (Section 6.2).
2. Model Setup
(b) We apply our GD-based model to language generation, training on combined corpus of Tiny Stories and Children Stories (Eldan & Li, 2023). We perform a quantitative analysis using GPT-4o scoring (Section 6.3), and demonstrate that the GD-based model yields similar performance as a comparable Transformer, with far fewer parameters.
2.1. Softmax model for token probabilities
Transformer-basee in-context learning with categorical data may be viewed from the perspective of two models: (i) the Transformer (z(2) with parameters 0 and input 2, and (fi) the model:
exp (w (w(()(())
Po(Y - y x-X=a) = x)- E exp (w (wE,Fo(x))
(I)
Our principal motivation is to gain insight into how the Transformer performs in-context learning when the obser- vations are categorical. It is hoped that such insights are relevant to mechanisms involved in the success of Trans- formers for language modeling. Additionally, the model we have developed may be of interest in its own right, for applications in fho-shot learning with categorical outcomes.
softmax representation for the probability of category E (1,..., for covariates x E ERd, where the context- dependent function fa(x) characterizes covariate depen- dence. This paper seeks to examine how a gradient-descent (GD) perspective on in-context inference of fz(z) may be implemented in the forward pass of Transformer (=).. To be explicit, we will primarily represent (=))) in terms of attention-based layers with skip connections.
1.2. Related Work
The broader context of our research is few-shot learning, where models are trained to generalize from limited ex- amples (Schmidhuber, 1987; Bengio et al., 1992). In this context, meta-learning aims learn "learning algorithm" that can quickly adapt to new tasks (Andrychowicz et al., 2016; Ravi & Larochelle, 2017: Finn et al., 2017). The Transformer, with its ability to perform in-context learning without parameter updates, can be viewed as powerful meta-learner (Brown et al., 2020).
The models T and p (Y are distinct, but they have close interrelationships. Specifically, they share parame- ters (We,e)e=1.C. where each We.c E Rdd" As developed below, the input encoding to [z) employs (we,eJe=i.c Further, in its forward pass (2) computes fa( based on contextual data encoded in and it then feeds the inferred fa( into P = y| = x) to predict the token/category of query. Therefore, the softmax representation in (1) composes the output element of Tz(2..
growing body of work specifically investigates the in- context learning capabilities of Transformers. One line of inquiry focuses on characterizing the types of functions that Transformers can effectively learn from context, ranging from simple linear functions (von Oswald et al., 2023; Ahn et al., 2023; Zhang et al., 2024; Schlag et al., 2021) to more complex nonlinear functions (Cheng et al., 2024) and even Bayesian models (Muller et al., 2022; Garg et al., 2022).
2.2. Functional GD for latent function input to softmax The input to Tz(2) is encoded representation of the context C = {(2,,, 31) li=i,N and the query IN+1-- We seek to understand (z) from the perspective of functional GD. In that context, function (x( is assumed to reside in func- tional class F, and following (Cheng et al., 2024) we assume F is a reproducing kernel Hilbert space (RKHS) (Schölkopf & Smola, 2002). Specifically, consider / (x:) Avb(2r) + b, with vi(x) a fixed mapping of covariates
Our work is closely aligned with research that interprets Transformer inference form of functional gradient de- scent. This perspective has been successfully applied to analyze in in-context learning in settings where the target func- tion is linear (von Oswald et al., 2023; Ahn et al., 2023) or resides in a reproducing kernel Hilbert space (RKHS) (Cheng et al., 2024).
a Hilbert space, and the parameters acting in that space are ø (A, b). If the dimension of the vector output from (xxx is m (which could be infinite, in principle), then A E Rdd xm while E IR'., As discussed below, while GD is performed on (A, the update of these parameters not performed ex- plicitly, rather the function for (x) is estimated directly, where Φr represent the parameters at GD step k.
The connection between attention mechanisms and kernels has been explored in various contexts. Linear attention can be seen special case of kernel-based attention, where the kernel is simply the inner product (Schlag et al., 2021; von Oswald et al., 2023; Ahn et al., 2023; 2024). More gen- eral ceenel-based attention allows for richer representations and has been investigated for real- valued data (Cheng et al.,
The matrix of embedding vectors W. E TRd'xC is learned, where column c, We,e E Rdd represents the embedding vector for category c E (1...., C} (details in Section 4). When performing inference of fa(z) for new contextual data (after the attention model is trained), W. is fixed.
2

---

On Understanding Attention- Based In-Context Learning for Categorical Data
The cross-entropy loss we wish minimize inference, to infer Ø- (A,b) based on context C- {(1,, i=,,, is
nonlinear B(we)) This motivates the new architecture
developed in Section 3, in which self attention is comple- mented by a form of cross attention.
N
exp
w c.gy (z,) b)]
(()) = N
log
exp W (T (z, + 6)]
2.4. Natural role of learned token embedding vectors If the matrix of paramaters A is initialized with all-zeros (and the bias b is all-zeros), then the function at each po- sition is initialized as fi,u = Oa,, where O2r is an all-zeros d' d'-dimensionl vector. With this initiation, E(we) ),,, =
As derived in Appendix B, for GD step k- + the following functional update rule is realized for for-(
fa = f3 + Afj.
N
Afjk* = a N
Σ
We.y. - E(wedlenn) k(())
(3)
2 EeE We,e, i.e., the average of the C' learned embedding vectors (subsequently denoted W.).
where ((wre)) A. Ec Po. (V =dx- mz). The kernel is represented as k( x3) (x)z)) which is endowed with a reproducing property (Wainwright, 2019). For simplicity, in (3) and in subsequent discussion we ignore the impact of the bias (the details of its inclusion are provided in Appendix B). The parameters are initialized at step k O.
Hence, from (3), the first functional update مʃ3.0 a wernee-weighted average of inputs {[we,s. k=..... If uh = Oar, this corresponds exactly to the widely used input encoding of tokens/categories in terms of the associated em- bedding vectors (Vaswani et al., 2017). This demonstrates that the ubiquitous means of encoding language tokens, in terms of learned embedding vectors, is consistent with GD perspective on the attention -based inference of fø(x)).
As developed further below, the kernel xT) in (3) man- ifests an attention weight between positions i and j. The requirement that k(zz, xt)) = (xx))² (aa;) ) for some vb(x) places restrictions on the form of 25) (e.g., symme- try), which are inconsistent with softmax-based attention in most Transformer implementations. In Section we extend the above formalism, also allow softmax-based attention. We also show there that the GD-based attention model (de- veloped in Section 3) is a stationary point of a more general construction.
3. A etention-Based Exact Multi-Step GD
(1x)))((())
3.1. Data encoding
The proposed architecture employs attention blocks, with each block composed of a self-attention layer followed by a cross-attention layer. Let 6i represent the vector input to attention block ks + 1. As depicted in Figure 1, for (3) we consider
(4) Within e,A,, the positions of B(we)ls, and J may be viewed as "scratch space" where these functions are sequen- tially updated, to implement GD. The importance of scratch space for Transformer implementation of GD was first noted in Akyurek et al. (2022) (see Sec. C.4.1 in the Appendix of that paper). Like in prior attention- implementations of GD (von Oswald et al., 2023; Ahn et al., 2023; Mahankali et al., 2024; Cheng et al., 2024), the latent function is up- dated at each self-attention layer, as Si. = Σ k-1 k==0 Afi, for k 2 1, with A defined in (3), and fi, = Oz Our novelty concerns development of a cross-attention mech- anism for updating the nonlinear expectation B(wx)) within each attention block.
C1,k = (I4, We.ge e B(we)lra fT,
2.3. Connections to prior GD Transformer research In most prior research motivated on understanding the fun- damentals of the Transformer based on a functional GD perspective (von Oswald et al., 2023; Ahn et al., 2023; Ma- hankali et al., 2024; Cheng et al., 2024), the observations y: were real-valued vectors. In such setup p(Y = yilX- = =N) o ). where the variance 6² is assumed constant (and not estimated) and La the d' x d' identity matrix (here the observations y/ Rdd ). As detailed in Ap- pendix G. by assuming a model fx(x) = (x)(2)) + within ((fo(z)) a Iw GD yields an update equation identical to (3), with the change We,y. - B(w=)),,, + y - B(YY) fe." Because of the underlying Gaussian model for Y, with mean fa(x)) i.e., B(Y) n. fi, This property allows multiple steps of GD with real vector observations be performed exactly with multiple steps of kernel-based self attention (plus skip connections) (Cheng et al., 2024), as the underly- ing expectation E E(Y)) ( is just the function itself.
Within each attention block, the self-attention layer is com- posed of two attention heads (left and center in Figure I), and the cross-attention layer by single attention head (right in Figure 1). The next two subsections summarize each of these layers, and Appendix B provides details of all Trans- former parameters.
For categorical observations of interest here, B(wa)in, is a nonlinear function of fi,., and the embedding vector We.y. encodes observed yi. Consequently, multiple steps of GD in this setting cannor be performed in terms of multiple self- attention layers alone; a mechanism is needed to address the
3.2. Self-attention layer 1. SELF ATTENTION FOR fi,k + fi.k + Afi.k
For this self-attention head, consider key and query ma- trices Wf = Wq for which W WKeLo = (r,, Osa )" and
3

---

On Understanding Attention-Ban Based In-Context Learning for Categorical Data
Self-Attention for f - fax + Afa
Attention for Erasing (we)
Cross Attention for E(W,) (Ww)) fike1
Attention for Kaya
Attention Used forkeys Kays 6 4.Quansee Via My and
andliig
Wess
MAy.
Values Sclected by
E E0
Hey Si
So actedd
Vislding RCMe)lo
ARA2
Bloddl Ad
Zeros Vector
W, softmas(W)) fraws)
fe Mix
Aix =
May,
EcHuRi.. ((6.)))
fiz
fiz+:
Attention
%
Figure 1. The left and middle figures depict the two self-attention heads, while the right-most figure reflects cross attention. A representative vector shown at position i. The portion the vector blue remains the same each layer kt, and the portion in yellow updated each attention block. The sum of ff. and A Fi. at left and center is manifested via the skip connection. The form of the vector shown at right, for input to cross attention, depicts the updated fi,441. and the crased expectation, these the result of the preceding self attention. and places it at its location within the vector (addition via skip connection).
Cross attention (right) computes the expectation E ( 12 One complete attention block, with skip connections, is depicted in Figure 2.
Woe30- (x,, 03d) )", respectively, and from which kernel- based attention yields x(Wker,o, Wges,o) k(x,, 2y). These the types of key and query matrices considered previously for this line of research for real observations (von Oswald et al., 2023; Ahn et al., 2023; Mahankali et al., 2024; Cheng et al., 2024). The keys are applied to inputs i = I.....N
(A., Wasy E(w E.N cloes)
(2, )i Di Rr)
Laper
Embhelding
Vactora, W.
N and the queries to j 1,.... + The value matrix Wv designed to achieve Wyee, (02+2d', Weyy ((.w))r,.)"
Self Layeniinnni
Using WK. Wo and Wy on inputs {er,oli=I,N. for queries j ....., N + 1, attention yields Wv Ei (3,, The N + output vectors tied to the queries are mul- tiplied by a projection matrix P E R(d+4dd)x PWv E (r,om((s :ry), where here P is just the identity matrix. The resulting vector from this head is
(s....,. ( fo)
Figure 2. Summary of interleaved attention for multi-step GD with categorical observations.
the Kronecker delta function. Further, Wv is set such that Wyei, = (0.4a's -E(we) Ox)) As shown at center in Figure 1, this attention head yields
#f1 i
(0.+2d's
Weey, ((we)l(wW In(In22)) -,,)),,,, N
where H1 denotes ceff-attention-head one (left in Figure 1). Note that this computes Afi,, as defined in (3).
H2 (0.4d' -E(we) Oz)²² i.k
3.2.2. SELF ATTENTION FOR ERASING B(we)la,
Combining the two attention heads, a + e and then adding this sum to the skip connection, yields the output from the self-attention layer reflected in Figure 2.
The second self-attention head is used erase E(we))
This erasure is necessary preparation for the subsequent cross-attention layer, where the expectation E[w.),,.. f..4 is computed, and placed in the location of the erasure.
3.3. Cross-attention layer for computing (Ewe)) Within attention block > 1, the output of the self- attention plus skip connection at position represented as
For this self-attention head, We and Wx are designed such that We = WK and WKeiA = A(x;, 0za))" Yr with > a large real number. With large and us- ing a softmax or radial basis function attention kernel, x k(Woez,, W n()xy, Axxe) as dez, where &.g is
ē (I; We.yr , Oz FT.
(5)
We now demonstrate that the expectation

---

On Understanding Attention- Basse In-Context Learning for Categorical Data
C
explwe.efi.t
one is only interested in a single GD step, we may simplify the design summarized in Figure 2. Since only one GD step is computed, there is no need to update the expectation B(we))lr.. = we + E(w=)),,, Consequently, within the self-attention layer we remove the expectation-erasuur atten- tion head. Further, since the expectation is not updated, the cross tross-attention layer is removed entirely. There is only one self-attention layer, using a single attention head. Further, there is no need to isolate the expectation within the encod- ing, as it not updated. Therefore, for implementation of a single step of GD, we consider simplified input
B(wellnn
Σ
We.cc
(6)
=l==1 explua..
We softmax fi,))
may be framed as noltmax-basee attention, with query fik- keys (we,e)e=1.c and values (We,el==1.C- We leverage this fact to update the expectation via softmax-bssed cross- attention layer.
Vectors (E,,) are used to constitute the queries associated with the cross-attention layer, with W = fi,411 (this Wq and related matrices are connected to cross- attention, and are distinct from those discussed above in the context of self attention). The cross-altention matrices Wx and Wv are applied to the C' columns of We, corresponding to the (learned) embedding vectors for the categories/tokens. Specifically, Wx = Wv are both the identity matrix, such that WK W = W.. Using softmax- attention, one exactly implements (6). The result of cross attention is added to the skip connection, with the updated expectation placed in the position of the aforementioned erasure.
6i = (2is, We.y -ee Oar)²
(9)
Details of the parameters for this single-step-GD self- attention layer are presented in Appendix C.
3.6. Extensions for modeling language
When presenting results for language modeling from the perspective of functional GD (in Section 6), we will em- ploy positional embedding vectors (Vaswani et al., 2017) as covariates Ty-- In that connection, concerning the func- tion update in Section 3.2.1, for language modeling we will consider multiple self-attention heads for the function up- date. This may be viewed as extending the discussion of M Section 2.2 to fa(x) = Σm=1 Am Am (1r), where here we assume M attention heads related to the function fe(xr), with distinct (z) for each. As discussed in Appendix H, in practice we consider s(Ammii, Amoz)) where An is head deaeddeeendent diagonal matrix (and k(.,-) is the same for all M heads). The same form of Wy as discussed in Section 3.2.1 is used for each of these heads.
At the output of the last attention block (Kth, for K GD steps), the vector at position i is
Ci,K- = (Ii We.ye E[ B(we)lss fixx)²
(8)
A d' x (d + 3d') matrix R is used after the last attention block, yielding Ji,k = ReiK.. At position i = N ++1 1, this is then input into the softmax in (I), to predict the probability that UN+1 c, for all c c {I....,
3.4. Connection to original language Transformer
In this setup. the erasure self-attention head and the cross- attention design remain unchanged. The multiple kernels for representation of fo (z) also connect to the multi-headed attention used in language models (Vaswani et al., 2017).
The original Transformer design (Vaswani et al., 2017) con- sidered a form of interleaved self-attention layers and cross- attention layers (with skip connections) in its decoder. This formed the same type of attention blocks we developed above. We have provided some perspective on such a for- malism through consideration of inference in terms of func- tional GD. The difference with that prior work is that here the cross attention is performed with the embedding vectors (columns of We), where in Vaswani et al. (2017) the cross attention was performed the encoded input text (prompt).
4. On Learning and Inference
For learning the parameters of the attention based model, it is assumed that we have access to L contextual datasets, ie., for 1 = 1,..., L, c(2) = ((x))) .M,
(1) (0) ()))=,, When
training, we seek to maximize the expected log-likelihood of (Y - each contextual dataset, it is assumed there is underlying fete (z) E F which is to be inferred approximately, based on the associated N contextual examples.
(1)
x (W+z), averaging across 1. For
From another perspective, the expectation as expressed in (7) is evocative of the multi-layered perceptron (MLP) layer employed in current Transformer-base (employing only a decoder arm (Brown et al., 2020), and no cross-attention). However, the softmax function is not a pointwise nonlinearity, as typically employed in MLPs.
language models
The model parameters include We. WK. Wy, and P for each attention head at each layer, We, and R at the Trans- former output, prior to the softmax model in (I). We follow terminology from von Oswald et al. (2023): a Trained TF is Transformer for which all model parameters are learned from a random initialization, without any assumptions or restrictions; and a GD Transformer is one for which we specify the form of the parameters to be consistent with the above GD-baaedd design. When performing learning for GD,
3.5. Simplification for a single step of GD
As demonstrated when presenting experiments in Section 6.3, Transformer inference based on a single attention block (corresponding to one step of GD) often yields results al- most as good as a multi-layered (multi-step GD) model. If
5

---

On Understanding Attention-Bann Based In-Context Learning for Categorical Data
since the form of the parameters is specified, the number of attention parameters is markedly reduced, and We is also learned.
regression setting, where the observations y; are equal to the target function f(x;) (up to a simple Gaussian perturbation). Their (Cheng et al., 2024) proofs do not generalize to the categorical setting, where the observations y; are sampled conditional on f(x) as in (1). Importantly, J(()) is a multi- dimensional unobserved latent function. The categorical setting is crucial for modelling language.
5. Theoretical Analysis
We analyze the loss-landscape of attention model weights for in-context-learning of categorical data. Consider an L- layer attention model, and let JN+ E Id denote our estimate of (xxN+1) based on the model's output layer L. See Appendix I for details on how fN+1,L is defined as a function of the model's output.
3. Importance of Softmax: enables the Transformer to compute E(ws)),,. as detailed in Section 3.3, which is crucial for implementing GD for categorical data, and plays a key role in the proof of The- orem 1. Softmax Softmax- activated self-attention generalizes the kernel-based attention considered in Cheng et al. (2024).
Softmax in cross-attention
Consider the population cross-entropy loss, defined as T WENNII [N+1.L]
exp
[())- = B.
log
6. Experiments
>e co=1 1 exp [W" [N+1LL]
All experiments were performed on a Tesla V100 PCIe 16 GB GPU. Code needed to replicate our experi- ments is at https::/githon vel_attention ion wwweeeetoonnnnooorrreee categgriia....
where 0 denotes attention model weights (e.g., Wo. WK K and Wv at each layer) and shows up on the RHS through fN+1.Z. Σ is the loss from our prediction of the query to- ken (1)+1), in expectation over possible realizations of the context (x,, 3))==L..NN+11 In the limit of infinite training data, the model weights 8, trained by gradient- based op- timization algorithm. will converge to a stationary point of L. In practice, discussed in Section 4, the expectation ((..))))..... associated with C()) is estimated with L contextual training sets.
comm/arrntwang:
6.1. Synthetic data
We consider synthetic data generated by the model p(Y cz)())) exp w fx)]=" /e explux (x)), for C 25 and w. R The data are designed such that f(x) is highly nonlinear and changes for each set of contextual data. Details on data generation are provided in Appendix E.
Theorem (Informal Statement of Theorem 2):
We compare test prediction performance of attention-basee models, based on Trained TF and GD. In all of these experi- ments, N = 10 data pairs C(.) = {xx : 0/
Assume that (1) the self-attention only involves the the feature vectors xe 's (2) the covariate vectors n,'s 's have distribution that is rotationally invariant (3) the cross-attention layer as described in Section 3.3 exactly computes E E(we)),.. For i = 1...N + 1, let fi, denote the readout of the model's prediction for (x;) at layer k. Then there exists a choice of attention-mode parameters 0, which exactly implement gradient descent, ie. fi,k evolves according to (2)-(3). Fur- thermore, is a stationary point of Z(0)-
(1) (e), define (1)
the context, with an associated query pair (x N4 1) (UN+1 is known when training, not when testing). We train the model with contexts c()) / 1,.,,,...44 L. and various training sizes L. are considered. Testing performance is per- formed using a distinct set of contextual data and query, with data generated as above, with performance averaged over contexts C ((L+m), for test sets m 1,..., M (M 2048), (L+m) each with an associated query x N
We present the formal version of Theorem I as Theorem 2 in Appendix I. We highlight aspects of Theorem 1 below:
For a single self-attention layer, there is close agreement between GD and Trained TF (Figure 3). From the theory of Section 5, we expect the Trained TF and GD to yield similar results. Considering the average Top-1 model pre- dictive accuracy and the average negative log-likelihood (on the test data), we observe close agreement between Trained TF and GD after a sufficiently large number L of in-context training sets. The GD-based self-attention model requires much smaller L to train well. We also examined linear atten- tion, for which K(T;,,)) 2,) = x. 24; as expected, it performed poorly because the underlying fo (z) is highly nonlinear.
1. Convergence to GD Implementation: As the number of training samples increases, the empirical loss over the training set approaches the population loss L. Consequently, Theorem 1 suggests that the K llyer model, after training converges, will implement K steps of functional GD. The- orem 1 is supported by Figure 3, which suggests Trained TF converges accuracy to GD with more training sam- ples. Additionally, Figure 11 in Appendix F shows that the attention parameter matrices training convergence do coincide with the stationary point parameters constructed in Appendix I. Figure 5(left) shows a similar match between GD and Trained TF for two layers.
Figure 3 shows close agreement in the predictive perfor- mance of the Trained TF and GD, but that doesn't 't necessar- ily mean that the underlying parameters learned by Trained
2. Challenges of Categorical Data: Previous analysis of the Transformer optimization landscape have focused on the
6

---

On Understanding Attention-Bal Based In-Context Learning for Categorical Data
04
Meggiie Log chudkeue
43 Clandey
0.0
62-
03
Figure 3. For one self-attention layer, comparison of accuracy and loss of GD and Trained TF with RBF (left) and softmax (right) attention, as a function of contextual training sets. Error bars are computed from five different random initializations.
Ge
Tha sing Spects
Figure 5. (Left) For two attention blocks for the model in Sec- tion 3, comparison of accuracy and loss of GD and Trained TF with softmax attention, as a function of contextual training sets, (Right) Model accuracy for the ImageNet dataset, for which the ditention-based model was trained once for 500 epochs, and the linear probing was trained for 500 epochs on cach test context.
Pre-Traned
Contextual Images (efew shots")
egg, CNWWT Fastues Fxtmerrr
X
Ape
New Classes
300
37:m
2003
1000
Figure 4. For the multi-layered model in Section 3, training curves of accuracy (left) and negative logllikeldhood (right) for the GD transformer with softmax attention. Error hars are computed from five different random initializations.
Quaryy
Mey...
Enceded Tranetorme lopart ((oueoollss
Functional GD Tranofomer Pare
TF agree with the GD theory. In Appendix F, we consider another synthetic-data example of lower dimension (for vi- sualization), and there we show close agreement between the Trained TF learned parameters and those of the GD framework, as predicted by the analysis in Section 5.
p( - efContertt spsi) L...5
Figure 6. Illustration of in-context learning with image data. Each image is mapped to a feature vector Jy Contextual data are com- posed of of images for new classes (here five), and query image that, in the Transformer forward pass, is assigned a proba- bility of being associated with each of the contextual labels.
Additional attention layers improve GD- pre- dictive accuracy. In Figure 4 we consider the same data as discussed above, and consider the GD-based attention model for up to six attention blocks (corresponding to up to six GD steps). In these experiments, we trained with L 2048 contextual blocks, and performance is shown on the test data as function of number of training epochs. We note that single attention layer achieves good Top-1 pre- dictive accuracy, but noticeable improvement in predictive accuracy is observed by adding a second attention block. The improvement in the (test data) negative Pog-likelihood more substantial with increasing attention- block layers.
such class 10 specific images are selected at random, and therefore N = 50 (image N + is selected at random from the 5 class types considered in the context data). When training L 2048, and test performance is averaged for M = 2048. The covariates z; (with d = 512) correspond to features from the VGG network (Simonyan & Zisserman, 2015). We set d' = 4 and learn [w:e)==.cc See Figure 6 for a depiction of ICL for image classification.
For in-context ImageNet elassification, GD-based at- tention yields Top-1 predictive accuracy comparable to "linear probing" (Figure 5 (right)). Importantly, with lin- ear probing a model must be learned anew for each test contextual block, to be contrasted with the attention- -aasd model, for which no fine-tuning of parameters are done after training. We observe that with one layer of GD with self- attention, the GD attention model achieves slightly lower accuracy and higher loss compared to the linear probing model. However, with two and three attention blocks, the performance of the GD transformer aligns almost exactly with the linear probing model. Note that with the VGG model for feature extraction, and our in-context model for classification, new image- contextual data (with la- bels never seen before) can be classified with no parameter
With sufficient training data for Trained TE, there is close agreement between Trained TF and GD for two attention blocks (Figure 5(eff)). However, note that large number of contextual data (e.g., L > 25, 000) are needed to train this multi-layered Trained TF model, while far less training data (e.g., L < 5000) are needed for the GD model. Comparing with the single-layer model in Figure 3, note that far more training data is needed for Trained TF when going from one-layer model to a two-layer model.
6.2. In-Context Image Classification
Using the ImageNet dataset (Russakovsky et al., 2014), we select 900 classes for Transformer training, and a separate 100 classes for testing. For each contextual set c()), 5 dis- tinct classes are selected uniformly at random, and for each
7

---

On Understanding Attention-Ban BBssd In-Context Learning for Categorical Data
Tokeooobbeeiiieeeeeee onteciond with oothoisseeeee seere
values only correspond indices i = 1, ,,..,, N for which we have tokens, and therefore the values are set to We,ger For predicting the token at position N + 1, PN+1 is used for the query. Hence, the covariates z; correspond to Pt, and attention is based entirely on token position. This is an important departure from the Transformer model, which performs attention using Wr.y, + pi for its queries, keys, and values, and makes a prediction based on the output at position N. Consistent with GD only using We,y. for the values, we note that many modern language models, such as LLaMa, also do not incorporate positional embeddings in their value vectors (Touvron et al., 2023).
Dioureed tategores (sRars) 7/ to
Eusercte Test WI E1
Garmee (fou Mexiiaa
Query fear
Poatioral pmbreed ng mactam anscearoat with cheerved takars, and be preditest rokin your.
Droode lopde Tharakereer (fuMocyloss replennnnt Funcriona df Touushemer Freword Rasss
play
p0 in Farteet
Figure 7. Summary of how language modeling is performed with ICL. Token y; is mapped to a learned embedding vector We:pr and positional embeddings are used as covariates z;. The input text constitutes the sequence of N tokens, and the next token is connected to the query position TN+11
Another key difference between the GD and Transformer is that the latter effectively does not initialize the underlying function as the zero vector (there is no scratch space with zeros), where the GD-based initializes the latent function fi, = Oar- Despite these limitations of GD, we show below it performs relatively well compared to the Transformer.
adjustment (fine tuning).
6.3. Language Modeling
In testing both the GD-based model and the Transformer, the model is tasked with generating the ending to a previously unseen story. To demonstrate the range in performance of each model, we show examples of both a "good" gener- ated ending and "bad" generated ending (for both models, about 50% of the examples were of each type). We have done extensive experimentation; many other examples are presented in Appendix K.
We examine language generation in the context of GD-based attention models. We make comparisons text generated by a Transformer architecture, following the design of (Rad- ford et al., 2019). For these experiments, we train on a combination of the Tiny Stories (Eldan & Li, 2023) and Children Stories!
We have performed experiments for multi-layered models. However, it has been shown that for such simple stories a single-layer model often generates good text (Eldan & Li, 2023) (and the experiments from the previous section on cat- egorical classification showed that one-layer model yields good "first-order" category/token prediction). For simplic- ity. in the aunuuge--ennratioon experiments we consider a single-layer GD model and a single-layer Transformer.
The generated text from the GD-based language model are of similar qualitative character as a comparable Transformer-
model. In Table 1 we present GD generated text based on softmax attention, and in Table 2 are shown results from the Transformer (we choose the softmax variant of GD because it aligns better with the Transformer's s attention; results for other attention kernels are presented in Appendix K). For both the GD and Transformer, the failure model of the bad examples is often repetition. The high quality of the simple GD-based "good" examples is deemed remarkable, given the relative simplicity of this model. In addition to the aforementioned limitations of the GD-based model, the Transformer employs a feedforward element after the attention and skip connection, while GD does not. For consistency, we included layernorm (Xiong et al., 2020) at the output of both models, before the softmax calculation.
For both the GD and Transformer based models, embedding vectors are learned for each token, with C' = 50, 257 unique tokens represented and an embedding dimension d' = 512; 8 attention heads are use for both models. Additionally, positional embedding vectors are learned for each of the 256 positions in our model's context window, with an addi- tional 257th position learned for the GD model (for position IN+1). We define P; as the positional embedding at po- sition and (as before) We,ys the corresponding token embedding for given sequence (y/ E {I...., C} is the token index at position i).
Recall that the purpose of this analysis is not suggest that GD-based model can replace the Transformer. Rather, it is believed that studying GD-based text generation may given insights into the Transformer, possibly suggesting ways it can be improved or simplified.
Concerning the GD-based attention model for language (summarized in Figure 7), recall that covariates I, are as- sumed for each position, including the query IN+1- To ensure are able to provide =N+1 for a context of size N, use exclusively positional embeddings for our covari- ates, we can easily learn 1 (referenced above). The
A quantitative analyis based on GPT-4o shows that the GD-based language model can achieve similar quanti- tative performance to comparable Transformer, with far fewer parameters. A similar GPT-4 based quantitative
Ittpps: cttp://aaaaccccooaattttt////s// aj ajibawa-2023/Chh Stori Strn--sttrrress-leeetos
8

---

On Understanding Attention-Baan Based In-Context Learning for Categorical Data
analysis was considered in Eldan & Li (2023) (more imple- mentation details in Appendix K). Each model was tasked with generating the ending to a story, above, but now these responses were graded by OpenAl's GPT-4o model. The grading model was told to evaluate the quality of the generated text based on the start of the story (see Appendix for details). It then provided a score out of 10 for the generation's 's grammer, its consistency with the start of the story, its general plot, and it's creativity. The table shows the average scores across 200 story generations.
Softmax GD Model (Good Example)
"Chnce upos time, small samed Harmonyvill,, live best friends, Max and Lily. Even though they were for from each other doe ciroumstances. they always found innovative ways to keep their bond strong. Max had moved to another city for his studies while Lily was still living in Homoosyville One sunny day, Max discovered how social media could make communication essier between them. He leamedd that messages sent over social media traveled almost instuarly across vast distances! These electronic signals cipped through wires airwaves incredible speeds, allowing Max Lily to talk real-time. It fclt they were truusponted into world full energy! Lily and Max conkdn't I wait shone their found knowledge with others. However, they soon realized something unexpected huppened. A storm hit Harmonyville, causing chans throughout the lown. Everyone was shocked and worried. They didn't 't know what to do. Lily and Max docided to talk to Dr. Maple. a wise old scientist known for her love of science and technology. She explained that when people feel overwhelmed, they night foet
Model
Grammar
Consistency
Plot
Creatiiy
Overall
Trie Ending
9.21
8.77
7.77
6.68
8.11
Softmax GD
4.25
3.40
3.08
3.76
3.62
Softmax GD (PF)
5,90
4.55
4.22
4.70
4.84
Transformer
5.97
4.98
4.43
4.75
5.04
Table 3. GPT-40 scoring of the generated story endings. Each
item is graded out of a maximum score of 10. Softmax GD and Softmax GD with FF have 8K attention-weight parameters, and the Transformer has 6M attention- parameters.
Softmax GD Model (Bad Example)
"Cnce there was a happy dog namtdd Spott Spot rus and play, but he was too big. His
weight was so heavy that be couldn't zun around he wanted to. Spot hated weight. wished could light the wind and be able run freety
The overall scoring of both the I-layer GD and Transformer models were low, relative to the true ending. However, while the Transformer yields better scores than GD, these scores are much closer to each other. Importantly, the GD model is able to capture much of the generation power of the Transformer model. To examine the importance of the feedforward (FF) element (plus associated skip connection) in the Transformer, which is absent from GD perspective, we kept all aspects of the GD-based model unchanged, but appended a FF network above the attention output from position N ++1,, 1, before the associated output goes to softmax over tokens. Note from Table 3, the GD formulation plus an added FF element yields performance almost the same as the Transformer! Associated examples of generated text are shown in Appendix K. This is deemed an important insight into the relationship between the GD perspective and the Transformer, and the importance of the FF element to Transformer performance. The performance of GD-based attention plus FF relative to the Transformer is particularly noteworthy, given that the GD-based model has far fewer parameters (see the caption in Table 3).
matter what Spot he couldn's it seem change. One day. Spot went adventure and found clear pond. Hie looked into the water and suw his reflection. Spot bated the way he looked. He wanted to see what be could do. He said, "Let's go to the pond and see if he could find it. le was so happy that he could help his friends. The moral of the story is to always be kind to help others. The moral of the story is to always be kind to help others. The moral of the story is to alwayys be help others. The moral the story to always kind to help others. The moral"
Table I. Generations from the GD model with softmax attention. The black font represents the prompt, and the red font represented generated language.
Single-Layer Transformer (Good Example)
"Onne upon a time, there was a big stake named Sam, Sam was flexible snake who could bend and twist in many ways. He lived in a forest with many triends. One day, Sam met little bird named Beety Betty was sad because she lost her to Sam said, "Don'' it warry, Berty. will help you find your toy!" They looked for the toy together. They looked for a su happy! You good friend." From that day on. Sam and Benty become best friends. They played together every day. Sam learned that helping each other is more important than having fun. And they lived hoppily ever after. The moral of the story is to always help your friends when they need it. Sam and Bent learned that it is important to
minute, them two, then theee, Fimally, they found toy under a big tree.
He saad. "Thank you,
7. Conclusions
be kind"
Through careful analysis and novel theory, as well de- tailed experiments, including on two real datasets, we have investigated in-context learning with categorical observa- tions. This analysis has led to a new attention- block frame- work that can perform multi-step functional GD exactly. Our language- generation attention-model perspective, with the addition of an ap- pended feedforward (FF) element, can generate text (trained on particular corpus) as well as comparable Transformer, but with far fewer parameters. More work is needed to in- vestigate how the FF element, coupled with the GD-based attention model, works with our new attention attenniio-bbsed model to yield such high-quality results.
Lingle-Layer Transformer (Bad Example)
"Onor spon a time, there was a yollow swan. The swan lived in small pond. The pond had pretty flowers and groen grees sround One day, the vane met a big frog. The frog uas not nice. The frug said, want play a gome." The swan not wam. play, but the frug vary big. The
language-generation experiments indicate that a GD-based
swim fister. The swan soam vory fast. The swan swam very fist. The swan swam very fixt. The swan saam and swam until reached the pund. The swan was happy to see the swan again. From that day on, the swan and the swam became best friends, They swam and played logether every day. The swan and the swin became best friends. They swam and swam together every day"
scaued. The suggest a game
Table 2. Generations from the single-layer Transformer model. The black font represents the prompt, and the red font represented generated language.

---

On Understanding Attention-Ban Based In-Context Learning for Categorical Data
Acknowledgments
functional gradient descent to learn non-linear functions in context. Int. Conf. Machine Learning (ICML), 2024.
The work reported here was supported by the Office of Naval Research under grant 31 310000130. We also thank the anonymous reviewers for constructive suggestions, that have impacted the final version of the paper.
DeepSeek.
Incentivizing reasoning capability in LLMs via reinforcement learning.
Deepseek-R1:
arXiv:2501.129481 2025.
Devlin, J., Chang, M.-W,, Lee, K., and Toutanova, K. BERT: Pre-training of deep bidirectional transformers for lan- guage understanding. North American Chapter of the As- sociation for Computational Linguistics (NAACL-HLLT), 2019.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning, particularly in the context of language modeling. There are many potential societal consequences of our work, particularly in terms of enhanced understanding of the mechanisms that underlic language models. Insights from this line of work may lead to better and/or more effi- cient language models in the future. This could, for example, reduce energy requirements for such systems, via enhance efficiency.
Eldan, R. and Li, Y. Tinystories: How small can lan- guage models be and still speak coherent English? arXin:305.007599, 2023.
Finn, C., Abbeel, P., and Levine, Model-agnostic meta- learning for fast adaptation of deep networks. Int. Conf.
Machine Learning (ICML), 2017.
References
Garg, S., Tsipras, D., Liang, P., and Valiant, G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Pro- cessing Systems, 2022.
Ahn, K., Cheng, X., Daneshmand, H., and Sra, S. Trans- formers learn to implement preconditioned gradient de- scent for in-context learning. Neural and Information Processing Systems (NeurIPS), 2023.
Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, 2010.
Ahn, K., Cheng, X., Song, M., Yun, C., Jadbabaie, A., and Sra, S. Linear attention is (maybe) all you need (to understand transformer optimization). International Conference on Learning Representations, 2024.
Liu, J., Shen, D., Zhang, Y., Dolan, B., Carin, L.,, and Chen, W. What makes good in-context examples for GPT-3? Proc. of Deep Learning Inside Out (DeeLIO): The 3rd Workshop on Knowledge Extraction and Integration for Deep Learning Architectures, 2022.
Akyurek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. What learning algorithm is in-context learning? inves- tigations with linear models. International Conference on Learning Representations, 2022.
Mahankali, A., Hashimoto, T., and Ma, T. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. Int. Conf. Learning Representations (ICLR), 2024.
Andrychowicz, M., Denil, M., Gomez, S., Hoffman, M., Pfau, D., Schaul, T., and de Freitas, N. On the opti- mization of a synaptic learning rule. Neural Information Processing Systems, 2016.
Muller, S., Hollmann, N., Pineda, S., Grabocka, J., and Hutter, F. Transformers can do Bayesian inference. Inter- national Conference on Learning Representations. 2022.
Bengio, S., Bengio, Y., Cloutier, J., and Geesei, J. On the optimization of a synaptic learning rule. Optimality in Artificial and Biological Neural Networks, 1992.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. OpenAl Blog, 2019.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J.,
Dhariwal, P., Neelakantan. A., Shyam, P., Sastry, G.,
Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray. S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. Neural Information Pro- cessing Systems (NeurIPS), 2020.
Ravi, S. and Larochelle, H. Optimization as a model for few-shot learning. International Conference on Learning Representations, 2017.
Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. Imagenet large scale visual recognition challenge. IJCV, 2014.
Cheng, X., Chen, Y., and Sra, S. Transformers implement
10

---

On Understanding Attention-Basn Based In-Context Learning for Categorical Data
Schlag, I., Irie, K., and Schmidhuber, J. Linear transform- secretly fast weight programmers. International Conference Machine Learning, 2021.
Schmidhuber, Evolutionary principles in selfreferential learning, learning how learn. Diploma thesis, Insti- fut Informatik, Tech. Univ. Munich, 1987.
Schölkopf. B. and Smola, A. Learning with kernels: sup- port vector machines, regularization, optimization, and beyond. MIT press, 2002.
Schölkopf, B., Herbrich, R., and Smola, A. J. A general- ized representer theorem. International conference on computational learning theory, pp. 416-426. Springer, 2001.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. In Int. Conf. Learning Representations (ICLR), 2015.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, F.,, and Lam- ple, G. LLaMA: Open and efficient foundation language models. arXiv.2.30.. 133771,,.11911,, 1971, 2023.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Attention is all you need. Neural Information Processing Systems, 2017.
von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. Transformers learn in-context by gradient descent. Int. Conf. Machine Learning (ICML.), 2023.
Wainwright, M. High-dimensional statistics: A non- asymptotic viewpoint, volume 48. Cambridge university press, 2019.
Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, I.., and Liu, T.-Y. On layer normalization in the transformer architecture. Int. Conf. Machine Learning (ICML), 2020.
Zhang, R., Frei, S., and Bartlett, P. Trained transformers learn linear models in-context. Journal of Machine Learn ing Research, 2024.

---

On Understanding Attention-Base In-Context Learning for Categorical Data
APPENDIX
where Weyg (i) is component j of wa,z Rd and S A1;x;;) + b. Therefore
A. Summary of Content in this Appendix
Ee explw aillwa.eilwee
This Appendix provides an extensive set details connected to material in the paper, as well as additional experiments. To aid the reader in navigating and using this Appendix, summarize what is provided, and in what section (with hyperlink to it).
VoC = N
We.gt
Ex=oexplul explw
Wa,y
B(we)lnn
(11)
We consequently have the GD update rule for b
Section B: Derivation of the GD Update Equation
N
bR+1 bn
We,Dr B(we))
(12)
Section C: Transformer Parameters for Single GD Step
Section D: Transformer Parameters for Exact Multi-
Similarly, let aj represent the jth row of A. Taking the gradient of L. wrt Aj:
Step GD with Interleaved Attention Layers
Section E: Details Generation of the Synthetic Data in the Main Paper
V =
We,sr (f)
exp Tww T. (Av)(z)) + b (1)
Section F: derstanding of Trained TF and Connection to GD
Dataset for Further Un-
Eeio exp(w? (Ab(z)) b))
Section G: Details on the Gradient Updates for Real Y and Gaussian p( X z)
We,gi (j)
N
Section H: Multiple Kernels for Function Modeling
Ee c=l
exp eo [il (b(xz))
Le exp(w" of fi)
Section I: Proof of Identity Stationary Point
Section K: Additional Language-Modee Experiments
We,y (i) B(wa(i))) ((z)). (13)
B. Derivation of the GD Update Equation
The gradient update step for aj is
The cost function for inferring the parameters ő = (A, b). with A E RRd xm and b E Red may be expressed as
aj.k+1 ajAh a aj
V.
C
We,ss (i) B(we ((we()))))))) ((x)).
agAA N
1
N
explw2 (Av(x;; + b)]
b(A,b) =
log
N
1= explw (I;( + b]]
Using the GD update rules for b and (asli=ı,d', we have
1=1
of af ((,)) + b1,841
1
Tw A1)(1;) + w.
N
fjtt1=
a k+ ((zz) + by ba
log
exp(w Avbes) + b)] (10)
Taking the partial derivative of L wrt bys component j of b:
= Sz, +
we,we B(we)len ((()))))))
0
N
N
L
We.gr (i)
+ a
We.y.
E(we)l:.a)
(14)
0b
N
1=1
Σ 1= explwa (Ap(w",.AA(( + 6)] b)]
C. Transformer Parameters for Single GD Step
Eezo expl
exp(wx w0 o, (A)(1;) + b)]
De 10 hoplwaffilwe...
Wee.g (i)
The GD update equation for the underlying function fo(x) is repeated here for convenience:
Ee explwe Si]
12

---

On Understanding Attention-Bal Based In-Context Learning for Categorical Data
After this initialization, the vectors are then input to the attention layer, which implements computation of Afio for all positions i, with keys and values inputs i 1, 1,...,. and queries for all j = 1..... N + 1. To implement this, we design the query and key matrices as follows
f3,411 fask
B(we)ls 18(22 xz))
N
Afaa
(15)
To effect a realization of the Transformer that implements this update rule in its forward pass, we consider the first attention layer, which can generalized an arbitrary number of layers. The desired input vectors the first layer
Laxd
Odxd'
Odxdd
Wq= WK
Od'xd
Oa' xd
Od'xd'
(21)
Og xd
0 xd'
such that WoELO (2r, 02a and then these are used in the kernel attention as k(Wqez,0, W Wxe.0)
It B(waDln,-0d
61,0 =
We,y
For the values, we need to extract N [Weey f(wo) from the input vectors, so the value matrix is
Od
Weyy
Le=1
We,e
(16)
Odxd
Oaxd
Odxd'
0g
Wv
Oa
Taxxa
0
(22)
N
0x'xd
Oa'
0g Od'xd'
We have initialized fr, Oz, ViE f1..... N+1), 1], which corresponds to the last d' elements of e4.0- The input for query N + is
such that Wye,0 (0s: a [We. Ee=1 We,el, Ow)""
This goes into the attention network, the output of which is (0a, Afj.,, Oz.)² g at each position
IN+1 Oa
EN+100 =
(17) Finally. we multiply the output of Wveso by matrix P. which corresponds to
0g'
Note that VN+1 not known the Transformer at input.
Oaxd O& xd Ox
Odxd
Odxa
The category embedding vectors reflected in the columns of W. are to be learned. They used in the input encoding, discussed above, and also within the softmax over categories at the Transformer output. Motivated by the need to learn W., in practice we encode the input vectors as:
P
Oa' xa'
0d xd 08
(23)
Ia Ia
This output is then added to the input, constituting skip connection, from which fr.o is updated to fr,1-
When we perform model learning for the GD-based form of the Transformer, based on ( we use the cross- entropy loss based on the softmax likelihood function. The parameters to be learned are W, E IR xC and the learning rate a (plus the kernel parameter). At the output of the last attention layer, the output at position N + 1 is multiplied by the matrix (Oæxx 0d'xa' La xa'xa)) which yields an ap- proximation for føx (#N+1) for K layers. This is employed with W. to manifest softmax probabilities categories (and cross-entropy loss).
Tz
6,, =
Lo 08
(18)
where Is is one-hot encoding of category yú i.e., D: vector of all zeros, except single 1, at position y; with y: E {O...., C}. and Le is C-dimensional vector of all ones. The input for query N + 1 (for which we do not have 3N+1 as Transformer input) is
xNN+1
eN+10 =
(19)
0.
D. Transformer Parameters for Exact Multi-Step GD with Interleaved Attention Layers
The first step at the Transformer input, prior to the first attention layer for implementation of GD, is to multiply e,,0 by the matrix
The input to the Transformer layer k is 6₁ (xi, Wear (i..) as summarized in Figures 1 and 2 of the main paper. Each attention block consists of self-attention layer, composed of two attention heads; one of these attention heads implements fik + fi,++1 like above (for which E(w.) is needed), and the second attention head erases E ((we)) Fe. preparing for its update by the sub- sequent cross attention layer.
Taxd Ox Oarxd
Oaxe
Oaxa
W.
O& x
(20)
Os'xC
Iaxxa")
where W. E df XC This implements the initial encoding of the vectors for implementing GD, and separates out the matrix W. for learning.
13

---

On Understanding Attention-Ban Based In-Context Learning for Categorical Data
D.I. Self-attention layer
K G4 X(2i, Oct 0s,
(32)
In matrix form, the input layer is
IN WOMNN E (lue) IN.K
NN+1 Oa
W S = M(zj, ,2, Ozr 0z))
(33)
We.y1 ([we) fiA
(24) where À > 1. With an RBF kernel, for example (similar things will happen with softmax), if À is very large,
B(we)ifx-LN fN+1K
f
The update equation for fA,41 is given by
x< x(W C W ej,,) dúg
(34)
fixt1 = fi.k + Afik
(25)
where 8.g = 1 if j. and it's zero otherwise.
where
The value matrix is designed as
w (2) (0s, Ov B(we)lnn 0z)
(35)
f
We - [[we)((w))))))) )x(2) (.
(26)
The output of this head
D.1.1. SELF-ATTENTION HEAD
0d 0g TB (B(w))
Oa
We design l() w and w such that
O
Our B(we)lsN1I
(
(F(w))
( K Cs,k (x; 0g, Oz, )
0
0
Of
(27)
(36)
We then add (())())) + ((2) with P(1) and P(2) de- (28) signed so as to yield the cumulative output of the attention
w.ll C, (zz, Osr, )
Od
Oa
Od
Of
Og
O& -(WW)lIN
w b = (0d, Og Ow N [eem - E(we)ifa,
(29)
O (total)
- -E(we)ln,
B(We)lNNN
Afi,
AfN.
Df
(37)
The output of this first attention head, at position j E (1....,N 1} is
This is now added to the skip connection, yielding the total output of this attention layer
(Oa, Oz, Od'
Weyy (..e)))..)))( (30)
.N
ØN+1
We.v
We.yx
Oa
T
(38)
0,
0.
O₂r fN+1A441
f1,++1
fN
The output of this first attention head at this first attention layer (before adding the skip connection) is
With the first attention layer, with two heads, we update the functions, and we also crase the prior expectations. In the next attention layer, we update the expectations, and place them in the locations of the prior expectations.
Os
Os
o
0g
O&
Our
(31)
0&
O&
Oa
AfLk
ANNA
AfN41E
D.2. Cross-Attention Layer
1.1.2. SELF-ATTENTION HEAD 2
The vectors connected T above will go into the next attention layer, where wish update
With the second attention head we want to add (0a, Oz. -Ew.) n,, )a from position j. so we clear out the prior expectation. This will provide "scratch space" into which, with the next attention layer type, will update the expectation, using f,,41 To do this, we design wx2) and (2( such that
w wh.le
E E(we)
We,c exp
(39)
w
in the aforementioned scratch space.
14

---

On Understanding Attention- Based In-Context Learning for Categorical Data
In this setup, the keys and values will be defined in terms of
(x) = Σ Wo(m) (o(m))BBB( mm]:
(46)
We (we,0:-..> -00 We.e)
(40)
where the RBF kernel parameter om for component m is selected such that KRBF [x 3mm); Fm] exp(-o || mmlll) equals 0.1 at the center of the other kernel to which it is closest (in Euclidean distance sense). We set À = 10 (41) (selected so as to have category c(m) be clearly most prob- able the region of (m)). Note that the attention- based inference assumes constant attention-kernel parameters for all keys/queries, and therefore these is model mismatch, in that f (1)) responsible for data generation employs different RBF kernel for the five components in.
There are many ways this can be done. Consider
Oa 0g'
WK WKW
0g
We,e
Oa
Ofr
WyWe =
(42)
We,c
0,
F. Lower-Dimennsonal Dataset for Further Understanding of Trained TF and Connection to GD
Let T; represent the ith column of T, then
0a
The following additional experiments consider data for which the covariates are two-dimensional. These experi- ments are considered so to provide enhanced visualiza- tion of the results. Further, because they of small scale, they allow explicit examination of the degree of match be- tween the Trained TF learned parameters and the associated GD theory.
Og
WoT=
(43)
0&
/
Then softmax attention with this setup yields the output for this softmax attention yields
Oa
Oa
Oa
Specifically, we consider data for which d 2, 20, and d' = 5. Two-dimensional covariates (d 2) allow visualization of predictions across the entire covariate space, and relatively small C and d' are also leveraged aid inter- pretation.
OZ
O
Ω
B(wa)In Of
B(we)1sNa4 0g
[[(w))
0₂
(44)
Via the skip connection for this attention layer, we have
The data are generated using (1) from the main paper, with We shared across all contexts, and with f(x) that is context- dependent. Matrix W. constituted by drawing each its elements iid from A00 1). For We so generated, a full set of contextual data (C() )t=,, ing the models (Trained TF or GD), five different random initializations are used for the model parameters (using ini- tialization methods as described in Glorot & Bengio (2010)). While results are shown here for one We, performance was found to be consistent across many different draws of W. (shared across contexts).
T40=
B(wa) f1,+11
WONN B(W.) f8
B(we)i fNALAH
Oge
(45)
generated. When train-
E. Details on Generation of the Synthetie Data in the Main Paper
We consider synthetic data generated by the model p(Y = clf(x)) = explw" (x)]// Ee exp[w of (x)), for C' = 25 and E RR. For data synthesis, the elements of the em- bedding matrix W. E 1 are generated (once) randomly, with each matrix component drawn i.i.d. from (0, 1). Af- ter W. is so drawn, different contextual datasets consider distinct function f( (ar), where I represents the context index. To constitute f()) (x). categories are selected uni- formly at random from the dictionary of C 25 categories. Let (1), (1)...., cll) (5) denote these categories for context I. We further randomly generate 5 respective "anchor po- sitions," ă(1).... ((5), each drawn i.i.d. from (O,,,, I,), where = 10 (for covariates z E R RO).. The function for context represented as
The 2D space of covariates is [-1, x [-1,1], with each of the two components of x; drawn uniformly over [- 1,1]; covariates were drawn similarly in (von Oswald et al., 2023). To define f(x) associated with a particular set of contextual data, the space of covariates is divided into four quadrants, and for given context (()) one of the = 20 classes associated with each quadrant (selected uniformly at random, without replacement). Let ( E {0...., C} the category associated with quadrant = 1,..., .,....,,4 for context x;; in quadrant g, we set f(x;) = w we
(1) of We. Using We and f(x) so where "eccf is column di
15

---

On Understanding Attention- Bastdd In-Context Learning for Categorical Data
0%0
RBF
Softmax
10.0
5M0
02-
30
30
Laplacian
Exponentiel
Skiing Capa
04
2003
3000
Figure 9. Depiction of p(Y = c|X = x) for queries x considered across the entire 2D space of covariates x. Results are for one attention layer and RBF attention, with the top row GD and the bottom Trained TF. The left, middle and right columns consider probabilities p(Y = c|X = z) for c = 7, c = 18 and c = 10, respectively. Contextual samples are depicted in white. where samples of categories 7, 10 and 18 are shown with respective numbers, and all other categories are depicted as circles.
Figure 8. For one attention layer, comparison of accuracy of GD and Trained TF, for attention based on RBF, softmax, Laplacian and exponential attention. The horizontal axes reflect Trained TF learning steps.
(1) for x; is drawn as in (I) from the constituted, category y) main paper.
For each contextual data ((), we expect the associated (xx) to have a nonlinear, step-like pattern linked to the quadrants. Note that within the data- generratnnntion process we have not explicitly defined (x) in terms of a function tied to any of the kernels (which is why, in the above, it was not ex- pressed (x))( However, for a particular type of kernel attention, in the forward pass of the Transformer the model effectively infers a context-dependent foo (a) approxima- tion to the function, tied to the particular contextual data C ((ag)) 3/ )}==1, observed and to the kernel used.
runs, and the lighter curves represent individual runs). These curves represent the frequency with which the true yN ()) was the most-probable prediction for query IN+1'
In Figure 8 the "mean max probability" is the average of the probability of the most-probable category from (1) in the main paper, averaged over all queries, and using the true underlying W, and the true stepped form of / /xx) employed when generating the data. This may be viewed as an upper bound on performance, and the capacity to achieve this bound is influenced by the context size (here N 100). The relatively low performance bound (less than 0.6) is due to setting d' = 5, which leads any column of We to have large inner product with a subset of other columns, yielding high correlation between associated categories.
We consider linear, radial basis function (RBF), exponential, Laplacian and softmax attention; each of these is kernel (with the reproducing property), except for softmax. Like (Cheng et al., 2024), we show results for softmax attention, for comparison; we show below that softmax attention per- forms well, consistent with the other nonlinear attention kernels considered. With the exception of the linear kernel, each kernel has single parameter, that is learned based (C()) We consider N 100 samples in each c()) and L = 2048 contextual sets are used for training. We present performance averaged over post-training contexts (C (t for K = 2048 and performance evalu- ated for the prediction connected to the query, given the context.
In Figure 8 we observe close agreement between the perfor- mance of the Trained TF and GD for all kemels. Similar agreement between Trained TF and GD was observed for linear attention, but the predictive accuracy was much worse, likely because of the highly nonlinear nature of f(x) (note as well the relative inferior performance of the exponential kerne)).
To further compare the predictions of GD and Trained TF, we plot the inferred probability for select categories, for queries covering the full range of the covariate space (here for one example set of contextual data C(D). For visualiza- tion, we consider a large number of queries for the same con- text, finely sampled across the 2D covariate space. Impor- tantly, in the Transformer implementation, all such queries can be considered in a single forward pass of the Trans- former (as discussed above). Results are shown in Figure
In Figure 8 we show a comparison of GD and Trained TF for one attention layer, for RBF, softmax, Laplacian and exponential attention. Five curves are presented for GD and Trained TF, corresponding to 5 different parameter initializations, and in most cases they overlap, although multiple Trained TF curves are evident for softmax attention (the bold curve represents the average across the training
16

---

On Understanding Attention-Bel Based In-Context Learning for Categorical Data
out, noting the large blocks of Wx and Wa that are near zero.
R
In Figure 10, the sub-matrices, A and B are identified, each in R 12122 Within the softmax attention, inner prod- uets (Ax;)² ((mr)) = x! Qr, are performed based on the learned form of We and Wx. where Q = A B. Columns and of A are found be very similar to the respective columns B (evident by a careful examination of the pixel colors connected A and B in Figure 10), and columns and 2 of these matrices are found to be almost orthog- onal. Consequently, for A and B learned by the Trained TF, Q E R is found to be nearly diagonal, with equal diagonal elements. The value of the diagonals may be seen as a parameter À within the softmax, e.g., as exp exxf æz), and hence the learned diagonal weight connected Q may be linked to learning the softmax attention parameter. There- fore, we infer that WQ and W K learned by the Trained TF are acting predicted by the GD analysis.
54
We
Wx K
Wv
E
fa
Figure 10. Parameters learned for one layer the Trained TF with softmax attention. Along each axis of the matrices, when applied to the ith data sample, elements (0.1) correspond to 2, E Rt,
elements (2...., 6) correspond We.s, De We,e R², and elements (7...., 11) correspond to Jick E CR.
Now consider the form of Wv, that has near-zero values on the clements that interact with 24, and hence Wy interacts almost entirely with War,g: C EE Wee.e as predicted by GD, Further, note that matrix P at the output of attention has only its last 5 rows with substantial values, and the last 5 elements of the output is where we anticipate the prediction of f(a is placed.
9 for the RBF attention kernel. These results are for one attention layer, considering the probability of 3 (of the 20) categories.
The three categories considered in Figure 9 were selected as follows. Using the true underlying data-generation process, we determine which category is most probable in each quad- rant of covariate space. Labeling the quadrants as 2, and 4, starting at the top-right and moving counter clockwise to the bottom right, the respective most-probable categories are == 7, 18, c 10 and с = 7 (c = 7 is most probable in two quadrants). In Figure 9, note the close agreement in predictions between GD and Trained TF, and that the high probabilities match the quadrant to which they are linked.
In Wy and P of Figure 10 we identify two sub-mattir matrices, E 122 and E E 12x12 These two sub-matrices are found to satisfy E D (evident in the colors in Figure 10). Further, each of the columns of D (rows of E) nearly orthogonal, with approximately the same lo norm. Q and the matrix product ED are shown to be very close to diagonal, with equal diagonal elements.
The updated f₁,1 manifested by the Trained TF may be ex- pressed as f R FDE [ 1. [x,)), 2g)) Since ED ~ EET R rTs. where is scalar and Is is 5 x 5 identity matrix, we have that the Trained TF represents fy, R tion for f,1 is placed in the last components of the atten- tion output. This is exactly how the GD model implements (4) in the main paper, assuming 7 plays role like a/N.
While Figure 8 indicates close agreement in the predic- tions of GD and Trained TF, further insight is gained by examining the underling parameters learned by Trained TF. Consistency between the learned parameters of the Trained TF and the properties of the GD-imposed parameters would indicate that the Trained TF learning to perform GD in its forward pass. We focus here on results for one attention layer, in that case GD-imposed parameters implement the exact update in (4) from the main paper. In Figure 10 we present the learned matrices Wo. WK. Wv and P for soft- max attention. The learned Transformer parameters were consistent across all synthesized datasets, and we discuss one representative example.
[e]xx(1s, 2y), and this approxima-
In the context of Figure 10, we now examine in detail the matrices Wq, WX. Wv and P learned by the Trained TE. In that figure, submatrices A, B, D and E were identified. In the left of Figure 11 we present Q = ATB 2g22 which ideally diagonal matrix with constant diagonal elements. Results are shown here for one of the five W. considered for data generation; results were similar for all data considered.
We first consider WQ and WK. The input at position is ei.0 = (z,, Wayyy - C De We,e. 0g)T, and via Wo and WK in Figure 10 attention directed only at Iy As predicted by GD, the term Wayy EE We,e is zeroed
Similarly connected to those results, in the right side of Figure 11 we show ED E R5x6 which also should be a diagonal matrix with constant diagonal elements, the
17

---

On Understanding Attention-Bann Based In-Context Learning for Categorical Data
f3N+1
fj
a
14
B(Y)
)x(r);
N
i=
filN +
Σ k==0
where DE, A As indicated in (49), A Aiy is the direct result of the attention mechanism.
Figure 11. Example of learned Q E IR (left) and learned matrix product ED 12555 (right). These based the softmax attention, but all kemnels provided similar results.
Real Y and Gaussian p(Y (x))) is special case, for which E(Y) S.2 ) is simply the latent function fi,k- More generally, and of relevance for the categorical data considered here, the needed expectation iss a nonlinear function of fi,
learned parameters are to be consistent with Trained TF performing GD step in its forward pass.
H. Multiple Kernels for Function Modeling
The results in Figure Il provide support to the perspective that the Trained TF, based on a single layer of attention, learns to perform one step of GD in its forward pass.
Assume that we have latent function
M
((x) =
A (amm(()
(50)
m=11
G. Details on the gradient updates for real Y and Gaussian X x>)
where each feature transformation vhm has an associated (erz))""mmm a diagonal matrix that places emphasis on different com- ponents of the covariates x. The set of feature maps or (equivalently) kernels are fixed, and the {Am) m= are context-dependent.
kernel (x;))" 4 (13) <(Am Amæ;), where Am is
We use the notation SJ.a fa (2z) to represent the function of interest after the kth step of gradient descent, evaluated at covariates zz. For real-valued Y and Gaussian p(Y|X = x)) (Ahn al., 2024; von Oswald al., 2023: Cheng et al., 2024; Ahn et al., 2023), have
Extending the discussion from the paper, employing func- tional GD, the update equation connected to the mth func- tion, m E (1,..., M). is fm, = Am, (x), we have
log )xlf(x;)) = 202 l11- f(x;)11
(47)
Assuming a constant, gradient descent for the param- eters A associated with the function fz) Axb(z) yields
fm,+1(())
N
= fmm)(()) +
O.m N
We.. - E(we ((we)l In(Am MAd Amzz)
N (y) - f()))((),
The total function is expressed as futi(I) Em=1 Pm.fm.41(0 where each Pn E Rd' The learned matrices standard way of merging multiple attention heads in the Transformer (Vaswani et al., 2017).
fa = fiN +
(48)
N
1=1
are consistent with the
Alik
Jik
(( (49)
N
In practice, the M kernels are connected to M self-attention heads. The use of the inferred function within (1) from the paper remains unchanged, and therefore the expectation E B(we)ln, is computed unchanged from the paper, using the cross- cross-attention layer.
where we typically initialize as fz. = Oz,, where Ozr a d' "d-dimensional vector all zeros. Consequently, we also have Afr. Og. The variance o² treated constant, and absorbed into the learning rate (effectively assuming that the variance o² may be approximated as constant for all contextual data of interest (von Oswald et al., 2023; Cheng et al., 2024)).
To make the connection to more- more-general Y and fx)),(x)).
we may re-express (49) as
18

---

On Understanding Attention-Baan Based In-Context Learning for Categorical Data
I. Theorem 2 and Proof
In this section, we present the formal version of the informal Theorem Section of the main paper. In Section I.I, we present the key setup and notation for Theorem which is similar to that of Section D. Section I.2, we present number of key assumptions.
I.I. Setup
For l E O...L, we let Ze E (Rddid+4)x(NN)) denote the input the t layer of the Transformer. For convenience, let us X2 I'
define Xę as the first d rows of Zz, L' the next d' rows of Ze, and Fy last d' rows of Ze, i... Ze
The "readout
Fe
of the Transformer's prediction for fa;) at layer k", referenced in Theorem 1, is thus given by the tt column of fi.
By the above definition, Zo is the input to the L-layer Transformer:
22
NN
IN+1
Zo
We,y1
We,y2
We,y
0&
0g
0g
0
0g
Let 5: Rd
:RA + Rdd x Rdd + Rd be defined as follows: for any f Rd"
1 E(f) = - 1/C E exp(w
exp(weed)
Bels[wank
Let h: (NN+1)x(N++) denote the column-wise activation (with 0 on the last column, to mask out the query). To be specific, each column corresponds to one of the N + 1 queries, and each row h corresponds to one the N keys. this analysis, is assumed that the attention, represented as f(.). is function of X BX,. The softmax attention is consistent with that, and hence this analysis extends beyond RKHS attention to softmax attention. In the context of RKHS kernels, linear attention may also be expressed in the assumed form, as can RBF-based attention, assuming that the covariates normalized for constant &g norm:. Concretely:
feoftmaz (K) = norm(M ex),
hRBF(K) == M cxp(K),
hiinear (K) == MK
(51)
INNN
ONx1
where exp(K) K) denotes element-wise normalization, M
and norm(.) normalizes each column of the
01xx
01x1
input matrix sum Note that plays a role like in the main paper. In the case of kernels, we consider the special case of kemeels for which the resulting function is dependent the set of inner products (x. ar i=,,, and it generalizes other functions that are dependent, such softmax.
Let E(X,T,F) BR (N+1)x + INN+1)xd" x R 1)xd' - R(N+1)xd denote the column-wise application of &, i.e.
E F)) E(FF) E [F)] =0
i 1...N
N=N+1
where in the above, [:]s denotes the ith column. The function E(Fe) computes the expectation (E(ee) for covariates x, for all covariates connected to the context, as in Eq (6) of the main paper. It is assumed that this expectation can calculated exactly; we show Section 3.2.2 how this is implemented with cross-attention. This theory evaluates the optimality of the multi-layered self-attention mechanism from the GD perspective, given the assumed exact computation of B(Fe) (which can implement exactly within the attention model).
Having defined all the above we define the forward pass of the Transformer as
X
Zef1 Z/ + Wv.t.
De
E(F))
M(z, Wo.eet)))
(52)
Fr
19

---

On Understanding Attention-Ban Based In-Context Learning for Categorical Data
We highlight that (52) differs from the construction in Section D. This is once again because, in 52, assume access which computes B(We))s,co)> whereas the construction in Section D) details how =\r may be implemented exactly via cross-attention layer. Our proof generalizes easily the setting of Section D, but writing things in terms of E significantly simplifies presentation.
Another minor difference is that Wy,e in 52 equivalent to PW in (22) and (23). This once again to reduce notation and does not affect the generality the proof.
Finally, under the above notation, the erosspyn loss
C
[(0)-
(FL(N+1) (FilN)))) log( 2
exp( (we,e) exp((w....FFilxxt)))
where 9 denotes the collection of parameters.
We verify that the cross entropy training loss is equivalent to
L(B(t, Rt) Exo. (oo, UN+1 B(t, R))]
1.2. Assumptions
Assumption 1. First, we assume that W k.e Wo, satisfies the following sparsity pattern
Br Ogr Oxxx x
Odxa
Oaxdd
W W.
0g xd
xd'
Og x
Oaxx
Og xdx
where Bg Raxdd We will use B = (Bo, Bt.. to denote the tuple of By's of all layers.
In words, the effect of Assumption 1 enforces that each self-attention module only computes the self-similarity matrix using the feature vectors (columns of X,). The embedding vectors r2 and function estimate F; are not involved in the computation of self similarity.
This assumption is deemed justified because We,y441 is unavailable for the query, and hence attention will focus on the covariates, which are observed for the query, the form of #N+1- The values only act on the same inputs as connected to the keys = ......... .N for which embedding vectors Weyy are available, and these that are tied to the language; therefore assume Wy acts the [We.g 1==1,N that encode the tokens connected to the context.
Assumption 2. We assume that Wy.e only updates the function estimates Fr:
Oaxd Oa Or
Oaxdd
Oaxa'
Wv.e
Oa
Oa'xd'
Tr
Ox
Furthermore, we assume that T, = / a scaling of the identity matrix, where L is a scalar. This enforces that the gradient update each coordinate of the function estimates (each row of Fe) have the same learning rate.
Assumption 3. We assume that Xo (which is the collection of N + 1 input covariates) has distribution that is rotationally invariant. In other words, Pr(X.) Xu), where Pr denotes the density or probability of Xo-
Assumption 4. We make a distributional assumption that all there exists a true latent function f. (:: Xo) from which all the class labels UI.. are drawn (recall that Xa denotes the collection of input feature vectors). In other words, we assume that for E (1..CC).
Pr(v: = c) x exp(w (.
(mg Xo)).
Furthermore, We will assume that
[f (2; [ x));
(53)
where ğ R+1 R-++ is any vector-valued function.
20

---

On Understanding Attention-Basee Based In-Context Learning for Categorical Data
One motivating example for assuming the form of f, the following:
(x x)
Σ any (y x)
where Vij R are arbitrary functions. This definition of motivated by the representation theorem (Schölkopf et al., 2001), when "Ay [2] 22y) kg) for some kernel In this case the function is in the RKHS family, in most of the paper. However, (53) is more general representation, and allows expansion functions like softmax.
L.3. Theorem Statement
Theorem 2 (Formal Version of Theorem 1)
Consider the setup in 1.1 and assumptions in 1.2. There exists stationary point of the cross-entropy loss E(0), where [Feli (the th column of F2) evolves (as l increases) according to the gradient descent sequence in (2)-(3).
1.4. Proof of Theorem 2
Under the sparse structure assumed Wy and W We in Assumptions I and 2 above, verify that for all t
X₂ = Xg = [1
22
EN
xN+1]
F7 = To = We.M. We.32
We.NY
O&
and
Fe+1 = - EFF)) B,Xo)
-
(54)
Next, notice that when Bg = 7ed (where by are scalars) for all e, the Transformer exactly implements gradient descent. This because (54) then simplifies
F = (1 tT) E(F:)))(b,X, X))
[Fe+1]: (1 te) E( ((we)) (bb Xo)
(55)
Notice that (55) above matches (2)-(3) up to naming of constants. h in (55) generalizes the role of in (3) (see discussion following (51)).
Therefore, it remains to verify that there exists stationary point of L under which B; b. for all layers f.
For the rest of this poof, it will be useful to explicitly write down the dependency of Fg as function of Xo, To, B:
Fe+1 (Xo, To, B) (I TT)
B,Xu)
(56)
We will also treat the parameters B as variables, and treat all other parameters as constants, i.e. E(0) =: L(B).
Our proof will proceed as follows:
1. Let S²(0) sylaxd for all é be some arbitrary initialization of B. Let S(0) = (S)
2. Let S(t) be defined by the dynamics
d Se(t) T- (5())==-rr(V,,(((()))
(57)
dt
We claim that for S(I) as defined above.
d (((()) < V1.
1.
(58)
di
1=
where || I| denotes the Frobenius norm.
21

---

On Understanding Attention-Bel Based In-Context Learning for Categorical Data
3. Notice that under (57), and under the assumed initialization, Se(O) are scalings of laxd for all I. Since E is lower bounded, (58) must then imply that Σz= BG ( - as + O. Furthermore, the limit of S(t) consists of scaled-identity matrices, and thus implements gradient descent as explained above.
It only remains to verify (58), which we do in the next subsection.
I.5. Proof of (58)
Let & be some fixed but arbitrary layer. Let R Idddd denote fixed but arbitrary matrix. Let Bg == belaxd for some arbitrary scalars by.
B(t, R) = {Bu, B2.., B By + tR, Bb+1 BL).
(59)
Le. in (B(t, R))e is equal (B(0, R)); B3 every layer f, except for layer k. where (B(t, R))s Bg tR.
We first verify the following three key facts:
1. For any orthogonal U E ddd and h Rd ((N+1) + E.
Weys+r)l B(oo. XU (h(To, We,go- )1
For any i = I...N + 1, let [Xo]: ne be the ith column of Xo. then
Pr(y c) exp(w f-
By definition of f. in (53), f. f,( g(X. U U[xo]:)
f(((l));;
The conclusion then follows by noticing that for each i = 1...N + 1, yt's are independent, conditioned on 24's.
2. For any and for any Xo, To.
Fe Io) = F,,
We will prove this by induction. The base is Fo 0, which holds trivially. Suppose the above holds for some l.
By (56).
=(T+T)((To EI B,UX))
To. BeXo)
=Xo,1(Xo, To. B).
In the above, the second equality is by the inductive hypothesis.
3. Let (4, O): (N++))(N++) x (N++))((++)) + (N++1)((++1) (A, C).
denote the Jacobian of i.c. =h(A+1C) =
We will prove by induction the following: for all E,
d
F,,UXo, To. B(t, RJ)
=
Fe(Xo, To. B (U,U)
(60)
dt
1=
dt
1=
For all E k, we verify by definition that
d FrUUX, To. B(t, B(t,
d
=0= -
FX To. B(t, T (RU)))
(61)
dt
1=1
dt
1=
22

---

On Understanding Attention-Baser In-Context Learning for Categorical Data
This is because Fg depends only on B, in layers O...e, which do change with L. Next, for layer k + 1,
d
dt Fo, Bo,
- - E(A((U B(, (X (Br (BC++
10
=(T+ (T++Ta((T) E(FA(U Tu, B(0, ( (X) U BUUX,, XOU RUXo)
-(1+ +(/+Th)) TT) - E(Po(Xo, To. B(O, (X) ByXo, XOU RU Xo)
IT)+T E(Fo(Xo, To, B(O, R)))) h(X (Be t0 RU)
t=
d To,, U RU))
dt
1=0
where the first equality is by chain rule and (61).
Finally, consider f k + 1. Let C=(A,C) R(N+) x R(N+1)xd R(N+1)x(d) denote the Jacobian of h i.e. = - 1C) 1=
= C).
d Lo,
dt
=(1 7b)( d =(F.(UXX. Po, B(I, U B,UX.)
dt
d
-(1 + - B(O, R)),
[,(UU,, To, B(1, R))
1)) B&Xo)
dl
dt =(1+ ( To. B(O, R)), Fi(Xo,
B(t, U RU)) ()) BeX+)
=(1 4 E(P((Xo To B B( RU))) )²(X, BeXo)
dt
di (Xo. To B(t, RU))
The first equality is by chain rule, the second equality is by definition of Jz. The third equality is by inductive hypothesis, and by item 2 above. This concludes the proof.
Let us define as follows:
g(Xo. To. B(t, = (Weoxen (F B(t, R))++1) - log Σ (F B(t, ((),)))))))
where B(1, R) is as defined in (59) above for some arbitrary matrix As a consequence of item above, we verify that
d
de gUU UN+I, B(t,
1= of oo. B(t,U RU))
(62)
1=0
23

---

On Understanding Attention-Bal Based In-Context Learning for Categorical Data
Then
d ((B(t, dt
10
di o( YN B(t, R))
Exo -oox
E BToNNN+INN
[ To, UN+ B(t, R))
Exo,U
E 1oNNNiL
g (UU To. UN+1, B(t, R))
dt
= Exo,U
E RoNNTIUUX
g( (o. UN, B(t, RU))
dt
d
= Exo,U
E BraNNILING
dt (Xoo, To UN+ B(t, RU))
= Ex.
ETo.NNII
dt To. UN+1 B Ey U RUT))
= Exo.To.NN+1
dt T(o. (oo,Too UN+1, B(t, EU RUT))
1=0
d E(B(t, RUT)
dt
1=
The second equality is by iterated expectation.
In the third equality, U is uniformly sampled random orthogonal matrix. This equality holds because we assume that P(X) P(UX).
In the fourth equality, we use (62).
In the fifth equality, we use item above, with h(To, We We. = d. 9 To UN+1. B(t, UT RU))
t=
In the sixth equality, we use the fact that Bro.ssil/o Lo(oo Po.1 B(t, BU RUT)) U RU, and is linear in U RU.
depends on U only via
1=
The seventh equality is again by iterated expectation.
Now let Bg = (t), where Se(t) is as defined in (57), and let R = - V ((S()) Then
C(S())]
(VB,C(S()),, VB, ((S()))
R)
d C(B(L R)) 1= dt C(B(L IUT RUD)
di
d
= (V D(S(t), EU RUI)
= V B. ((S()) Sull)
dl
where the last equality uses (57) together with the fact that Ey (UT RUJ (T((R). Since k is arbitrary, we conclude that
d ((S()) Σ Vn
d sa(0))
Σ
Vn, ((S())
k=0
24

---

On Understanding Ailention-Basce In-Context Learning for Categorical Data
which proves (58).
25

---

On Understanding Attention- BBsse In-Context Learning for Categorical Data
Model
Grammur
Consistency
Plot
Creativity
Overall
The rtudent coeds to ceoplets this story.
Io the folLowing coNccise, the ptudent ta glon peeracasien begiscing ot stery.
True Ending
9.21
8,77
7.77
6.68
8,11
laagsage
creativity.
Linear GD
4.77
3.64
3.30
3.81
3.88
Hore the prorecssion beginsing!
RBF GD
3.83
2.99
2.72
3.32
3.21
Softmax GD
4.25
3,40
3.08
3.76
3.62
Linear GD (FF)
5.94
4.57
4.20
4.44
4.79
</3E3TTIC 100r
RBF GD (FF)
5.11
4.09
3.72
4.07
4.24
Aod bacw La the atodicts pwappeow:
Softmax GD (FF)
5.90
4.55
4.22
4.70
4.84
Transformer
5.97
4.98
4.43
4.75
5,04
</STIEEIT RESPOESE:
Figure I4. GPT-4o scoring of generated story endings for all mod-
First, proxioe a concine qallitetive coossenens shous the otetens's veising give uriting grate
cls. Each item is graded out of a maximum score of 10.
follossin
drawaar
Linear GD Model (Good Example)
the ariniin gromatirally correct?
2raloste systis, puactsatioon und suntencce siructure.
Cassiakancy:
the stadont's uriting coseletent with the procided begtaxing the stooy?
"One day, a boasy cat and a nice dog were playing a sport. They were kicking ball. The was very bassy and wanted to win.The dog said, "Let's play and have fun. We can both be good!" The cat did not like this idea. The wanted to be the best. The cat said, "Okay, let's play together" They played and had fün. But then, something unexpected happened. A big wind came and blew the ball for away. The cat and the dog tried to catch the ball, but they could not catch it. Then, something unexpected happened. A hig wind came and took the ball away. The cat and the dog were very surprised. They did not know what to do, They tried catch it, but they could not" "
Bae will dons the stadeet coaplate the fital seesence of the proseribed beginasin?
7.001
Jeen plos of the stodest's writing make astse Cregerdiess of the previdel teginningl
creative the atudent's writing?
Forest your reggesso as follsce:
rrattin asserzintt grassr]
SQUATPAR GAREE>
[Orada WNE of [X]
(Rakiitative aasesaaa of coastasency]
Linear GD Model (Bad Example)
Draate 102
</OUSSSTTESYY Clate>
<P50T>>
"Once upon a time, there was a little girl called Lucy. She really liked playing with the board. She had lots fun with every day. One day, Lacy was playing with her board in her backyard. Suddenly, she beard a loud noise. It was the delivery track. Lucy yelled as she spat, "Oh no, the delivery is here!" !" Lacy frantically ran inside and quickly grabbed her board. She wanted to keep it away from the delivery driver. But, was too late. The delivery driver saw her and yelled, "Hey, give me!" Lucy was so happy to have new friend. The next day. She said, "Let's go!" Lacy was so excited okay!" Lucy happy. She thanked her mom and said, "It's okay." Lucy smiled, She said. "It's okay.""
plas]
[drada sit of 13]
4/2131 mAED
<DSERTIVETT>
(Raaciaative sasemmat creativity)
</CREATE/ETD>
0SSF4TIVITY_ JBRA1R>
fGrade out 13)
Prossde your acseraent beloz
the package. She said, "It's
Figure/ 12. The user prompt given to GPT-4o to generate the model scores.
You are a writing worluster dexl designed to somes acndent story complacions. You will be provided children's stariee vristan for 3-4 your ald audiance. Your role provide constractive, faix, and detailed evaluations based apecific rubric criteria.
Table 4. Generations from the GD model with linear attention
Figure 13. The System prompt given to GPT-4o.
RBF, and softmax GD models with the addition of the feed- forward element, as found in the transformer model. While this deviates from gradient descent, it allows for an even closer comparison with the transformer.
J. Language Model Evaluation Details
Our quantitative evaluations used the OpenAl API with the GPT-4o model. For each story, GPT-4o was given the user prompt in Figure 5, with the BESTORY_BEGTN++ key- word replaced with the true beginning of the story (used as the evaluated model's input) and the **STORY_ENI END+* keyword replaced with the generated ending to the story. The numeric scores were then parsed from the GPT-4o model responses. Additionally, the system prompt in Figure was provided for all evaluations.
K. Additional Lodguage-Model Experiments
In the main paper, we considered only doltmax-basse at- tention to allow for a more direct comparison to the Trans- former model. Here we include results for the linear and RBF kernels In addition, we provide results for the linear,
26

---

On Understanding Attention- Baatdd In-Context Learning for Categorical Data
RBF GD Model (Good Example)
"Once upon time, the city of Columbus, there professional soccer team called the Columbus Crow. This team had lots of dedicated fans who loved watching them play every week. One day, two best friends named Alex and Jamie deci ded to go to Columbus Crew game together. They were so excited bocause they heard that the team had something extra special planaed for their fans! soon as they arrived at the stadium,, they received free t-shirts hals with the Columbus Crew logo on them. "Wow!" said Alex, putting on his new hat. "This cool! I love getting free stuff. Alex and Jamie decided to visit the local library. As they walked, they noticed something strange happening around them. There were many peuple wearing masks, wearing masks and gloves, hats, gloves and masks, Mr. Johnson explained, "These masks can help protect our hodies need to protect us from harmful germs. We need to protect our masks and protect ousselves against harmful genns." Feeling inspired, Alex Jamie decided to spread awareness througbout neighhorhood. Soon
Linear GD Model with Feed-Forwaid (Good Exam- ple)
"One day, a little girl named Lily wanted to play dress-up. She went to her room and looked at all her costumes. She saw a big box with a new costume inside. Lily was very happy opened box Lily took out costume examined It superhero costume with red cape. She put and said, "Now I can be superbero!" Lily around the house, laughing and having fun. But then, something unexpected happened. The costume started to talk! It said, "Thank you. Lily! ! am a magic superhero! I can grant you one wish." Lily was so surprised and happy. She wished for a new friend to play with. From that day on, Lily and the magic faicy bocame best friends. They had lots of fun together, And they lived happily ever after. The moral of the story to always"
Linear GD Model with Feed- Ford-Forwwrrd (Bad Exam- ple)
enough,"
RBF GD Model (Bad Example)
"Once upon time, there loncly tree. The had one olive. The olive was green and small. A little gidl named Lily lived near the tree. Every day, Lily would to the free and gaze at the olive. One day. Lily saw a bird on the tree. The bird looked at the olive too. Lily said, "Hi bird, do you like the olive?" The bird said , "Yes, I can help you." Lily and the bird became good friends. They played together every day. They had lots fun. And they lived happily ever after. The moral of the story is be kind and help others when they need And they lived happily ever after. The moral of the story always be kind and help others when they need The moral of the story is to always be kind and help others when they need it"
"Once upon time, there small bird named Tim. Tim was very proud of his language. He could talk and sing better than any other bird in the forest. One day, Tim met a big cat named Max. Max was not nice. He liked to demand things rom others. Max said to Tim, "Sing for me, little bird." Tim was scared, but he sang his best song for Max. Max was happy with Tim's 's song. but be wanted more. "Sing again." Max domanded. Tim sang again, but he was tired wished could
way to stop Max from tree. The moral of the story to he kind help others. The moral of the story is to he kind help others. The moral of the story is to be kind to help others. The moral of the story is to be kind to help others. The moral of the story is to be kind to help others. The moral of the story is to be kind to help others. The moral of the story is to be kind help others
Table 6. Generations from the GD model with linear attention (with the foed-forward)
Table 5. Generations from the GD model with RBF attention
27

---

On Understanding Ailention-Basce In-Context Learning for Categorical Data
Softmax GD Model with Feed-Forward (Good Ex- ample)
"Once time, in a town called Meladyville, there lived two hest friends named Rhythm and Harmony. They loved exploring different types of music from all around the world! One day they heard that the annual Musicland Awards were ha ppening soon, and they wanted to learn more about what made musicians special. Rhythm
teacher, Tempo, the Musicland Acadomy choose winners?" Mr. Tempo explained, "Well, like in real life, many things play part including cultural context and political climate." "aruutee"aaaaa "Hamony was curious, so she asked, " "How do we know if we can find out more?" Mr. Rhythm explained, "Well dear friends, it depends on various factors such as sound waves, frequency, amplitade, and pitch. When we listen each other, creates sound waves that travel through the air and our ears." Melody addded, "That sounds umazing! Mr. Melody smiled and replied, "That's right! Sound made up of tiny
does work?"
Softmax GD Model with Feed-Forward (Bad Exam- ple)
"Once there was generous little girl named Amy. She wanted to to the fountain so she could play. She zipped in her rod shoes and ran towards the playground. When she gol to the lountain, she saw two boys splashing around. Amy asked if she could play too, but the boys said no. They were not being very generous. Instead, Amy decided to fun. The moral of the story that it's important
toys with friends. They
played together and had be kind and share with each other. The moral the story to always be kind and share with each other. The moral of the story is to always be kind and share with your friends. The moral of the story is to always be kind and share with others. The moral the story"
Table 7. Generations from the GD model with softmax attention (with the feed-fore (orwaorr
28


Reviews:
[{"review_id": "ZD3FGDLFmK", "reviewer": "Reviewer_4VEb", "summary": "The paper explores theoretical understanding of the In-Context Learning in Transfomer-stack models while dealing with categorical data. \nIt attempts to design a transformer block that can do gradient-descent in-context.\n\nAuthors try to construct a transformer stack that can, in theory perform ICL on categorical data.\nThey use a softmax transformer for this. The task is to ‘learn” a function “f” from the context, from categorical data using transformer’s forward pass.\nThey start by expressing the required GD equation for this function, in terms of output embeddings generated by the transformer\nThey next show in the special case of the real-valued function (regression tasks), this task is easily handled by simply self-attention. This work has already been published.\n\nThis paper then engineers input representations and Q, K, V weight matrices to achieve the above derived GD update equation to “f”.\nThey go on to show the performance of the GD-enabled transformer stack on synthetic, image and language modeling tasks.\n\n## update after rebuttal\nAuthors have proposed adding another diagrammatic representation of their tecnnique. That might help fix the fact that the paper is hard to understand and follow. No other update from my end", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very careful review and helpful feedback.\n\nWe will work to make the paper more readable and understandable in the main body of the paper. We constituted Fig 1 with the goal of trying to summarize the setup in a figure, but we can do more. Please see here [\\[LINK\\]](https://anonymous.4open.science/r/tmprepo-4447/transformer_rebuttal_additional_plots.pdf) a draft figure we propose adding to the body of the paper, to summarize things and hopefully provide more intuition without getting too much into the technical details. We will revise the final version of the paper, with the goal of enhancing readability. Thank you for looking carefully at the Appendix. As you suggested, we will try to move as many insights from there into the main paper. \n\nWe will fix the typos and will implement your suggestions.\n\nOn your question: Each attention block implements one step of functional GD. An attention block corresponds to a self-attention layer and a subsequent cross-attention layer, as in Fig 1. Multiple attention blocks implement multiple steps of functional GD. In Fig 3, for example, the multiple layers of attention blocks corresponds to multiple steps of GD.", "claims_and_evidence": "They claim that the particular structure of transformer stacks suggested in this paper is capable of showing GD-learning type behaviour for in-context data and produce proper probability distributions for categorical class prediction\n\nTheir claims are well backed by the experimental results they’ve show with varied objective - synthetic class prediction task, image task and next-token prediction too. The additional heat maps for learned-matrices showing confirming the predicted values for weights is a good cross-check on the overall learning happening in the network", "methods_and_evaluation_criteria": "Their evaluation strategy for the tasks look good.\nThey also claim to have taken care to use distinct contextual data for testing the models after training. The evaluation setup is really sound in the sense that the test-time classes are not trained on. New set of classes & corresponding embeddings are provided directly in context and the model appears to give good test-time quality after sufficient training.", "theoretical_claims": "The paper is extremely dense in theoretical proofs. They claims and proofs mostly looked sound to me. Pointing out some typos etc in \"comments\" section. Some errors/gaps may have slipped by me though", "experimental_designs_or_analyses": "Experimental setup and evaluation looks sound, if the claims are accurate.\n\nNo issues found on my end.", "relation_to_broader_scientific_literature": "Resources\n- Transformers Implement Functional Gradient Descent (Cheng et al): https://arxiv.org/pdf/2312.06528\n- Transformers Learn in-context by GD (https://arxiv.org/pdf/2212.07677)\n- WHAT LEARNING ALGORITHM IS IN-CONTEXT LEARNING? (https://arxiv.org/pdf/2211.15661)\n\n\nWent through some papers listed above that analyze implicit learning algorithms implemented by trained transformers for tasks like linear regression.\n\n\nNormal proof methodology involves showing that:\n1. Attention mechanism _can_ emulate learning algorithms for task T\n\n2. the data transformations involved in GD/other learning algorithms for task “T”, etc can be performed by transformers stack. OR, designing Q, K, V weight matrices such that the forward-pass can create equivalent learning outcome\n\n3. Trained transformer stack do seem to emulate learning algorithms on task “T”\n\n4. By creating special data for task “T”, comparing performance of trained transformer stack to GD, etc\n\nThis paper uses a special transformer block to make the inference pass behave like a GD update.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths\nPaper has a highly mathematical approach to designing a GD-emulating transformer stack. It’s exhaustive, discuss all details required for understanding the hypothesis and the model setup. The extra material in the Appendix is definitely useful and aids understanding of the parameter setups. Lot of work exists around linear transformers emulating real-valued outputs - this work builds on top of that & extends that understanding to categorical outputs with softmax (albeit with a big constraint on the transformer block)\nCategorical outputs are especially useful.\n\nThey have made good effort to exhaustively test the methodology on a wide array of tasks and are able to show that their approach manages to emulate learning via GD. The technique has been tested on sufficient number of very different scenarios. The evaluation criteria is pretty strong. The details in Appendix are exhaustive and cover theoretical aspects pretty well.\n\n\nWeaknesses\nPresentation:  The paper is extremely dense.  This is likely because of 8-page limit, but lot of details could have been omitted and the paper could have been made more readable/accessible with some helpful diagrams. \n\nSome important explanations didn’t find place in the main paper. As an example in the same domain, [Akyurek et al] https://arxiv.org/pdf/2211.15661 is lot more readable at around the same length", "other_comments_or_suggestions": "TYPOS:\n\n- in Section 3.2; line 2: \"k\" in \"W_k \" needs to be capitalized\n- Large number of Typos in Appendix \"E\" (x \\in R^{10}, incomplete sentence in the end, ...)\n\n\n\nSUGGESTIONS:\n- In appendix, around equation 35: just stating P1, P2 = +-1 might make it clearer\n- Param setup with Q, K, V matrix sizes involving 0_{3d} etc look sprung-up on the reader at first-pass. Not sure if there's any way to better justify that, but I couldn't get past without Appendix read in detail.", "questions_for_authors": "1. Do you have a hypothesis on why we need multiple layers of the transformer block, if a single block can emulate GD perfectly well?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E5PymiXiPY", "reviewer": "Reviewer_GFKD", "summary": "The paper seems to be dealing with learning the implicit function embedded in the examples within prompt in in-context learning setting. \n\nThe paper is poorly written and difficult to understand. The introduction does not clarify what the proposed method is aiming at, there are disconnected/disparate components stated in methods section that make it very difficult to follow. The experimental setting and description are equally murky. \n\nMore details below", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper. Concerning your comments:\n\n> The introduction ....\n\n- lines 20-21 state that the goal of this paper is \"extending the functional GD framework [for ICL], to handle categorical observations\".\n- Reviewers 4VEb and Hsdu both provided accurate summaries of the goal of the paper in the \"Summary\" portion of their review.\n\n> ... disconnected/disparate components ....\n\n- Sec. 2.1: we describe the problem -- the labels y conditioned on covariates x are drawn proportional to $\\exp(w_{e,y}^T f_{\\phi}(x))$ for some unknown $f_{\\phi}$ that the Transformer learns in-context (implicitly). \n\n- Sec. 2.2: we describe the equations for updating $f_{j,k}$ to $f_{j,k+1}$ in a single functional GD step (lines 110-113). The key formula is Eq (3), which consists of two parts: (A) computing $w_{e,y}-\\mathbb{E}[w_{e}|f_i]$, and (B) computing a weighted RKHS average $\\sum_{i=1}^N (\\cdots) κ(x_i,x_j)$. \n- Sec 3.2: we show that a self-attention layer can exactly compute (B): the weighted average $κ(x_i,x_j)$. (line 184), and in Sec 3.3 we show the cross attention can exactly compute the expectation in (A).\n\n>  .... designed for explicit learning of this function. Is my understanding correct?  ....\n\n- This is incorrect. Our goal is also to learn $f_{\\phi}$ implicitly. \n- Ahn and Zhang considered **linear** latent functions with **real** observations. This led to a gap between theory and practice, and prior theory has **only been validated on synthetic data**. \n- We address a gap in the literature: Our central contribution is extending existing analysis to categorical data. There are significant complexities, such as the fact that the GD update is no longer linear, and contextual demonstrations are no longer $y_i=f(x_i)$ but rather a random draw $y_i\\sim \\exp(w_{e,y}^\\top f(x_i))$.\n\n> The experimental setting  ....\n\n- Our goal is **not to propose a new model**. We have not alluded to or suggested that the proposed model is superior to existing Transformers. \n    - The model in Sec 3 consists of only 3 components: self attention, cross attention, and skip connections. All of these modules are **standard components of a Transformer**. We noted on line 119 that (Vaswani et al 2017)'s original Transformer architecture also contains interleaved self-attention and cross-attention layers. \n    - In 1(b), when we say \"introduce a new attention-based architecture\", we mean with respect to the existing linear-Transformer design of previous theoretical papers such as (von Oswald, Ahn, Zhang, Cheng). **While it is true that an important innovation of our paper is the analysis of how cross-attention perfectly facilitates GD with respect to categorical data**, we again emphasize that **cross-attention is a very standard component of modern Transformers**.\n    - 2(a,b) are simply experiments to validate our theory -- that the Transformer construction in Sec. 3 is capable of matching in-context GD for tasks on real-world datasets. We are unsure why this \"alludes to a new model proposal\".\n\n> ... do not qualify for a theory paper ...\n\n- Can reviewer GFKD please **elaborate on why \"the method description and analyses do not qualify for a theory paper\"?** \n    - In Secs. 2 and 3, we **show rigorously that our proposed construction indeed implements exactly one step of gradient descent for categorical data**. Proof details are provided in Appendix B,C,D.\n    - In Theorem 1 (also 2), **we prove that our proposed construction is a stationary point of the training loss**.\n    - The strength of our results are no weaker than analogous results in earlier works that analyze multi-layer Transformers on real-valued data.\n    \n> ... (Ahn, Zhang) already demonstrates ... \n\n- This is incorrect. Our goal is also to learn $f_{\\phi}$ implicitly in the Transformer forward pass, like Ahn/Zhang, but for the first time here for categorical observations.\n- We acknowledged the importance of existing work, such as Ahn and Zhang. However, most practical applications of Transformer networks are over **categorical-valued** data. This led to a gap between theory and practice, and prior theory has **only been validated on synthetic data**. To the best of our knowledge, our experiments in Sec 6 are the **first time** that the ICL theory for Transformers has been validated on **any** real-world tasks and data. \n- The central contribution of this work is extending existing analysis to categorical data. There are significant complexities, such as the fact that the GD update is no longer linear, and contextual demonstrations are no longer $y_i=f(x_i)$ but rather a random draw $y_i\\sim \\exp(w_{e,y}^\\top f(x_i))$.\n\n> ... some experimental results  ...\n\n- The goal of our experiments is **not to propose an architecture change for existing LLMs**. Our experiments **empirically validate our theory on real-world data and tasks** -- something which was **not possible in previous papers** as they **do not handle the setting of categorical-valued observations**.", "claims_and_evidence": "First of all, is this a theoretical treatment of function learning during in-context training or a more practical model proposal to explicitly learn the aforementioned (implicit) functions? The title, Lines 041 ~ 014 next column and section 5 give an impression of theoretical analysis whereas contribution 1(b), 2(a,b) and Section 3 alludes to a new model proposal. \n\nThe method description and analyses do not qualify  for a theory paper and the experimental results do not seem to pass the bar of a practical model proposal paper.", "methods_and_evaluation_criteria": "I am also confused what are the methods section claiming. Is a method to *explicitly* learn the function represented implicitly in the prompt during in context training? Section 2.2, and subsection in Section 3 appear to be designed for explicit learning of this function. Is my understanding correct?\n\nIf so, why would we need to learn the function explicitly when multiple studies (Ahn 2023, Zhang 2023) already demonstrates that these implicit functions are learned organically during in context training?", "theoretical_claims": "---", "experimental_designs_or_analyses": "Although the paper shows some experimental results on public datasets, the problem setting for Imagenet classification Section 6.2, and Language modeling in Section 6.3 qualifies for a demonstration on a toy problem, not a full fledged evaluation needed for a new proposed model architecture. \nBut, then again, I am not sure if this is a theoretical analysis paper or not.", "relation_to_broader_scientific_literature": "--", "essential_references_not_discussed": "--", "other_strengths_and_weaknesses": "--", "other_comments_or_suggestions": "---", "questions_for_authors": "---", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QQaCa2yOfY", "reviewer": "Reviewer_Hsdu", "summary": "This paper investigates how transformers perform in-context learning on categorical data, extending prior work that has largely focused on in-context regression tasks.\nThe authors provide a kernel-based functional gradient descent perspective, wherein each layer of the transformer can be interpreted as performing one step of gradient-based inference using in-context examples.\nThey first demonstrate an expressivity result: that attention-based architectures can implement kernel-based gradient descent for categorical inference.\nThey further prove that their idealized construction corresponds to a stationary point in the loss landscape.\nMoreover, they empirically validate their framework, showing that trained models converge to solutions that closely align with their theoretical predictions.\nTo support their findings, they conduct experiments on datasets such as ImageNet, Tiny Stories, and Children Stories to demonstrate the predictive power of their approach across both language and vision tasks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review and for your valuable feedback. On your questions:\n\n1. \"Trained TF\" and \"GD\" (we agree that it would be good to re-introduce them here, despite the mentions in Section 4; we will do that). \nThe detailed derivation of the GD updated equation is in Appendix B and summarized in Eq (13). This is functional GD in an RKHS space for the latent function, based on cross-entropy loss for that function (with the loss minimized via functional GD, effected in the Transformer inference stage, based on the observed contextual data). In (20) we show the form of the $W_Q$, $W_K$ matrices, and for these there are no learnable parameters -- within a permutation, these matrices come directly from GD theory. In (21) we show $W_V$, and for that there is one learnable parameter, learning rate $\\alpha$. There is in general a different learning rate at each layer. In addition, for GD we learn the category-dependent embedding vectors $\\\\{w_{e,c}\\\\}\\_{c=1,C}$. By contrast, in Trained TF all parameters are learned, without restrictions (including the category-dependent embedding vectors). In Figs 8 and 9 of the Appendix we show close agreement (on a low-D example, for interpretation) between Trained TF and GD parameters, indicating (as supported by our theory) that the Trained TF does indeed learn to perform inference like functional GD. However, Fig 4 shows that Trained TF needs a lot of data to achieve such learning, while GD-based composition of the model parameters trains well with far less data (as it has far fewer parameters).\n\n2. Synthetic data generation: In Appendix E we provided more details. In short, there is a latent, context-dependent function $f(x)$, which the Transformer infers (at the $N$ sample points), in context. To generate synthetic data, we first synthesize the embedding vector $\\\\{w_{e,c}\\\\}\\_{c=1,N}$, with each component of each vector drawn from $N(0,1)$. These vectors drive the softmax on categories, given $f(x_i): SM(f(x))=\\exp[f(x)^Tw_{e,c}]/\\sum_{c^\\prime} \\exp[f(x)^Tw_{e,c^\\prime}]$. With {w_{e,c}} set, we generate $f(x)$ as a summation of kernels, to yield a highly nonlinear function (kernel centers positioned at random). Then, for $N$ samples, we draw $N$ covariate vectors $\\\\{x_i\\\\}\\_{i=1,N}$, from N(0,I). We stick each $x_i$ into $f(x)$ to yield $f(x_i)$, and this goes into the $SM(f(x))$ over categories. We draw a category $y_i\\sim SM(f(x_i))$. Now we have $\\\\{(x_i,y_i)\\\\}\\_{i=1,N}$ which is the context sent into the Transformer, along with query $x_{N+1}\\sim N(0,I)$. Each contextual set has a uniquely (randomly) generated underlying $f(x)$.\n\n3. In-context image classification. Consider context that is characterized by 5 image types (classes) that have never been seen by the ICL model. In the contextual data, we are given 10 examples images of each of the 5 image types (e.g, type of dog, type of cat, etc). So we have 50 contextual samples, each one of 5 classes. The image is represented in feature space by an arbitrary (but good) feature representation. Here we have used VGG features from a CNN, but they could also be from a ViT. For image $i$ let the features be represented as $x_i\\in R^d$. For the contextual samples, we also have the label $y_i\\in\\\\{1,\\dots,5\\\\}$. We have five classes, but the details of those labels and image types is arbitrary. Key is that the ICL model has never seen these image classes before. The ICL is given context=$\\\\{(x_i,y_i)\\\\}\\_{i=1,50}$, manifested in terms of 10 (random) examples for each of the 5 classes. It is also given a query $x\\_{N+1}$ from one of the 5 image classes, and is asked to label it, or more specifically $p(Y=c|X=x_{N+1},\\text{context})$, for $c=1,\\dots,5$.\n\n4. Language modeling. In-context classification has $C$ classification labels. For language, the $C$ labels correspond to all the (discrete) tokens ($C$ is the number of unique token types). For ICL, we have context $\\\\{(x_i,y_i)\\\\}\\_{i=1,N}$, x_i covariates and y_i category (token) label. In ICL, the category/token is encoded by a learned embedding vector $w_{e,y_i}$ for token $y_i$, just like in NLP. To model language, we use the positional embedding vector to represent the covariates $\\\\{x_i\\\\}\\_{i=1,N}$, and also $x_{N+1}$. With this mapping ICL can be used directly for language modeling.\n\nAs recommended by you, we will add a figure in the main body of the paper, to help with readability. Please see here [\\[LINK\\]](https://anonymous.4open.science/r/tmprepo-4447/transformer_rebuttal_additional_plots.pdf) a draft figure (which we will seek to further improve) to aid readability and provide clarity.", "claims_and_evidence": "Yes, it seems so.", "methods_and_evaluation_criteria": "Yes, I think so.", "theoretical_claims": "Yes, they make sense to me. I checked some proofs in Appendix B and C.", "experimental_designs_or_analyses": "The experiments (Section 6) were not clear to me, and I would appreciate clarification on the following points:\n* I understand that you are comparing two models: \"Trained TF\" and \"GD\" (it would be good to re-introduce them here, despite the mentions in Section 4). Could you please clarify how \"GD\" is constructed --- you would need access to some $\\psi$ to make $f_\\phi$, and how is this $\\psi$ derived/learned?\n* I'd appreciate more detail + exposition on why you chose to generate synthetic data in this way (Section 6.1)\n* I don't understand the problem setup for in-context image classification (Section 6.2). Could you please clarify what an example task might look like?\n* I also don't understand the language modeling experiment (Section 6.3). Similarly, could you please provide more clarification on example tasks --- in addition to those in Table 2? Could you please make the connection between in-context classification and text generation more explicit?", "relation_to_broader_scientific_literature": "This expands on our understanding of in-context learning, particularly for categorical data. Much of the literature focuses on regression, so looking at in-context classification is quite important.", "essential_references_not_discussed": "Not that I'm aware of.", "other_strengths_and_weaknesses": "Strengths:\n* The theoretical contributions look good, and I appreciate the mathematical rigor!\n\n\nWeaknesses:\n* I had trouble getting through the notation, particularly in Section 2. However, because these notations are so important to the paper, I implore the authors to generously use diagrams and examples. Some discussions and experiments may be added to the appendix if space is lacking.\n* As mentioned above, some experiment descriptions were not clear to me. For example, in-context image classification is not a \"natural\" task for me, so I would appreciate some examples --- even if they may be toy.\n\n\nOverall, I found this to be an interesting read and a solid paper, but I believe that its clarity could be refined.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "See above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper investigates how transformers perform in-context learning on categorical data. The core idea is to introduce a new attention-based architecture, built upon self-attention and cross-attention blocks, that enables exact implementation of multi-step gradient descent with categorical data. Experimental results and analysis look interesting. As pointed by all reviewers, the paper is dense and can be made more readable. As promised in the rebuttal, I recommend authors to incorporate the diagrams and provide more intuitive explanations through out the paper to help readers follow the work more easily. Overall, I am leaning towards accepting this work.
Decision: Accept (poster)