---
date: Thursday, December 5^th^, 2024
katex-css: true
subtitle: a simple and streamlined proof
title: Transformers are universal in-context function approximators
---

This is a short technical note adapted from a presentation I gave,
apropos of nothing, at the [metauni](https://metauni.org/) [SLT
seminar](https://metauni.org/slt/) on Thursday, December 5^th^, 2024.

In the seminar, I presented a simplified and streamlined version of the
results from section 2 of the paper "Vocabulary in-context learning in
transformers: Benefits of positional encoding," currently under public,
double-blind review at ICLR 2025 ([OpenReview
link](https://openreview.net/forum?id=YE6N8htoFQ)).

```{=tex}
\newcommand\Reals{\mathbb{R}}
\newcommand\Nats{\mathbb{N}}
\newcommand\eps{\varepsilon}
\newcommand\Tsoftmax{T_{\mathrm{softmax}}}
\newcommand\Tsigma{T_{\sigma}}
```
## Definitions: Transformers

A *transformer* is a function $$
  T
  :
  \left(
    \mathbb{R}^{d} \times \mathbb{R}^{d'}
  \right)^{\ast}
  \times
  \mathbb{R}^d
  \to 
  \mathbb{R}^{d'}
$$ where $\ast$ is the [Kleene
star](https://en.wikipedia.org/wiki/Kleene_star). The first input to
this function is a *context*---a sequence of pairs of vectors of a
variable, unbounded length $n \in \mathbb{N}$,
$(x_1, y_1), (x_2, y_2), \ldots, (x_n, y_n)$ where each
$x_i \in \mathbb{R}^{d}$ and each $y_i \in \mathbb{R}^{d'}$. It then
also takes a *query* $x \in \mathbb{R}^d$ and returns a *completion*
$y \in \mathbb{R}^{d'}$.

Given a particular context $(x_1, y_1), \ldots, (x_n, y_n)$, the
transformer [can be thought of](https://en.wikipedia.org/wiki/Currying)
as a function from queries to completions: $$
  T(x_1, y_1, \ldots, x_n, y_n, -)
  :
  \mathbb{R}^d
  \to 
  \mathbb{R}^{d'}.
$$ It is this function (modulo some auxiliary dimensions here and there)
that we want to show is a universal function approximator, in the sense
that for any continuous function on a compact domain we can find a
context such that the transformer as a function from queries to
completions is approximately equal to that function.

It may be unclear why we are considering a transformer that operates on
sequences of pairs of vectors, rather than a sequence of tokens.

-   First, imagine we are studying the part of the transformer between
    then embedding and the unembedding. That takes us from tokens to
    vectors.
-   Second, imagine we arbitrarily factor the residual stream of shape
    $\mathbb{R}^{d+d'}$ as $\mathbb{R}^{d} \times \mathbb{R}^{d'}$. That
    gives us pairs. Well... that's *what* we are doing. I don't actually
    know exactly *why* we want pairs, except to say it makes it simpler
    to talk about approximating a function, and that I think this is a
    common approach in theoretical work on transformer expressivity. I
    don't completely buy it, but let's roll with it for today.

With that said, let's fill out the transformer type signature with the
concrete mappings we will show are universal in-context function
approximators. Consider the following two "transformers":

1.  The *softmax attention module,* $T_{\mathrm{softmax}}$, with the
    mapping $$
      T_{\mathrm{softmax}}(x_1, y_1, \ldots, x_n, y_n, x)
      =
      \sum_{i=1}^{n} 
      \frac
        {\exp( x_i \cdot x )}
        {\sum_{j=1}^{n} \exp( x_j \cdot x )}
      \, y_i.
    $$

    This is a simplified version of a single attention head within a
    real transformer. See that the query $x$ is compared to previous
    $x_i$ values in the context, and these comparisons are used to
    create a softmax distribution. Then, the completion is an average of
    the $y_i$ values from the context weighted by this softmax
    distribution.

2.  The *elementwise attention module,* with the mapping $$
      T_{\sigma}(x_1, y_1, \ldots, x_n, y_n, x)
      =
      \sum_{i=1}^{n} \sigma( x_i \cdot x ) \, y_i
    $$ where $\sigma : \mathbb{R}\to \mathbb{R}$ is an arbitrary
    activation function.

    It's not so clear to me that this second option should be called a
    transformer. The completion is not generated as an average weighted
    by a normalised probability distribution, but by the outputs of an
    arbitrary activation function. But, it is a useful conceptual
    stepping stone for the proof in the softmax case---if you look
    closely you might be able to anticipate the way the context plays a
    similar role to the weights of a feed-forward neural network
    already.

Even the softmax attention module is substantially simplified compared
to a full modern transformer architecture, or even a single attention
head within such an architecture. The key missing piece is a triple of
parametric linear transformations (the *query, key,* and *value*
transformations). The authors consider some such parametrised
transformers but I have removed these parameters for this note because
they lead to more cumbersome notation while not being necessary to see
the key ideas in the proof.

## Universal (in-context) function approximation

Recall the following result from Leshno, Lin, Pinkus, and Schocken
(1993):

**Theorem** (Universal approximation, Leshno et al., 1993)**:** Let
$\sigma : \mathbb{R}\to \mathbb{R}$ be a non-polynomial, locally
bounded, piecewise continuous activation function. For any continuous
function $f : \mathbb{R}^d \to \mathbb{R}^{d'}$ defined on a compact
domain $\mathcal K \subset \mathbb{R}^d$, and any $\varepsilon> 0$,
there exists $n \in \mathbb{N}$, $a_1, \ldots, a_n \in \mathbb{R}^{d}$,
$b_1, \ldots, b_n \in \mathbb{R}$, and
$c_1, \ldots, c_n \in \mathbb{R}^{d'}$, such that for all
$x \in \mathcal{K}$ $$
  \left\lVert
    f(x)
    −
    \sum_{i=1}^n \sigma (a_i \cdot x + b_i)\,c_i 
  \right\rVert_\infty
  < \varepsilon.
$$

The "Vocabulary in-context learning" paper proves as a corollary of this
result a universal approximation property for the above two transformers
(and more complex, parameterised transformers) *as a function of the
context:*

**Corollary:** For any continuous function
$f : \mathbb{R}^d \to \mathbb{R}^{d'}$ defined on a compact domain
$\mathcal K \subset \mathbb{R}^d$, and any $\varepsilon> 0$, there
exists $n \in \mathbb{N}$, $x_1, \ldots, x_n \in \mathbb{R}^{d+1}$, and
$y_1, \ldots, y_n \in \mathbb{R}^{d'}$, such that for all
$x \in \mathcal{K}$ $$
  \left\lVert
    f(x)
    −
    T_{\mathrm{softmax}}\left(
      x_1, y_1, \ldots, x_n, y_n, (x, 1)
    \right)
  \right\rVert_\infty
  < \varepsilon.
$$ Moreover, the same is true for $T_{\sigma}$ if
$\sigma : \mathbb{R}\to \mathbb{R}$ is non-polynomial, locally bounded,
and piecewise continuous.

Note that compared to the function to be approximated
$f : \mathbb{R}^d \to \mathbb{R}^{d'}$, the context $x_i$ vectors have
an additional dimension, and the query is $(x, 1) \in \mathbb{R}^{d+1}$
where $x \in \mathbb{R}^d$ is the function input. These are the
"auxiliary dimensions here and there" to which I referred earlier.

I present the proof first for the case of the elementwise attention
($T_{\sigma}$) and then the slightly more complex proof for softmax
attention ($T_{\mathrm{softmax}}$).

## Elementwise case: Equivalence of NNs and contexts

*Proof* ($T_{\sigma}$ case): By the theorem there exists
$n \in \mathbb{N}$, $a_1, \ldots, a_n \in \mathbb{R}^{d}$,
$b_1, \ldots, b_n \in \mathbb{R}$, and
$c_1, \ldots, c_n \in \mathbb{R}^{d'}$ such that for all
$x \in \mathcal{K}$ $$
  \left\lVert
    f(x)
    −
    \sum_{i=1}^n \sigma (a_i \cdot x + b_i)\,c_i 
  \right\rVert_\infty
  < \varepsilon.
$$ For $i = 1, \ldots, n$, put $x_i = (a_i, b_i) \in \mathbb{R}^{d+1}$
and $y_i = c_i \in \mathbb{R}^{d'}$. Then $$
    T_{\sigma}\left(
      x_1, y_1, \ldots, x_n, y_n, (x, 1)
    \right)
    =
    \sum_{i=1}^{n} \sigma( x_i \cdot (x, 1) ) \, y_i
    = 
    \sum_{i=1}^{n} \sigma( a_i \cdot x + b_i ) \, c_i
$$ QED.

**Remark:** The above proof reveals an equivalence between:

(1) biased single-hidden-layer neural networks with $d$ input units, $n$
    biased hidden units, and $d'$ output units, and
(2) transformer contexts with length $n$ and dimensions $d+1$ and $d'$.

The paper shows this correspondence also exists for any *parameterised*
elementwise attention module, as long as the query, key, and value
transformations have a certain structure that is pretty unrestrictive
but not very easy to describe without introducing a lot of cumbersome
notation (and, unfortunately, also a bit buried in the paper).

## Softmax case: Approximation of NNs by contexts

*Proof* ($T_{\mathrm{softmax}}$ case): Unfortunately, normalisation
prevents a similar correspondence for softmax attention. But for the
corollary it suffices to approximate neural networks with contexts.

In particular, the function $\sigma = \exp$ is non-polynomial, locally
bounded, and piecewise continuous. Therefore by the theorem there exists
$n \in \mathbb{N}$, $a_1, \ldots, a_n \in \mathbb{R}^{d}$,
$b_1, \ldots, b_n \in \mathbb{R}$, and
$c_1, \ldots, c_n \in \mathbb{R}^{d'}$ such that for all
$x \in \mathcal{K}$ $$
  \left\lVert
    f(x)
    −
    \sum_{i=1}^n \exp (a_i \cdot x + b_i)\,c_i 
  \right\rVert_\infty
  < \frac\varepsilon 2.
$$ I will show there exists $x_1, \ldots, x_{n} \in \mathbb{R}^{d+1}$
and $y_1, \ldots, y_{n} \in \mathbb{R}^{d'}$ such that for all
$x \in \mathcal{K}$ $$
  \left\lVert
    \sum_{i=1}^{n} \exp (a_i \cdot x + b_i) \,c_i
    −
    T_{\mathrm{softmax}}(x_1, y_1, \ldots, x_n, y_n, (x, 1))
  \right\rVert_\infty
  < \frac\varepsilon 2,
\tag{$\star$}
$$ from which the corollary follows by the triangle inequality.

For $i = 1, \ldots, n$, put $x_i = (a_i, b_i - C) \in \mathbb{R}^{d+1}$,
and $y_i = \exp(C)\, c_i \in \mathbb{R}^{d'}$ where $C$ is to be defined
later.

Then $$\begin{align*}
  T_{\mathrm{softmax}}(x_1, y_1, \ldots, x_n, y_n, (x,1))
  &=
    \sum_{i=1}^{n}
    \frac
      {\exp( x_i \cdot (x,1) )}
      {\sum_{j=1}^{n} \exp( x_j \cdot (x,1) )}
    \, y_i
  \\[3ex]&=
    \frac
      {\sum_{i=1}^{n} \exp( a_i \cdot x + b_i - C )\exp(C)\, c_i}
      {\sum_{j=1}^{n} \exp( a_j \cdot x + b_j - C )}
  \\[3ex]&=
    \frac
      {\sum_{i=1}^{n} \exp( a_i \cdot x + b_i)\, c_i}
      {\sum_{j=1}^{n} \exp( a_j \cdot x + b_j - C )}
\end{align*}$$

Importantly, we have recovered the exponential neural network as the
numerator. We can now factor it out. $$
  (\text{LHS of }\star)
  =
  \left\lVert
    \sum_{i=1}^n \exp (a_i \cdot x + b_i)\,c_i 
  \right\rVert_\infty
  \left(
    1 - \frac{1}{\sum_{j=1}^{n} \exp (a_j \cdot x + b_j - C)}
  \right)
$$

Then, noting that for $z > 0$ we have $z > 1-1/z$, $$
  (\text{LHS of }\star)
  <
  \left\lVert
    \sum_{i=1}^n \exp (a_i \cdot x + b_i)\,c_i 
  \right\rVert_\infty
  \left(
    \sum_{j=1}^{n} \exp (a_j \cdot x + b_j - C)
  \right).
$$

To finish, take $C$ large enough so that for all $x \in \mathcal{K}$ and
$j \in \lbrace1,\ldots,n\rbrace$, $$
    \exp(a_j \cdot x + b_j - C)
    <
    \frac
      {\varepsilon}
      {2n\left\lVert
        \sum_{i=1}^n \exp (a_i \cdot x + b_i)\,c_i 
      \right\rVert_\infty}.
  $$ Such a $C$ is well-defined (finite) because $\mathcal{K}$ is
compact.

QED.

**Remark:** I think the above proof for the softmax attention case is
cleaner than the one in the paper (though I followed their approach).

-   This proof is substantially streamlined---in the paper, the authors
    use the triangle inequality twice, going from $f$ to
    $T_{\mathrm{softmax}}$ first via an exponential neural network (as
    above) but then also via a normalised 'softmax neural network' (the
    argument is split over two lemmas in the paper).
-   The use of the triangle inequality is also implicit in the paper's
    argument(s), whereas my personal style encouraged me to frontload
    this strategic information so that a reader knows what bound we are
    trying to achieve and why.
-   In each of the two steps, the authors introduce an auxiliary element
    into the approximating architecture: one extra unit in the softmax
    network to approximate the exponential network, and one extra
    context pair to approximate the softmax network with a context. The
    result is that the constructed context has length $n+2$ where $n$ is
    the width of the exponential neural network derived via Leshno et
    al. (1993). In contrast, the above proof constructs a length-$n$
    context.
-   The above proof corrects a minor flaw in the paper's proof, where
    the authors effectively define $C$ as a function of $x$, instead of
    constructing one prompt that achieves the required epsilon bound for
    all $x \in \mathcal{K}$ (as is to be demonstrated).

The proof in the paper is also more elaborate due to accounting for
query, key, and value transforms. However, this does not prevent my
streamlined approach---the above proof can easily be elaborated to
account for those.

## Finite vocabularies and positional encodings

This result is actually not the main focus of the paper. That would be
the following:

-   In section 3, the authors show that, for both elementwise attention
    modules and softmax attention modules, this universal approximation
    property does not hold if the context is restricted to be drawn from
    a finite alphabet (a "vocabulary"). By "finite alphabet," I mean
    contexts are confined to $(\mathcal{X} \times \mathcal{Y})^\ast$ for
    finite $\mathcal{X} \subset \mathbb{R}^{d}$ and
    $\mathcal{Y} \subset \mathbb{R}^{d'}$, as opposed to
    $(\mathbb{R}^d \times \mathbb{R}^{d'})^\ast$.

    I checked this proof carefully and I think it holds, though there
    are some minor errors remaining in the most recent revision.

-   In section 4, the authors show that, in the case of elementwise
    attention modules, the universal approximation property can be
    restored in the finite-alphabet setting if one considers contexts
    with a dense additive positional encoding. By "dense additive
    positional encoding," I mean they consider a sequence of vectors
    $p_1, p_2, \ldots \in \mathbb{R}^d$ such that
    $\{ x + p_i | x \in \mathcal{X}, i \in \mathbb{N}\}$ is dense in
    $\mathbb{R}^d$ (or some bounded subsets of $\mathbb{R}^d$ in the
    case of certain activation functions). Then, $x_i \in \mathcal{X}$
    in the context is replaced by $x_i + p_i$. They also assume
    $\{-1, 0, 1, \sqrt{2}\}^{d'} \subseteq \mathcal{Y}$.

    I haven't studied these proofs in detail, but from these
    assumptions, they could scan the positional encoding sequence until
    they find an arbitrarily close approximations of the $x_i$s in their
    previously-constructed real contexts, and likewise assemble
    arbitrarily close approximations of the $y_i$s by merging many such
    examples with different $y \in \mathcal{Y}$. Then since their
    architecture is continuous in the context this should suffice to
    achieve universality again.

To be honest, I don't think these results are very compelling, even
aside from the fact that universal approximation is not much more than
theoretically interesting.

-   The finite-alphabet setting is motivated by the fact that in a real
    language model, the input sequence is indeed assembled from a finite
    vocabulary of tokens. However, in a multi-layer language model, the
    outputs of the first layer aggregates information from multiple
    tokens, and through successive transformer blocks I conjecture that
    we can reach an exponential number of embedding vectors, bringing us
    closer to the dense setting.

-   As for the benefits of positional encoding, the assumption of
    density of the positional encoding is so strong as to substantially
    cheapen the result. I guess the proof involves some very long
    contexts to find the right rational approximations. Moreover,
    authors were only able to prove this result for their elementwise
    "transformer" architecture (I would guess that the long contexts
    creates problems via normalisation). Yet the authors force a
    positive framing, and even go as far as to suggest empirical
    research should explore the benefits of positional encodings that
    enumerate the rational numbers.

Will this paper be accepted? Two of the reviewers liked the paper (low
confidence accept, fair confidence weak accept). The other two reviewers
recommended weak reject, one (low confidence) is not on board with the
idea of optimising over contexts to begin with; the other (high
confidence, very in-depth review) is on board but not impressed with the
clarity of the presentation or the strength of the assumptions. I
predict this paper will not be accepted.

## Bibliography

Once again, this note is based on the following paper:

1.  Anonymous author(s), "Vocabulary in-context learning in
    transformers: Benefits of positional encoding," under review at
    ICLR 2025. [OpenReview](https://openreview.net/forum?id=YE6N8htoFQ).

The proof techniques build on foundational, classical work in universal
function approximation with feed-forward neural networks by Leshno et
al.:

2.  Moshe Leshno, Vladimir Ya. Lin, Allan Pinkus, and Shimon Schocken,
    "Multilayer feedforward networks with a nonpolynomial activation
    function can approximate any function," *Neural Networks,*
    6(6):861--867, 1993.
    DOI:[10.1016/S0893-6080(05)80131-5](https://doi.org/10.1016/S0893-6080(05)80131-5).

The observation that in-context learning is a universal approximator in
the above sense has apparently been made before, using different
techniques, according to the anonymous author(s):

3.  Aleksandar Petrov, Philip Torr, and Adel Bibi, "Prompting a
    pretrained transformer can be a universal approximator," at
    ICML 2024.
    [PMLR](https://proceedings.mlr.press/v235/petrov24a.html).