Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context

Xiang Cheng, Yuxin Chen, Suvrit Sra

TL;DR

This work shows that non-linear Transformers can learn to implement gradient descent in function space by performing functional gradient descent in RKHS via carefully configured attention modules. When the attention nonlinearity matches a generating kernel, the Transformer can converge to Bayes-optimal predictors in in-context learning, with multi-head attention enabling a broad class of composite kernels. The authors provide a constructive RKHS-based mechanism, theoretical results on stationary points under sparsity and full-attention regimes, and extensive simulations demonstrating the learned mechanisms and the dependence of optimality on the chosen nonlinearity. They also discuss the implications for representation power, including kernel composition and the potential for improved guarantees, suggesting practical and theoretical avenues for future research.

Abstract

Many neural network architectures are known to be Turing Complete, and can thus, in principle implement arbitrary algorithms. However, Transformers are unique in that they can implement gradient-based learning algorithms under simple parameter configurations. This paper provides theoretical and empirical evidence that (non-linear) Transformers naturally learn to implement gradient descent in function space, which in turn enable them to learn non-linear functions in context. Our results apply to a broad class of combinations of non-linear architectures and non-linear in-context learning tasks. Additionally, we show that the optimal choice of non-linear activation depends in a natural way on the class of functions that need to be learned.

Transformers Implement Functional Gradient Descent to Learn Non-Linear Functions In Context

TL;DR

This work shows that non-linear Transformers can learn to implement gradient descent in function space by performing functional gradient descent in RKHS via carefully configured attention modules. When the attention nonlinearity matches a generating kernel, the Transformer can converge to Bayes-optimal predictors in in-context learning, with multi-head attention enabling a broad class of composite kernels. The authors provide a constructive RKHS-based mechanism, theoretical results on stationary points under sparsity and full-attention regimes, and extensive simulations demonstrating the learned mechanisms and the dependence of optimality on the chosen nonlinearity. They also discuss the implications for representation power, including kernel composition and the potential for improved guarantees, suggesting practical and theoretical avenues for future research.

Abstract

Many neural network architectures are known to be Turing Complete, and can thus, in principle implement arbitrary algorithms. However, Transformers are unique in that they can implement gradient-based learning algorithms under simple parameter configurations. This paper provides theoretical and empirical evidence that (non-linear) Transformers naturally learn to implement gradient descent in function space, which in turn enable them to learn non-linear functions in context. Our results apply to a broad class of combinations of non-linear architectures and non-linear in-context learning tasks. Additionally, we show that the optimal choice of non-linear activation depends in a natural way on the class of functions that need to be learned.
Paper Structure (42 sections, 15 theorems, 145 equations, 5 figures)

This paper contains 42 sections, 15 theorems, 145 equations, 5 figures.

Table of Contents

  1. Introduction
  2. Main Contributions
  3. Related Work
  4. Setup: In-context learning with non-linear Transformers
  5. Input Data for In-Context Learning
  6. Transsformers with general non-linear attention
  7. The In-Context Loss
  8. Examples of Attention Modules
  9. Transformers can implement gradient descent in function space.
  10. Transformers can implement gradient descent in function space.
  11. Case Study: Linear Kernel
  12. Case Study: Exponential Kernel, and connection to Softmax activation
  13. Optimality of ${\tilde{h}}$ for matching $\mathcal{K}$.
  14. Experiments for Proposition \ref{['p:matching_h_k_optimality']}
  15. Composing multiple attention heads with different activations
  16. ...and 27 more sections

Key Result

Proposition 1

