On Understanding Attention-Based In-Context Learning for Categorical Data

Thank you for your careful review and for your valuable feedback. On your questions:

1. "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). The 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).

2. 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)$.

3. 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$.

4. 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.

As recommended by you, we will add a figure in the main body of the paper, to help with readability. Please see here $LINK$ a draft figure (which we will seek to further improve) to aid readability and provide clarity.

Thank you for reviewing our paper. Concerning your comments:

The introduction ....

• lines 20-21 state that the goal of this paper is "extending the functional GD framework [for ICL], to handle categorical observations".
• Reviewers 4VEb and Hsdu both provided accurate summaries of the goal of the paper in the "Summary" portion of their review.

... disconnected/disparate components ....

• 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).

• 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)$.

• 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).

.... designed for explicit learning of this function. Is my understanding correct? ....

• This is incorrect. Our goal is also to learn $f_{\phi}$ implicitly.
• 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.
• 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))$.

The experimental setting ....

• 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.
  • 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.
  • 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.
  • 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".

... do not qualify for a theory paper ...

• Can reviewer GFKD please elaborate on why "the method description and analyses do not qualify for a theory paper"?
  • 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.
  • In Theorem 1 (also 2), we prove that our proposed construction is a stationary point of the training loss.
  • The strength of our results are no weaker than analogous results in earlier works that analyze multi-layer Transformers on real-valued data.

... (Ahn, Zhang) already demonstrates ...

• 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.
• 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.
• 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))$.

... some experimental results ...

• 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.

Thank you for your very careful review and helpful feedback.

We 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$ 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.

We will fix the typos and will implement your suggestions.

On 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.