Let $\mathcal{K}$ be an arbitrary kernel. Let $\mathbb{H}$ denote the Reproducing Kernel Hilbert space induced by $\mathcal{K}$. Let ${z^{(i)}} = ({x^{(i)}},y^{(i)})$ for $i=1\ldots n$ be an arbitrary set of in-context examples. Denote the empirical loss functional by $L(f):= \sum_{i=1}^n \left( f({

Figures (5)

  • Figure 1: Plot of log(test ICL loss) against number of in-context demonstrations. The labels are generated using a $\mathcal{K}$ Gaussian Process (Definition \ref{['d:k_gaussian_process']}) Each sub-figure corresponds to one of three choices of $\mathcal{K}$, defined in \ref{['e:3_kernel_choices']}. Each sub-figure contains 4 plots corresponding to 4 choices of ${\tilde{h}}$, as defined in \ref{['e:4_kernel_choices']}. Black line denotes Bayes Loss.
  • Figure 2: Plot of log(test ICL loss) against number of layers. The labels are generated using a $\mathcal{K}$ Gaussian Process (Definition \ref{['d:k_gaussian_process']}), for $\mathcal{K}^{relu}$ and $\mathcal{K}^{exp}$ as defined in \ref{['e:3_kernel_choices']}. Each sub-figure contains 3 plots corresponding to three choices of ${\tilde{h}}$, as defined in \ref{['e:4_kernel_choices']}. The two plots in the top row have $n=14$ demonstrations. The two plots in the bottom row have $n=6$ demonstrations. Black line denotes Bayes Loss.
  • Figure 3: Plot of log(test ICL loss) against number of layers. Each sub-figure samples data from a different distribution ($\mathcal{K}^{\diamond}(u,v) := \alpha \mathcal{K}^{linear}(G_1 u, G_1 v) + (1-\alpha) \mathcal{K}^{exp}(G_2 u, G_2 v)$). We compare the performance of three kinds of Transformers. The labels are generated using a $\mathcal{K}^{\diamond}$ Gaussian Process. Context length $n=14$.
  • Figure 4: Plots of $\log({\sf dist}(M,I))$ for $M=\Sigma^{1/2}\left\{ B_0^\top C_0, B_1^\top C_1, B_2^\top C_2 \right\}\Sigma^{1/2}$ against number of training iterations. Each plot coincides with a different experiment setup, where we vary the generating distribution and the architecture. The subplot title is ($\mathcal{K}$, $\tilde{h}$), where $\mathcal{K}$ defines a Gaussian Process for labels, as described in Definition \ref{['d:k_gaussian_process']}, and ${\tilde{h}}$ is the non-linear map in the Transformer's attention module. In all cases, the corresponding matrix appears to be converging to identity, which is the stationary point from Theorem \ref{['t:informal_master_sparse']}.
  • Figure 5: Plots of $\log({\sf dist}(M,I))$ for $M \in \left\{ A_0, A_1 \right\} \cup \left\{ \Sigma^{1/2} B_i^\top C_i \Sigma^{1/2} \right\}_{i=0,1,2}$ against number of training iterations. Each plot coincides with a different experiment setup, where we vary the generating distribution and the architecture. The subplot title is ($\mathcal{K}$, $\tilde{h}$), where $\mathcal{K}$ defines a Gaussian Process for labels, as described in Definition \ref{['d:k_gaussian_process']}, and ${\tilde{h}}$ is the non-linear map in the Transformer's attention module. The definitions of each $\mathcal{K}$ and ${\tilde{h}}$ can be found in \ref{['e:3_kernel_choices']} and \ref{['e:4_kernel_choices']} respectively. In all cases, the corresponding matrix appears to be converging to identity, which is the stationary point from Theorem \ref{['t:informal_master_full']}.

Theorems & Definitions (35)

  • Remark 1
  • Example 2: ReLU Transformer
  • Example 3: Softmax Transformer
  • Proposition 1
  • proof : Proof of Proposition \ref{['p:rkhs_descent_transformer_construction']}
  • Remark 2
  • Definition 1: $\mathcal{K}$ Gaussian Process
  • Proposition 2
  • proof : Proof of Proposition \ref{['p:matching_h_k_optimality']}
  • Proposition 3
  • ...and 25 more