Universal In-Context Approximation By Prompting Fully Recurrent Models
Authors: Aleksandar Petrov
Tom A. Lamb
Alasdair Paren
Philip Torr
Adel Bibi
Venue: NeurIPS 2024 poster (2024)
Keywords: prompting
universal approximation
in-context learning
recurrent models
rnn
ssm
Primary area: generative_models
TL;DR: We show that fully recurrent models, e.g., RNNs, LSTMs, GRUs and SSMs, can be universal in-context approximators

Abstract: Zero-shot and in-context learning enable solving tasks without model fine-tuning, making them essential for developing generative model solutions. Therefore, it is crucial to understand whether a pretrained model can be prompted to approximate any function, i.e., whether it is a universal in-context approximator. While it was recently shown that transformer models do possess this property, these results rely on their attention mechanism. Hence, these findings do not apply to fully recurrent architectures like RNNs, LSTMs, and the increasingly popular SSMs. We demonstrate that RNNs, LSTMs, GRUs, Linear RNNs, and linear gated architectures such as Mamba and Hawk/Griffin can also serve be universal in-context approximators. To streamline our argument, we introduce a programming language called LSRL that compiles to these fully recurrent architectures. LSRL may be of independent interest for further studies of fully recurrent models, such as constructing interpretability benchmarks. We also study the role of multiplicative gating and observe that architectures incorporating such gating (e.g., LSTMs, GRUs, Hawk/Griffin) can implement certain operations more stably, making them more viable candidates for practical in-context universal approximation.

Universal In-Context Approximation By Prompting Fully Recurrent Models
Aleksandar Petrov, Tom A. Lamb, Alasdair Paren, Philip H.S. Torr, Adel Bibi
Department of Engineering Science
University of Oxford aleks@robots.o ox. ac.uk
Abstract
Zero-shot and in-context learning enable solving tasks without model fine-tuning, making them essential for developing generative model solutions. Therefore, it is crucial to understand whether pretrained model can be prompted to approximate any function, i.c.. whether recently shown that transformer models do possess this property, these results rely on their attention mechanism. Hence, these findings do not apply fully recurrent architectures like RNNs, LSTMs, and the increasingly popular SSMs. We demon- strate that RNNs, LSTMs, GRUs, Linear RNNs, and linear gated architectures such as Mamba and Hawk/Griffin can also serve as universal in-context approximators. To streamline our argument, we introduce a programming language called LSRL that compiles to these fully recurrent architectures. LSRL may be of independent interest for further studies of fully recurrent models, such as constructing inter- pretability benchmarks. We also study the role of multiplicative gating and observe that architectures incorporating such gating (c.g., LSTMs, GRUs, Hawk/Griffin) can implement certain operations more stably, making them more viable candidates for practical in-context universal approximation.
universal in-context approximator. While it was
Introduction
Until recently, solving a task with machine learning required training or fine-tuning a model on a dataset matching the task hand. However, large foundation models exhibit the ability to solve new tasks without being specifically fine-tuned or trained for them: often it is sufficient to simply prompt them in the right way. Prompting has been especially successful because of in-context learning: the ability to modify the model's behavior with information provided within the input sequence, without changing the underlying model parameters (Brown et al., 2020). Yet, we know little about the theoretical properties of prompting. It is not even clear if there are limits to what can be achieved with prompting or, conversely, whether is possible prompt your way into any behaviour task.
This can be framed as a universal approximation question. Classically, universal approximation results show how class of tractable functions, such neural networks, approximates another class of concept functions, c.g., all continuous functions on a bounded domain, with arbitrary accuracy. This often done by showing that one can choose model parameters that approximate the target function. However, in-context learning poses different challenge as the model parameters are fixed. Instead, part of the input (the prompt) is modified to cause the model approximate the target function. Hence, we define universal in-context approximation to be the property that there exist fixed weights such that the resulting model can be prompted to approximate any function from a concept class. Understanding whether model can be universal in-contexi approximator is especially important as most commercial models are accessible exclusively via a prompting interface (La Malfa et al., 2023).
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

In-context learning has been almost exclusively studied in conjunction with the transformer archi- tecture (Vaswani et al.. 2017). This likely because in-comtext abilities uppear once the models are large enough (Wei et al., 2021) and most large models have been consformer-bssed On the subject of universal in-context approximation, Wang al. (2023) were first to show that transformer possesses this property by discretising and memorising all possible functions in the model weights. Memo- risation is not needed, though, and even small transformers can be universal approximators when prompted Petrov et al. (2024). Both results, however, critically depend on the attention mechanism of the transformer architecture (Bahdanau et al., 2015).
Still, generative models are not restricted to attention-basee architectures: there are the "classic" recurrent neural networks (RNNs, Amari, 1972), long short-term memory models (LSTMs, Hochreiter and Schmidhuber, 1997) and gated recurrent units (GRUs, Cho et al., 2014). Recently, Linear RNN models (also known as state-space models or SSMs) were proposed as scalable alternative to the transformer architecture (Orvieto et al., 2023; Fu et al., 2023a) and have started to outperform similarly-size transformers when multiplicative gating is added (Gu and Dao, 2023; De et al., 2024; Botev et al., 2024). Furthermore, despite in-context learning being associated with the transformer, recent empirical results show in-context learning in SSMs, RNNs, LSTMs and even convolutional models (Xie et al., 2022; Akyürek et al., 2024; Lee et al., 2024).
Yet, despite their ability to be in-context learners, there little known about the theoretical properties of these fully recurrent architectures. As these architectures become more and more widely used, understanding their in-context approximation abilities is increasingly more important for their safety, security and alignment. We show that, in fact, many of these architectures, similarly to transformers. can be universal in-context approximators. Concretely, our contributions are as follows:
i. We develop Linear State Recurrent Language (L.SRL): a programming language that com- piles to different fully recurrent models. Programming in LSRL is akin to "thinking like a recurrent model". LSRL programs can then be implemented exactly as model weights.
ii. Using LSRL, we construct Linear RNN models that can be prompted to act as any token- to-token function over finite token sequences, or to approximate any continuous function. These results also hold for RNNs, LSTMs, GRUs and Hawk/Griffin models (De et al., 2024).
iii. We present constructions with and without multiplicative gating. However, we observe that the constructions without these gates depend on numerically unstable conditional logic. iv. Nevertheless, we show that multiplicative gates lead to more compact and numerically stable models, making it more likely that universal in-context approximation properties arise in models utilising them, such LSTMs, GRUs and the latest generation of Linear RNNs.
2 Preliminaries
Fully recurrent architectures. In this work, we focus exclusively fully recurrent neural network architectures. Recurrent models operate over sequences. Concretely, consider an input sequence (xI...., æN) with 2, E X' X' being some input space. We will refer the elements of the input sequence as tokens even if they are real-valued vectors. A recurrent model : X - y maps a sequence of inputs to an output in some output space y. These models are always causal, namely:
y/ g( a.:).
(1)
We will abuse the notation and refer to (y1...., yr (9()))) g(...,,g(........ x()) as simply g( .... xx).
We will also separate the input sequence into query (qi,... qn) and prompt (PI.....Py) P.). The prompt specifies the target function ff that we approximate while the query designates the input at which we evaluate it. Contrary to the typical setting, we will place the query before the prompt.
There various neural network architectures that fall under the general framework of Eq. (1). The quintessential one is the RNN. It processes inputs one by one with only a non-linear state being passed from one time step to the other. A model g can thus be stacked RNN layers, each one being:
8; a (AAs + Bx, + b),
(Classic RNN)
y: ((s)).
'That necessitated by the limited capacity of the state variables. As the model fixed, order to increase the precision of the approximation, we can only increase the prompt length. If the prompt is before the query, it would have to be compressed into fixed-size state, limiting the approximation precision even with increased prompt lengths. But if the query has fixed size, it can be stored in Exed-size state variable exactly.

---

with A, B. b and the initial state value su being model parameters. o non-linear activation function and @ a multi-layer perceptron (MLP) with ReLU activations. We assume that always a ReLU to keep the analysis simpler. The non-linearity in the state update can make the model difficult to train (vanishing and exploding gradients, Bengio et al., 1994). Therefore, Linear RNNs have been proposed as regularizing the eigenvalues of A can stabilise the training dynamics (Orvieto et al., 2023). Linear RNNs also admit a convolutional representation, making them trainable in parallel (Gu et al., 2021; Fu et al., 2023a). Linear RNNs drop the non-linearity from the state update in Eq. (2):
8: = A8 + Bx; + b, y/ ((s)..
(Linear RNN)
The fully linear state updates do not affect the expressivity of the models, as non-linear activations are nevertheless present in the MLP layers e between the linear state update layers (Wang and Xue, 2023; Boyd and Chua, 1985). The state-of-the-ar Linear RNN models also utilise some form of multiplicatiie gating (Gu and Dao, 2023; De et al., 2024; Botev et al., 2024). While specific implementations can differ, we can abstract it as the following Gated Linear RNN architecture:
8₁-- A + Br, + b,
(Gated Linear RNN)
y: (x,,) ② ((s.),
with y being another MLP and O being the glement-wist multiplication operation (Hadamard product). Eq. (4) encompasses a range of recently proposed models. For example, one show that any model consisting of L stacked Gated Linear RNN layers, with y and o with k layers, can be represented as [(k+2)-layer Hawk Griffin model (De et al., 2024). The conversions are described in detail in App. E. We can similarly add multiplicative gating to the classic RNN architecture:
8, = (A8; + Bx, + b),
(Gated RNN)
(5)
y: = (xx)) @ ((s)),
Eq. (5) may appear unusual but is related to the well-known GRU (Cho et al., 2014) and LSTM (Hochreiter and Schmidhuber, 1997) architectures. Same as the case with Griffin/Hawk, any Gated RNN can be represented as a L. (k++2)-layer GRU or LSTM model (details Apps. C and D). As a result, if there exists Gated RNN model that is universal in in-context approximator (which we later show to be the case), then there also exist GRU and LSTM models with the same property.
Theoretical understanding of in-context learning. Beyond the question of universal in-context approximation, there have been attempts to theoretically understand in-context learning from various perspectives. The ability to learn linear functions and perform optimization in-context has been extensively explored in the context of linear regression (Garg al., 2022; Akyürek et al., 2022; von Oswald et al., 2023a; Fu et al., 2023b; Zhang et al., 2023; Ahn et al., 2023), kernel regression (Han et al., 2023) and dynamical systems (Li et al., 2023). Furthermore, studies have explored how in-context learning identifies and applies the appropriate pretraining skill (Xie et al., 2022:
Coda-Forno al., 2023; Bai al., 2023). It has also been shown that transformers can construet internal learning objectives and optimize them during the forward pass (von Oswald et al., 2023b; Dai et al., 2023). However, these studies almost exclusively focus on the transformer architecture, and the applicability of their findings to fully recurrent models remains unclear.
Approximation theory. Let X' and y be normed vector spaces. Take a set of functions C C y from X' y called concept space. Take also set of nicely behaved functions H C yx called hypothesis space. H. could be any set that we have tools to construct and analyse, e.g., all polynomials or all neural networks of a particular architectural type. Approximation theory is concerned with how well functions in H. approximate functions in C. We say that H. universally approximates C' over compact domain D (or that H dense in C) if for every ECC and <>0 there exist a hEH such that suP=ED [xx)) ( There is a long history of studying the concept class of continuous functions and hypothesis classes of single hidden layer neural networks (Cybenko, 1989; Barron, 1993) or deeper models (Hornik et al., 1989; Telgarsky, 2015). The concept class of sequence-to-sequeete functions has been shown to be universally approximated with the hypothesis classes of transformers (Yun et al., 2019), RNNs (Schäfer and Zimmermann, 2006) and Linear RNNs (Wang and Xue, 2023).
The hypothesis spaces in this work are different. The model is fixed and only the prompt part of the input is changed, i.c.. all learnable parameters are in the prompt. Take a recurrent model g as in Eq. (1) with fixed model parameters and a query length n. The hypothesis class is all functions that result by calling g on the user query followed by the prompt and taking the last n' outputs:
HD" (qq.........................((.. ..... .q.) + giq..... qn PI..... p.)..-' :] Vp, D,N 0}.
(6)
3

---

LSRL program
Original DAG
Simplified Path Graph
Compiled Linear RNN
ForEach:
input (nput(din=1)
ctris=LinState (input,
Deca 11:13
y
f-10 1 ReLL
-10
10
-1.1
A=[[1]]. B=[[1]].
100
init_state=((e1l)
100
Example input and output
equivalent to 1-input
Eate
is_zero = f_not (input))
Z mn
ctrês LinState( is_zero,
126
0
A=[[1]], B=[[1]].
18 init_state=((a1))
0
Yt
9h
11 equiv. f_step(ctris-etre
hitpu=ff_large
1E Seo ACD
Node types
13 ctris,ctr8s nu=18) output
IFRS 10
Twvut
LIEPAY
Compilation of an LSRL program to Linear RNN. An example of simple LSRL program that takes a sequence of and Is as an input and outputs there have been more Is than Os and 0 otherwise. The LSRL compiler follows the rules in App. A to simplify the computation DAG into a path graph. The resulting path graph can be represented as a Linear RNN with one layer.
Figure 1:
The domain D of p. and qi be continuous embeddings in R2 or discrete tokens V {1, V).
Note that each ke is identified by a prompt (PI,..,,N) but is a function with domain all possible queries (q),..., q.). Therefore, finding hypothesis hEHg that approximates a target function is equivalent to finding the prompt that hypothesis. The approximation properties of Hg Eq. (6) depend on the architecture of g. well as its specific parameters.
We study the recurrent architectures in Eqs. (2) to (5) and their ability approximate contin- uous functions over real-valued vectors and to represent discrete maps over tokens (which cor- responds how language models are used in practice). We consider the following classes of functions. Cvee_ (Te==[[[<w])001111 contains all continuous functions from the unit hypercube to Rda, while C²ak_ ( classes are 7000 (g) corresponding to Eq. (6) with D= n=n'=] and some fixed model of one of the four architectures in Eqs. (2) to (5), and Ho0k (g) with D=V and n=='=1
causal all causal functions from / tokens to / tokens. The hypothesis
Linear State Recurrent Language (LSRL)
We can construct the weights for universal in-context models with the architectures Eqs. (2) to (5) by hand but this is labour-intensive, error-prone, difficult to interpret, and the specific weights would be architecture-dependee
Working at such a low level of abstraction can also obfuscate common mechanisms and design patterns, making it more difficult to appreciate both the capabilities and the constraints of fully recurrent architectures. Instead, we propose a new programming language: Linear State Recurrent Language (LSRL).? LSRL programs compile Eqs. (2) to (5). Conversely, any Linear RNN can be represented as an LSRL program, making LSRL versatile tool for studying the capabilities of recurrent models. Later, in Secs, to 6 make use of LSRL. to develop programs that are universal approximators for (ave and (lok, thus showing that all four architectures can be universal in-context approximators.
the four architectures in
LSRL syntax. An LSRL program specifies how a single element is processed and how the recurrent states are updated for the next element. LSRL programs always start with an Input(æ) with an æ of fixed dimension. Only one Input can be declared in a program. Linear layers and ReLUs are also supported: Lin[A, b)(x) Ar + ReLU(2)) = max(0, x). The unique component of LSRL, however, is its LinState operation implementing the linear state update in Linear RNNs (Eq. (3)): LinState[A, B, b, so) (xx) == As + Br, + b, where the state s is the output of the call this node at step 1. LinState is the only way information can be passed from previous tokens the current one. We also provide a Concat operation that combines
Our implementation of LSRL is available https / thub.

---

ForEach:
input (dui:=llddd_irdddoo
counter needed knos whether looking query the prompt
const_1 (_constanttinnuut 1) center_vector Linstate(tnput=-cos::
"ones(d_in,d_in), B=ones(d_in.1) inis_state=zerossdd.kk copy the query in state (only when the counter is i) (cond senerfff_mlllerrcuuteeeeeeeee q LinState (inpatte(nttttqquuueee
q_update
1.5), t-input(: d_in], frinput(: d_in]se) siit.statezerros(d.in..)
Areye (reye(d_in),
the rollowing operations will only change the output when countor
18 t the step size is the rirst elexent of every pronpt element
11 step_size Linear (inputtintt-
Adonea(d_in,)), b=zaros(d_in,1))
conpute upper bounds 13 imput[1 : d_in]
12 using it
current pronst cell
14 ab lb step_size
15 * поw check if is in this cell (the busp should be on all dimensions)
16 ho__bun__coneeeette
= f_bump 1b, ub)
17 bump_.un -iinear(nputtq__.0 wttputooii.bupp.coppooo cow..un...opeeentiieee.... A=ones(1,d.in), braeroz(1.1))
18 in_cell
19 in m_ardd_rrooessing (buap_sun, d_in 9.5)
(.larger(buap.110 f_and(In_ f.larger 0.5])
20 t if counter>1 and this cell contains add the value to the cutput state
update r_ifelse(cand=r. (aadseccondrrrageerr (in (ln_and_preeessie .8.5), t=inputI-d_outt rsinputC-d_out: :J+=
22 y Linstate (Inputsupdate, 4=eye(d_out), s=eye(d_out), intt_state=zoroscd.ou...
23
return y
Listing 1: LSRL program for universal approximation in-context for continuous functions. The inputs are q q= Iq" 0 +1]+ with q' E 1] being the query value at which we want to evaluate the function, then followed by prompts describing the target function as in Eq. (8).
variables: Concat(a, y) == (I1. ..., 2 am) yi...., y v1). Finally, to support gating architectures we also implement rudimentary Multi operation that splits its input into two sub-arrays and returns their element-wise multiplication: Multi(x)) == x[ : lel/2] @ x[|ml/2: J. Naturally, Multi requires that æ has even length. These six operations can be composed into direct acyclic graph (DAG) with a single source node (the Input variable) and single sink node (marked with a return statement).
Such program operates over single token o; passed to Input, while recurrent model needs to operate over sequences. Thus, we wrap the program into a ForEach loop that passes each element individually for the DAG to output variable denoted by return clause. Each element is processed by the exact same program, with the only difference being that the state of the LinState variables is changing between iterations. You
an example of small LSRL program in Fig.
Expressiveness limitations. ForEach does not behave like the typical for loop: only the states are accessible between iterations, i.e., you cannot use the output of a linear layer at step t in any computation at step t + 1. Furthermore, the program is DAG and only states of LinState nodes are passed between iterations, variables computed in latter operations of a previous time step are not accessible inputs in earlier layers (with respect the topological sorting of the computation graph). This leads to a key programming paradigm in LSRL: a LinState update cannot depend non-linearly on its own state. That includes it depending on variable that depends the LinState itself and conditional updates to the state. Such dependency would break the DAG property of the program." This poses serious limitations on what algorithms can be expressed in a Linear RNN and makes programming them challenging, Still, in Sec. 4 we show how carefully constructing state updates and auxiliary variables can nevertheless allow to program some limited conditional behaviours.
Compilation. Any LSRL program without Multi nodes can be compiled a Linear RNN (Eq. (3)) or to a Gated Linear RNN (Eq. (4)). If the program has Multi nodes, then it cannot be compiled to Linear RNN as the multiplicative gating cannot be implemented exactly. However, it can be compiled to Gated Linear RNN. To compile an LSRL program to Linear (Gated) RNN, we first parse the program to build a computation graph. This is DAG with a single source (the Input node) and a single sink (the return statement of the ForEach loop). At the same time, Linear (Gated) RNN can be represented as a path graph (no branching) with the six basic operations as nodes. Therefore, the compilation step needs to transform this DAG into path graph. We achieve that by iterativly collapsing the first branching point into a single node. The exact rules that achieve that are described in App. A. Later, in Sec. 6, we will show how any Linear (Gated) RNN can be converted into non-linear (Gated) RNN, hence, how we can compile LSRL programs to these architectures as well.
Forr example, we cannot implement an operation that adds one to the state and squares it at each time step: set1 = (s: + 1)² operation that performs conditional assignment st+1 0 if (st > 5) else st.
5

---

Target function
Approximated function
x2)
p1 (1)
g( T2, 0,, 0,0)
(q1.92
1.0
10
0.0
0.0
22
0.0
0.0
3/2
11
0.0
1.0
States:
y: 0
5(2 3
f( D
Vodel
(:(0.0)
gr 22)
(n:4b)
(g.
Inputs:
i 0,0)
1..000 0G
18.3.0
18.0
Figure Intuition behind the LSRL program for universal in-context approximation for continuous functions in Lst. 1. Our target function f has input dimension dm = 2 and output
dimension dou = 1. Each input dimension is split into two parts, hence 1/2. We illustrated an example input sequence of length 5: one for the query and four for the prompt tokens corresponding to each the discretisation cells. The query (q144)) falls in the cell corresponding to the third prompt token. We show how the two LinState variables in the program are updated after each step. Most notably, how the state holding the output y is updated after p3 is processed.
Syntactic sugar. To make programming easier, we define several convenience functions. For instance, we can Slice variables az [l:u] via sparse Lin layers. We can also sum variables and element- wise multiplication with scalars (implemented Lin layers). For logical operations also need step functions which can be approximated with ReLUs: f f_step[[u[x>) == ReLU(ua)) - pReLU( - 1/p). where yu is positive constant controlling the quality of the approximation. We can also approximate bump functions (1 between 7 and u and O otherwise): f f_bump(t. u, xa(())= == f_step[[u( 1) f_ f_step]u] (z u). Similarly, we can approximate conjunction (f ___ndd, disjunction (f_or), negation (f_not), and comparison operators (f faaaage and f_smaller). See App. F for the definitions.
Critically, we need also conditional operator that assigns value t(æ: if certain condition is met and another value f(x) otherwise. One way to implement this
f f_ifelse(cond, t, (() = ReLU(-X cond()))f() + ReLUU)/ALLLAAALLoot((((o)))ta)))
(BLL((AAood(z))--) - ReLU( - (nno((coodo=)))t(aaa/ (7)
where À is constant that is larger than any absolute value that t(æ) and f(x) can attain. This construction, however, is not numerically stable and we will study alternatives in Sec. We provide both numerical (SciPy. sparse, Virtanen et al. 2020) and symbolic (SymPy, Meurer al. 2017) backends with the second being crucial for programs that are not numerically stable.
Prior work encoding algorithms in model weights. A similar approach to developing a pro- gramming language that compiles to model weights was already done for the transformer architecture with the RASP language (Weiss al., 2021) and the Tracr compiler (Lindner al., 2023). They were predominantly created as tool for interpretability research. In sense, RASP is to transformer as LSRL (Linear) (Gated) RNN. Hence, can be used to develop benchmarks for interpretabil- ity methods for fully-recurrent architectures. However, while RASP can only express a subset of transformer models, LSRL isomorphic the set of all (Gated) Linear RNNs (though not to the non-linear ones). That means that any (Gated) Linear RNN can be represented and analysed as an LSRL program and vice versa. Hence, the limitations of what you can express in LSRL are also limitations of what Linear (Gated) RNN can do. Namely: (i) we cannot have exact multiplicative interactions between inputs without multiplicative gates, and (ii) we cannot have state variable updates depending non-linearly their previous iterations or in any way variable that depends them.

---

discadeed outbuEs
globsl_cte
prompt
Figure 3: Intuition behind the LSRL program for universal in-context approximation for discrete functions in Lst. 2. Our keys and values have length n=3 and represent countries and capitals, e.g., AUStria-VIEnna, BULgaaria>SOFia, and so on. The query is CAN for Canada and the final n outputs are OTT (Ottawa). We show the values of some of the variables in Lst. 2 at each step. with the LinState variables being marked with arrows. For cleaner presentation are tokenizing letters as 0-+?, 1 A 2-+B, etc. Vertical separators are for illustration purposes only.
4 Universal In-Context Approximation with Linear RNNs
We proceed with building LSRL programs that are universal in-context approximators: one for approximating continuous functions (C) and one for maps between token sequences (Clok).
4.1
Approximating continuous functions in C
The idea behind the approximation for continuous functions is to discretise the domain into grid and approximate the function as constant in each cell of the grid. This technique is commonly used for showing universal approximation using the step activation function (Blum and Li, 1991; Scarselli and Tsoi, 1998). However, is not obvious how to implement this approach in-context when information across input tokens can be only combined linearly. Consider target funetion f [o, 1] [0, 1 Pou and a discretization step 8. Our approach is to describe the value of ʃ in each of the discretization cells single prompt token. For the cell with lower bounds , ld, and their respective upper bounds I +o.,.. Id, +6, the corresponding prompt token is a (cin 1 vector:
p = [6, 14.11....ee Id.. Fi.....dd..!'
(8)
where ÿ is the value of f at the centre of that cell: ÿ = [11 Ia, +5/2). Each prompt token describes the size of the cell (the discretisation step 6), its starting lower bound, and the value of the target function at the centre of the cell. Thus, [3/8] du such tokens, one for each cell, sufficient to describe the piece-wise constant approximation of f. A query q' E [0. 1j can fall in only one of the cells. We pad it with zeros and encode it as the first input element: q [q 0 +1 followed by the prompt. Our program will extract and save q' to a state and then process the prompt tokens one time until it finds the one whose cell contains q'. The target function value for this cell will be added to an accumulator state. If the current cell does not contain q', then 0 is instead added. Hence, the accumulator's 's final value corresponds to the value of f at the centre of the cell containing q' The full LSRL program is provided in Lst. 1 and an illustration for din 2, deut 1,,8 1/2 is shown in Fig. The prompt length required to approximate an I--Lipschitz function I (w.r.tt the &g norm) to precision is N = (2/LLvT.) die (le (see App. B for the proof). Asymptotically, this is as good as one can hope without further assumptions on the target function. This is also better than the best known result for the same problem for transformers: O -00 10-14d,--d6; -df) in Petrov et al. 2024.
4.2 Approximating functions over token sequences in Ctok
Sec. 4.1 focused on continuous functions but recurrent architectures are often used model natural language whose domain is tokens. Thus, we also look at modelling maps over discrete domain. Any function from n tokens to n tokens taking values in V = .l, V) can be represented as a dictionary whose keys and values are in y". Therefore, a simple way to represent this function in-context is to first provide the n tokens corresponding to the query and then a sequence of 2n tokens corresponding to key and value pairs (see Fig. 3 for an illustration of the setup). The model stores the
7

---

ForEach:
Input = (ngut(din=1)
counter needed knos whether looking the query the prompt
const_1 [_constantiinnut,
kiobal_ctr ctr [Lnnsate(Amput=ooostt. A=[[1]], a-(())) sttt==
* counters nod[n] and and[2n]
rod_n_ctr f_
(input, n)
vod_2n_ctr
2*n)
which mode are ие in (looking at the query, comparing query with key, or copying value to stace)
18 is.pronpt [.larger(slebal_ttr n-8.5)
11 s_coopprr__ddr =__aarge (rod_2n_ctr n-8.5)
12 1s_copy_nede r_and(1s_prrrpt, [_not(1s_corpare_mog
13 fs_firsttttken__orrr
f_and(1s_copy_nod) wowey..ddd,,...oo f.aog_nodd,..ssleeeeeeee F_analler(toddnn.l 8.5))
14 upcate the state holding the query if this is one of the first tokens
15 @g=f_ffeleeff__ff wes.lleerrrssprror... 9.5), t=input. f=inoutag)
16 (ff_iifelleef.addfff.. lnrgerrnoo..n.err 11-8.)) f.smaller 11...))....8.88. frggoe)) for 1..0] (n)y (eeye(n), intt_state-zerost. a query
17 (Lnntaateinuutttoooaaaat
18 if we are in conpare node (looking at keys), check if this token natches the correspending one in the query
19 (s=t__ifelsef__nn((( (f (argerrmmd___.trr.
Laedllerr (+8.5)) t- Feq(i+44)) for I in 1..n]
28 cor_q.el (inpart (opputtooott(ss)) Arones(),n), (rxeros(1,1))
satching (f.nd (ingut, cor-9_el-8.53 f. (input, cor-q_el-8.3 +))),, is_cooparee_noe))
23 buffar (Inpaatte ne1 natch values. the because we can only read butfer after writing Aeshift_ratris natrix, B-[[A] for in I..nd, [EI1]), init_atate-zerros 24 buffer_sue= Linear (inuutttnutsuufreer A=[[1 for in I..ol, [o]], bizeros(1, 1))
22 keep buffer
all_natching [_Larger
if all matching
it's first token the value part of the (key, value) pair, then nork this the iterations when не start copying to state
27 natching, firtt,ftt.
f_and natching, is.first_token, foren__ 28 t_started_or_ und_tt f_iffelsemaacchins.. wtt.annddstttrrreeeoo
teglobal_ctr, (=xlobal_ctrae) 29 started_on Linstate sttuuttttstttttt.ttmm on_update, AFeye(1), Bleye(1), sttt_staate,eeoac(, 1))
38 copylng to state
iterations after started_on.
31 copy_and_on = (sppppseop__ddee f_saller aetaaalsse, started_an-n))
32 rod_n_eq_i
(f (redd cned i--.5, [ff_and (topyaaccoy,,,nd..o.. rod_r_eq_i[[1]) for L_ppéates_should_updee (fffelse(f_.arger(tt. updates. sto.sshhuud.uuat((ii e.5), t-input, f-inputsd) for i in 1..n]
for 1..n
34 t_updates
35 output_regs + (Linstate (input- Areye(1), (reye(1), init_state
for update in [_uddates]
36 + rrinally read out the value fron the correspording output register in order to output fros the model 8,5), reoutbut_regelt). reoutput_ return Linear (input-Concnt (tooupputs), (rones(1,n) bex1ros(1,1))
t_outouss- dfffelse(f_largerrnddd.oo
in I..n]
Listing 2: LSRL program for universal in-context approximation of discrete functions. The inputs are q1,..., In (the query tokens), followed by pairs of keys and values from the map we are approximating. The last n outputs are the value corresponding to the key matching the query.
query state and processes the key-value pairs one by one by comparing the key (the first n tokens) with the query. If they match, then the value (the next tokens) is copied into a state that keeps it and repeatedly outputs it. This continues until the end of the prompt, at which point the last n outputted tokens will be the value corresponding to the key matching the query. This is essentially a dictionary lookup. However, as shown in Lst. 2, implementing dictionary lookup in a linear recurrent model is much less straightforward than executing di dicy] in general- purpose programming language.
Lst. 2 can appear daunting at first so we would like to elarify the non-trivial aspects. First, we need to count how far we are into every set of or 2n tokens. This can be done with mod n and mod 2n operations but implementing modulo for arbitrary large inputs is not possible with ReLU MLPs (Ziyin et al., 2020). Therefore, we implement this with LinState as f_mod_ll_ooooo _oolo__ccceerrrr which has a unit-length state that rotated 1/2 or 1/2m revolutions per iteration, with the angle corresponding to the modulo value (App. F.7). Second, we need to do dynamic indexing to copy the i-th input in a subsequence to the i-th element of state and vice-versa. Dynamic indexing, however, cannot be succinetly represented in a Linear RNN. We work around this with temporary variables that are non-zero only at the i-th coordinates (see Lines 16, 17, 19, 20. 32 to 35, 37 and 38). Finally, in order to compare whether all elements in the query and the key match, need to remember whether the previous n pairs were matching. As RNNs do not have attention, we implement this short-term memory buffer as LinState with shift matrix (Line 23).
5 Stable Universal In- In-Context Approximation with Gated Linear RNNs
The ReL.U- conditional operator is not numerically stable. The LSRL programs in Lsts. 1 and 2 for approximating functions in respectively Ch and Crok rely on the f f_ifelse conditional assignment operator in Eq. (7) in order to implement different behaviours depending whether
8

---

Continous functions
Functions over token sequences
10
100
Originl
10"
Optrreed
Exp- ieptriiad
Average difference
10
Mulippiaazze
Wrong tokers [M
10"
10
10"
10"
10
10
10
No noien
10"
10"
Amount noise added to the parametters
Amount of noise added to the parameters
Figure 4: Robustness of the various f_ ffifelse implementations to model parameter noise. We show how the performance of the two universal approximation programs Lsts. and 2 deteriorates as we add Gaussian noise of various magnitudes to the non-zero weights of the resulting compiled models. As expected, the original f_ifelse implementation in Eq. (7) exhibits numerical precision errors at the lowest noise magnitude. For the token sequence case, numerical precision errors are present in all samples even in the no-noise setting. Hence, the original f f_ifelse implementation is less numerically robust while the implementations with multiplicative gating are the most robust. For Lst. (approximating ( we report the Euclidean distance between the target function value and the estimated one over 10 queries for 25 target functions. For Lst. 2 we report the percentage of wrong token predictions over 5 queries for 25 dictionary maps. Lower values are better in both cases.
are processing the query specific parts of the prompt. This operator is not numerically stable. The first term in Eq. (7) relies on cond(æ) being exactly zero if the condition is not met. In this way, multiplying with -À would be O and f(æ:) would be returned. However, if cond(x:) is not identically 0 but has a small positive value, then -Acond(ar) can "overpower" f(æ)) resulting in the ReLU output being O. In our experience, this is not a problem when processing inputs through the LSRL program step-by step-by-step. However, de-branching the DAG into a path graph -which is necessary in order to uncover the equivalent Linear RNN- appears introduce such numerical instabilities which occasionally result in wrong outputs as conditional assignments will be O when they should not. This problem is more prominent in Lst. 2 which is longer (more debranching steps) and has more f f_ifelse operations: it gets most tokens wrong because of that instability (see Original, No noise in Fig. 4). To this end, we support LSRL with a symbolic backend (based on SymPy) that performs the debranching steps exactly. Using it, both programs always produce the correct output.
This numerical instability highlights critical practical limitations of the universal approximation results in Sec. 4: if the models are not numerically stable, it is unlikely that they occur in practice by training models using gradient descent. This section shows how to improve the numerical stability of Eq. (7) and obtain more realistic recurrent models that are universal approximators in-context.
Removing unnecessary terms Eq. (7). Eq. (7) has separate ReLU terms. The first two handle the cases when t(æ) and f(x) are positive and the second two when they are negative. Therefore, if we know that one or both of these will always be non-negative, we can drop the corresponding terms. Additionally, if f(x) is always O, then the first and third terms can be safely dropped. Similarly, the second and fourth are unnecessary fx) = 0. All f_ifelse in Lsts. 1 and 2 fall in this case and hence can be simplified. We will refer to this f_ifelse implementation that is aware of the attainable values of t(x) and f(x) as optimized As reduces the number numerically unstable ReLU operations in the model, expect that it will improve the stability of the compiled models. We experimented with adding various levels of noise to the non-zero model parameters, and, as the results in Fig. 4 show, aptimized is indeed more numerically robust than original.
Step-based implementation.
We can get rid of the input sensitivity of Eq. (7) using f_step:
f_ f_ifelse[cond, t, f, f)(r)) = ReLU(-A+Af ReUUK--NAALssto((//-omdd)))))))) + BelUl--AAAf_sseep(oomo/-///////
(9)
ReLU(->++AF_ssee////=ooodd
ReLU(-X=Af_steee MeUU--AAAf__to(((omm())))))--1
We also apply the optimisation strategy here. While this implementation is robust to noise in the input it appears to be more sensitive to parameter noise, as shown in Fig. 4.

---

Numerically stable f_ifelse with multiplicative gates. Removing the unused ReLU terms in the original f f_ifelse reduces the opportunities for numerical precision issues to creep but does not solve the underlying problem. The multiplicative gating present in the Linear Gated RNN (Fq. (4)) and Gated RNN models (Eq. (5)) can help via implementing numerically stable conditional operator:
f f_ifelse|comme t, fa( := cond(x) @ t(z) + f_not(cond(z)) O f(æ),
(10)
where the element- wise product is implemented in LSRL with Concat and Multi. We will refer to the implementation of f_ifelse in Eq. (10) as mul multiplicative Similarly to original implementation of f_ifelse in Eq. (7), we can drop the t(x) and f(z) term if they are equal to zero (mul wolltpplcatiiii optimized). If (a)d(a:) is not exactly zero, cond(a)) @ t(x:) will result in small error the output but, in contrast to the original implementation, is not going to cause a discontinuity in the output of the operation. Therefore, Eq. (10) should be more robust to numerical precision issues than Eq. (7). Fig. 4 shows that this is the case in practice with Lsts. I and 2 being more robust to parameter noise when using multiplicative gates compared to the ReLU-based implementations. Therefore, Linear Gated RNNs (Eqq (4)) -to which models with multiplicative gates can be compiled- are more likely than Linear RNNs (Eq. (3)) to exhibit universal approximation properties in practice.
6 Universal In-context Approximation with Non-linear (Gated) RNNs
Secs. 4 and 5 showed how universal approximation of continuous and token-to-token functions can be implemented in LSRL and compiled to respectively Linear RNNs and Linear Gated RNNs. This section aims to address the situation with non-linear state updates, that is, the cases of classic and gated RNNs (Eqs. (2) and (5)). Concretely, we show how every linear (Gated) RNN can be converted to a non-linear (Gated) RNN. The key idea is that the ReLU applied to the state updates in the non-linear architectures is an identity operation its inputs are positive. Hence, we can split the states in positive and negative components, flip the sign of the negative component, pass them separately through the ReLU-which will act as an identity as all elements will be non-negative- and then fuse the positive and negative components back together in the A matrix at the next time step:
(1
A
A
B
6
$1 As + Ba, +b y: ===00 (s:).
8,
ReLU
A
A
8
+
B
2
(1I)
y: s
[I
I]
s
sz
Using Eq. (11) we can compile any LSRL program to an RNN (Eq. (2)) or a Gated RNN (Eq. (5)). This includes Lsts. and 2. Hence, RNNs and Gated RNNs can be universal in-context approximators for continuous and token-to-token functions. As any Gated RNN can be represented as GRU model
(App. C) or LSTM (App. D), these models are too universal in-context approximators.
7 Discussion and Conclusions
We developed LSRL: a programming language for specifying programs expressible with recurrent neural architectures. We then used LSRL to show that various architectures -from the humble RNN to the state-of-the oh-the-art Linear Gated RNNs can all be universal approximators in-context.
Safety and security implications. If a model can be prompted to approximate any function, then preventing it from exhibiting undesirable behaviours (i.c., alignment) might be impossible. Therefore, it is important to further study the safety and security implications of these properties.
Limitations. In this work we provide constructive existence results: that is, we show that there can exist models with various recurrent architectures that are universal in-context approximators. However, the present theory is not sufficient to analyse whether a given model has this property. That is a much more difficult question that would require a very different approach. We also assume no restrictions on the matrix in the state update equations. However, many state-of-the-an models impose structural constraints on A (e.g., it being diagonal) for the sake of fast training and inference (Gu et al., 2020, 2021: Gupta et al., 2022). is not directly obvious whether such structural restrictions would affect the universal in-context approximation abilities of these architectures. In practice, however, the compiled matrices are very sparse and often diagonal. Therefore, is highly likely that our results translate to models with structural restrictions.
10

---

Acknowledgements
We would like thank Simon Schug for pointing relevant works fully recurrent models. We are also extremely grateful to Juuso Haavisto for his insight on building compilers. This work supported by UKRI grant Turing AI Fellowship (EP/W002981/1) and the EPSRC Centre for Doctoral Training in Autonomous Intelligent Machines and Systems (EP/S024050/1), Adel Bibi acknowledges funding from the KAUST Office of Sponsored Research (OSR-CRG2)) -4648) and support from Google Cloud through the Google Gemma 2 Academic Program GCP Credit Award.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. 2023, Transformers learn to imple- ment preconditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. 2022. What learning algorithm is in-context learning? Investigations with linear models. In The Eleventh International Conference on Learning Representations.
Ekin Akyürek. Bailin Wang, Yoon Kim, and Jacob Andreas. 2024, In-context language learning: Arhitectures and algorithms. arXiv preprint arXiv:2401.12973
Shun-ichi Amari. 1972. Learning patterns and pattern sequences by self-organizing nets of threshold elements. IEEE Transactions on Computers, C-21(11):1197-1206
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 2015. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. 2023, Transformers as statisticians: Provable in-context learning with in-context algorithm selection. In Advances in neural information processing systems.
Andrew Barron. 1993. Universal approximation bounds for superpositions of sigmoidal function. IEEE Transactions on Information Theory, 39(3):930-99) 945.
Iz Beltagy, Matthew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv preprint 15X
Yoshua Bengio, Patrice Simard, and Paolo Frasconi. 1994. Learning long-term dependencies with gradient descent difficult. JEEE Transactions on Neural Networks, 5(2):
Edward K Blum and Leong Kwan Li. 1991. Approximation theory and feedforward networks. Neural networks, 4(4):511-515.
Aleksandar Botev, Soham De. Samuel L Smith, Anushan Fernando, George-Cristiand Cristian Muraru, Ruba Haroun, Leonard Berrada, Razvan Pascanu, Pier Giuseppe Sessa, Robert Dadashi, et al. 2024. RecurrentGemma: Moving past transformers for efficient open language models. arXiv preprint arXiv:2404.07839.
Stephen Boyd and Leon Chua. 1985. Fading memory and the problem of approximating nonlinear operators with Volterra series. IEEE Transactions on Circuits and Systems, 32(1 1):1 1:115011166 161.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. In Advances in Neural Information Processing Systems.
Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcchre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. 2014. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP).
Julian Coda-Forno, Marcel Binz, Zeynep Akata, Matt Botvinick, Jane Wang, and Eric Schulz. 2023. iett-inn-contett learning in large language models. In Advances in Neural Information Processing Systems, pages

---

George Cybenko. 1989. Approximation by superpositions of sigmoidal function. Mathematics of control, signals and systems, 303)3303-314
Damai Dai, Yutao Sun, Li Dong, Yaru Hao, Shuming Ma, Zhifang Sui, and Furu Wei. 2023. Why can GPT learn in-context? Language models secretly perform gradient descent meta-optimizers. In Findings of the Association for Computational Linguistics: ACL 2023.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murant Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. 2024. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402.1 194227777.
Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. 2023a. Hungry Hungry Hippos: Towards language modeling with state space models. In International Conference on Learning Representations.
Deqing Fu, Tian-Qi Chen, Robin Jia, and Vatsal Sharan. 2023b. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. arXiv preprint arXiv:2310.17086.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. 2022. What can transformers learn in-context? A case study of simple function classes. In Advances in Neural Information Processing Systems.
Felix A Gers, Jürgen Schmidhuber, and Fred Cummins, 2000. Learning to forget: Continual prediction with Istm. Neural computation, 12(10)2441-
Albert Gu and Tri Dao. 2023. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. 2020. HiPPO: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems.
Albert Gu, Karan Goel, and Christopher Re. 2021. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations.
Ankit Gupta, Albert Gu, and Jonathan Berant. 2022. structured state spaces. In Advances Neural Information Processing Systems.
Diagonal state spaces are as effective as
Chi Han, Ziqi Wang, Han Zhao, and Heng Ji. 2023. In-context learning of large language models explained as kernel regression. arXiv preprint arXiv:2305.12766
Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural Computation, 9835
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. 1989. Multilayer feedforward networks are universal approximators. Neural networks, 2(5):359- 366.
Emanuele La Malfa, Aleksandar Petrov, Simon Frieder, Christoph Weinhuber, Ryan Burnell, Raza Nazar, Anthony G. Cohn, Nigel Shadbolt, and Michael Wooldridge. 2023, Language Models as Service: Overview of a new paradigm and its challenges. arXiv preprint arXiv:2309.16573
Ivan Lee, Nan Jiang, and Taylor Kirk-Kiikpatiick. 2024. Exploring the relationship between model architecture and in-context learning ability. In International Conference Learning Representations.
Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. 2023, Trans- formers algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning.
David Lindner, János Kramár, Sebastian Farquhar, Matthew Rahtz, Tom McGrath, and Vladimir Mikulik. 2023. Tracr: Compiled transformers as a laboratory for interpretability. Advances in Neural Information Processing Systems.
12

---

Aaron Meurer, Christopher P. Smith, Mateusz Paprocki, Ondřej Certík. Sergey B. Kirpichev, Matthew Rocklin, AMiT Kumar. Sergiu Ivanov, Jason K. Moore. Sartaj Singh, Thilina Rathnayake, Sean Vig, Brian E. Granger, Richard P. Muller, Francesco Bonazzi, Harsh Gupta, Shivam Vats, Fredrik Johansson, Fabian Pedregosa, Matthew J. Curry, Andy R. Terrel, Stěpán Roučka, Ashutosh Saboo, Isuru Fernando, Sumith Kulal, Robert Cimrman, and Anthony Scopatz. 2017. SymPy: Symbolic computing in Python. PeerJ Computer Science, 3.
Antonio Orvicto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. 2023, Resurrecting recurrent neural networks for long sequences. International Conference Machine Learning.
Aleksandar Petrov, Philip HS Torr, and Adel Bibi. 2024, Prompting pretrained transformer can be universal approximator. In International Conference on Machine Learning.
Franco Scarselli and Ah Chung Tsoi. 1998. Universal approximation using feedforward neural networks: survey of some existing methods, and some new results. Neural networks, 11(1):15- 37.
Anton Maximilian Schäfer and Hans Georg Zimmermann. 2006. Recurrent neural networks are uni- versal approximators. In Artificial Neural Networks-ICANI 2006: 16th International Conference, Athens, Greece, September 10-14, 2006. Proceedings, Part 16, pages
Springer.
Noam Shazeer. 2019. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150.
Matus Telgarsky. 2015. Representation benefits of deep feedforward networks. arXiv preprint 10Xiv:1509.08001.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017, Attention is all you need. In Advances in Neural Information Processing Systems.
Pauli Virtanen, Ralf Gommers, Travis E. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C Carey, Îlhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perktold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. 2020. SciPy 1.0: Fundamental algorithms for scientific computing in Python. Nature Methods, 17:261-272.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. 2023a. Transformers learn in-context by gradient descent. In International Conference on Machine Learning.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zuc- chet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pascanu, and João Sacramento. 2023b. Uncovering mesa-optimization algorithms in transformers. arXiv preprint 05858
Shida Wang and Beichen Xue. 2023. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. In Advances in Neural Information Processing Systems.
Yihan Wang, Jatin Chauhan, Wei Wang, and Cho-Jui Hsieh. 2023. Universality and limitations of prompt tuning. Advances in Neural Information Processing Systems.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu. Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. 2021. Finetuned language models are zero-shot learners. In International Conference on Learning Representations.
Gail Weiss, Yoav Goldberg, and Eran Yahav. 2021. Thinking like transformers. In International Conference Machine Learning.
13

---

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2022. An explanation of in-context learning as implicit Bayesian inference. In International Conference Learning Representations.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. 2019. Are transformers universal approximators of eequence-to-seqeene functions? In International Conference Learning Representations.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. 2023. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.0992)
Liu Ziyin, Tilman Hartwig, and Masahito Ueda. 2020. Neural networks fail to learn periodic functions and how fix it. In Advances Neural Information Processing Systems.
14

---

A Computation Graph Debranching Rules
We convert the computation DAG resulting from the LSRL program into path program by attending the first node whose output the input for multiple other nodes, iee, the first branching node.
Preparation step. Before we even start debranching we first pre-process the graph by fusing consecutive nodes of the same type together. The specific rules are:
If Lin node is followed single other Lin node, then fuse them together. This follows directly from the classical result that composing linear functions is a linear function.
If a ReLU node is followed by another ReLU node, we can drop one of them as ReLU is idempotent.
If a Lin followed by LinState, we can subsume the weight matrix A of the linear node in the B matrix of the LinState, and the bias b of the Lin node in the bias b of the LinState.
If all inputs of Concat node are the same, then this node only duplicates the input and hence can be safely replaced with a Lin layer.
The debranching process goes through the following cases in order. And iterates until there are no branching nodes left, in other words, until the graph has become path graph. We will refer to the nodes whose input is the branching node as subsequent nodes.
Case 1A: If all subsequent nodes are Multi. As all Multi nodes that have the same input (the branching node) they must all be producing the exact same output. Hence, only one can be kept. This removes one branch.
If subsequent nodes are a combination of Multi and other nodes. We add a single Lin layer that acts as bypass for the non-Multi nodes using the fact that multiplicatin by 1 is identity. This is followed by a single Multi layer. We then add Slice operators between the new Lin layer and the non-Multt ti nodes. This keeps the number of branches unchanged but removes the Multi. ti node and the new branch can be handled by the other rules.
Case 1B:
Case 2: All subsequent nodes are LinState. LinState nodes can be fused into a single L LinState node by combining their states and update matrices. As each LinState may have different subsequent nodes itself, we add Slice nodes to extract the respective subspaces of the state. This keeps the number of branches unchanged but puts the graph into Case 5A.
Case 3: All subsequent nodes are ReLU. We can replace them by a single ReLU node. This removes one branch.
Case 4: All subsequent nodes are Concat. One complication is that Concat nodes can depend on other Concat nodes. So, will restrict ourselves at this step by only treating the Concat nodes that depend only on the branch node directly by replacing them with a single Lin node. The rest will be handled by the Lin and Concat case (Case 10) or the only Lin case (Cases 5A and 5B). See the following example:
BranchNode
BranchNode
Concat
Lin
Concat
Concat
Hence, this operation either reduces the number of branches by one or will be followed by a case that reduces the number of branches.
15

---

Case 5A: Only Lin nodes and they are all Slices. This is one of the more challenging cases. While the Slice nodes are simply Lin nodes with special structure, we cannot treat them like standard Lin nodes (see Case 5B). While we can merge them into single Lin node, we will then need further Slices to extract the relevant subspaces for the subsequent nodes. Therefore, we would be simply replacing Slice nodes with Slice nodes. Instead, we use the observation that Slice nodes can be fused with subsequent Lin and LinState nodes and can be pushed after ReLU and Concat nodes. Therefore treat each subsequent node differently, depending its type:
If there are Multi nodes after any of the Slice nodes, they can all be fused into a single Lin node followed by single Multi ti node.
If there Lin LinState nodes after any of the Slice nodes, the Slices be fused with the A matrix of the Lin nodes and the B matrix of the LinState nodes. This uses the fact that composing linear functions results
linear function.
If there a ReLU after Slice node, their position can be switched without changing the nodes. That is because ReLU commutes with linear operations with b = 0 and A with non-negative eigenvalues is the case for Slice nodes.
If there is Concat node after Slice node, we similarly push the Slice a new Lin node after the Concat.
This step does not reduce the number branching nodes but prepares the graph for a removal, with the specific case depending on the remaining nodes.
Case 5B: Only Lin nodes and they are not all Slices. node and then add Slices to extract the relevant subspaces for the subsequent nodes. These Slices can then be pushed into the next operations using Case 5A.
We can combine them into single Lin
Case Both LinState nodes and other nodes. If both LinState nodes and other nodes are present, we can pass through the other variables with dummy LinState variables using zero matrices for A and identities for B. Then, Case 2 can be used to fuse all the LinState variables together.
Case 7A: Only Lin and ReLU nodes where all Lin nodes are followed by only one node which is ReLU. we add Lin bypasses to the ReL Us we will have only nodes left. Each one of them would be followed by a ReLU. Hence, Case 5B can be first applied, followed by Case 3.
Case 7B: Only Lin and ReLU nodes where some Lin nodes are not followed by only one node which is ReLU. In this case cannot apply the above strategy. Instead, fuse the ReLUs by placing ReLU-based bypasses before the Lin nodes. We do this in a similar spirit to Eq. (11), by splitting the positive and negative components and treating them separately. See App. F.6 for the LSRL implementation. Our DAG will then be in Case 7A first, then Case 5B, and, finally, in Case 3.
Case 8: Only Lin and Concat nodes. We add Lin bypasses for the Concat nodes which can then be merged using Case 5B and then Case 5A.
Case 9: Only ReLU and Concat nodes.
Same strategy as for Case but with ReLU bypasses.
Case 10: Only Lin, ReLU or Concat nodes. We introduce ReLU bypasses to all Concat nodes and to the Lin branches which are not immediately followed by a ReLU. This will be followed by applying Case 5B and then Case 3.
The above 13 cases cover all possible branching configurations. After repeated application, they reduce any DAG corresponding to an LSRL program to path graph that can be compiled to one of the recurrent models in Sec.
B Error Bound on the Approximation Scheme for Continuous Functions
In Sec. we outlined strategy to perform universal in-context approximation for continuous functions with Linear RNNs. The full program is in Lst. 1 and illustration of the scheme is
16

---

presented in Fig. 2. In Sec. 4.1 we claimed that the prompt length required to approximate an L-Lipschitz function f (w.r.t the & norm) to precision e N (2t/L/d.)) (m)) The present appendix offers the formal proof of this claim.
The program in Lst. approximates the value of function y = f(q) with the value y the centre of the cell that contains Therefore, the error of our approximation is the maximum difference between f()) and f(e): (g) fc)]l2. First, the length of each side of the cell is 6, that means that || cl x I 8/2. Thus, q c||2 S Vd.6/2. Therefore, thanks to f being L-Lipschitz we get:
(1 - S()]l2 - 8LvdM
If we want to upper bound this approximation error by e, we need to have small enough:
2r
8 Σ dv
Finally, the number of cells we need cover the whole domain N (1/5)d,, this corresponds to us needing sufficiently long prompt:
du
T. [Lddm
da
N
2e
Therefore, if we want our approximation to have error at most anywhere in the domain, we need prompt of length at least (Lv/d/22)))
C Gated RNNs are GRU models
GRU layer (Cho et al., 2014) with input a; Rdd and hidden state h4-1 1 E Rediaasss and output
he E Reviaer can be described as follows:
2t Sigmoid( (Waa, + U,h1-1 bb),
(update gate vector)
(12)
Tt Sigmoid( W.a: + U, + b,),
(reset gate vector)
(13)
h tanh(Wna, (n(r, h--1) bn),
(candidate activation vector)
(14)
h = (1 21) ① h,-1 + 2₁ C h
(output vector)
(15)
In this section, we show a conversion of single Gated RNN layer (Eq. (5)) to k + GRU layers. Here, A is the number of layers the y and h MLPs in Eq. (5). We first show that a single GRU layer can be used to compute the updated state 8, and the output of the first layer of when applied to
2t. Then, every pair single layers of (xx)) and d(s;) can be represented as an individual GRU layer. Finally, single layer can be used to compute the element- -iise multiplication ((x)) @ ((s)).
For simplicity, we assume the Sigmoid and tanh nonlinearities are replaced by ReLUs. If not, they can each be approximated with MIPs and hence also with additional GRU layers. Additionally, for convenience we will assume d dhidden-
C.1 Representing the state update GRU layer
For this layer we set by = 1, W; = 0, U. = 0 giving 2, = Similarly, we set b, 1, W, = 0 U, 0 giving Ti 1. Thus, Eq. (14) reduces to:
h, a WWaa + Un h ba.:
(16)
[0]
8(-11
B
Setting a;
where xe E Iddo/2 h
where $ E R Anissom /2
Wn =
0
I
A by
U, =
where kig is a vector where every element in k is lower bound on
0
kib
Ce,. results in Eq. (15) becoming:
h = a (D
[0 B
0
A
0
8 0
b kuo
Jo [As) + BB, b)] O(x k1b)
a(si)
+
+
I
z
0
Dt
kib
(17)
Note: if we do not want to assume a compact domain for æ., it would be possible to use the same trick as in Equation (11) rather than subtracting k in this layer and adding the next. However, we omit this approach for clarity of presentation.
17

---

C.2 Representing each MLP layer GRU layer
In these layers, similarly to the recurrent layer, we set by = 1, W; = 0, U, 0 giving 2 = 1. In the same way, we set b,
W U,
giving rr 1. Here, however, we set
Whe h:
0
Dkc b
Wh.
Un 0 and by W 72
except for the first of such layer where bn.
bhs by. + Wy kib
a1,t
Thus, for an input az
the layer output (Eq. (15)) for layer i is:
a2,t
he (WVo
a1 02,
bo: byy
)
(18)
W
Here, de and b are the i-th layers (including the ReLU) of respectively & and y in Eq. (5).
C.3 Representing the multiplicative gating with single GRU layer
The only thing left is to model the element-wise multiplication of the outputs of & and y in Eq. (5). We
0 0
do this using a GRU layer with b. = 0, W = 0, U
We set b, = 0, W, = 0, U,
0
0 0
[ai
giving тı = 0. We also set b, = 0, W =
Un = 0. Thus, for an input a =
the
I
0
02.
output h, (Eq. (15)) of this GRU layer becomes:
1
0
0
1
h4
(19)
22,1
I
02,
( о az,1
If a, is the output of GRU layer constructed in Eq. (18) (as is in case), then it must be non-negative. This is due to the ReLU application in Eq. (18). Hence, the application another ReLU a1,t in Eq. (19) be safely removed ReLU is idempotent and Eq. (19) simplifies
(20)
a1
a221
Thus, this construction computes element-wiss multiplication of a1,t and ag,t.
C.4 Composing the operations to model single Gated RNN layer
In order to represent Eq. (5), one GRU layer for the recurrence (as described in App. C.I), followed by GRU layers modelling a pair of the k MLP layers of ø and (App. C.2), completed with single mixing layer (App. C.3). This stack of k + 2 layers models exactly the Gated RNN layer (Eq. (5)):
0
+ +b)
8; a
81-1
+
0
y: (xx)) O ((si)
With this, we have shown that any Gated RNN (Eq. (5)) can be expressed as a GRU-based model. Hence, the two universal approximation programs in Lsts. and 2 can implemented also in GRU-based models. Thus, the GRU architecture can also be universal in-context approximator.
D Gated RNNs are LSTMs
single I.STM layer (Hochreiter and Schmidhuber, 1997; Gers et al., 2000) with input a, Rdi,, hidden state h E geaiaie candidate memory cell é R dadaasn memory cell c gddiatio and layer
18

---

output ha E Rdiat can be expressed as:
ft Sigmoid ja + Uyh1-1 + by),
(forget gate vector)
(21)
& Sigmoid a + U;h + b.).
(input gate vector) (output gate vector) (candidate cell vector)
(22)
0: Sigmoid a. + Ught-1 bo),
(23)
C = tanh(Ween + Uch + bo),
(24)
C₁ f. ct-1 + 4 ② Ct,
(memory cell vector)
(25)
h o, o tanh(c:),
(output vector)
(26)
where hg 0 and co 0.
In a way analogous to App. C, we show that a single layer of a gated RNN (Eg. (5)) can be expressed using k + 2 LSTM layers, where k is the maximum depth of either of the MLP networks y. We again follow the setup of replacing all Sigmoid and tanh activation functions with ReLU activations which we denote and we again assume that dın dnidden. The set up follows the same structure as in App. C. First, we show that the non-linear state update computing 8 can be expressed as a single LSTM layer. We then show that single LSTM layers. Finally, a single layer can compute the Hadamard product between (xxe) and ((x)). Therefore, any Gated RNN with ReLU activations can be expressed as LSTM with ReLU activations.
represent the layers in MLP networks ((::) and Φ(sx) using
For clarity of the exposition, once again assume that our inputs belong to compact domain X of real vectors. This implics that the set is bounded and, in particular, that we can find a vector ky such that kis. < (x;); for E [din] for all Φ; X. In other words, we have (æ, kw); 0 for for i E 1, 1,..., din- We will make use of this fact several times when dealing with ReLU activations.
D.1 Representing the state update as an LSTM layer
We first represent the non-linear state update in Eq. (5) using single layer of an LSTM. In particular, we set Wi = 0, U, = 0 and b = so that ft O. We also set W, 0, U, 0, b; and W. 0, U, = 0, b. = 1. This results in é; = 1 and é, 1. We see from this that the LSTM layer
with these weight settings reduces to
h = =o + Uohi-1 + bo).
(27)
We now set a₁ =
0
where D4 E Rdia/2 h--1 =
s/--
where 8 E TR drussmm / W. =
Zt
0
0 B
A
0
b
0 1
U. =
0
b,
=
kitb
so that
h = a
0
B
I
zt
0
+
A
0
0
81-1
0
+
kib
b 1)
=
o [o(As + Ba, + b)
[ kbb)
= xt kbb
(28)
D.2 Representing each MLP layer as an LSTM layer
Now we want to use an LSTM layers to model the MIP layers of both y and / simultaneously. We
set W, 0, U, 0, b and W, = 0, U: 0, b and W. 0, U.
= 0, b.
as
W&
before. We make a change for these LSTM layers by setting W.
de
W U. and
0
by. b
b
a
b.
except for the first layer where by- by,
Thus, for an input a;
+ W
a2.1
the layer output is:
W (WW
0
a1
b
(a((1.))) (
h4
+
(29)
222
by.
Here, de and 7 again refer to the i-th layers (including the ReLU) of respectively and y in Eq. (5). Note that, without loss of generality, if we have that & has m layers whereas 7 has & with m < k. then can also model this by simply adding additional layers to model additional layers for y whilst
19

---

simply passing on & unchanged. Specifically, we set set the weights to ensure that f: = 0 and that žt and é, are 1 so that h; = or. The input to this layer for >> is then given as d,
()()) The we
0222
set the weights to compute o
o = σ (66
[01 boo))
)
σ(5)) (72(1)))
(30)
W 7
02
D.3 Representing the multiplicative gating with an LSTM layer
Finally, we model the element-wise multiplication of the outputs & and y Eq. (5). To do this we set the weights of the input gate and candidate cell vectors for the final layers of of o and ø to be as follows:
0
0
1,
0
0
I
0
22.2
+
0
a1..
(31)
and
0
0
a
0
0
è
(32)
0
I
a2.
a22r
Then by setting w 0, y, 0, b, and W. 0, U. 0, b, 1 to force fi and
Oj 1, we get
0 y1 (c:) a a o 02)))
o(ai, o a2.)
(33)
D.4 Composing the operations model single Gated RNN layer
To model the gated RNN described in Eq. (5), again follow the same lines described in App. C. In particular, we use one LSTM layer for the recurrent state updated as described in App. D.1. We then stack LSTM layers as described in App. D.2 to model the A MLP layers of and y in parallel. We then use one final layer to both give the final MLP layer of & and and to compute their Hadamard product as set out in App. D.3 in order to match the output of the gated RNN in Eq. (5). Now, since we are working with = ReLU, both ((x,) and ((s;) are positive and therefore so is their product. Hence, applying a to the product components Eq. (33) leaves the the components invariant. Therefore, we output is
y: ((x))) O )(s)))
(34)
as required.
Hence, we have shown that a single layer of a gated RNN as described by Eq. (5) can be represented using + 2 LSTM layers where k is the maximum depth of d and 7. Therefore, once again, the two universal approximation programs in Lsts. and can also be implemented for LSTMs. Hence, LSTM models are also universal approximators in the sense described in Sec. 4.
E
Gated Linear RNNs are Hawk/Griffin Models
single residual block of a Hawk/Griffin model (De et al., 2024) consists of two components, a recurrent block for temporal mixing which makes of one-dimensional temporal convolution, as well as real-gated linear recurrent unit (RG-LRU) and a gated MLP block. Specifically, we consider an input a, di inputs to the blocks of dimensions d,, and outputs from each block of dimensions din- Within blocks, all vectors have dimensionality dnidden = Edin, where is denotes an expansion factor. Below, formally describe the form of the recurrent and gated MLP blocks which the two main components making up the residual blocks used for Hawk and Griffin.
Recurrent block. The recurrent block consists of two branches. The first applies one-dimensional temporal convolution followed by RG-LRU. The second branch simply performs linear transfor- mation followed by non-linearity, i.c. applies single layer of MLP.
20

---

m, W Wme + bm- Consider the first branch of the recurrent block with an input a, The one-dimension temporal convolution can be written as;
a/ Wa
(35)
g: GeLU(W + bg):
(36)
M = [dk-(demn 1) 2/-2,,0-1:
(37)
dam-
WM i] + bonn
(convolution with window size dean)),
(38)
where Deony is a bias vector and WM = [B, B. ... ..B.. is the convolutional kernel for the one-dimensional temporal convolution.
...
The output of this convolution is then fed into a RG-LRU. We can write this down concretely using as an input from the one one-dimensional convolution and with recurrent state h, E Rwwd
r; Sigmoid di w, b,),
(39)
& Sigmoi d W, + b:),
(40)
a Sigmoid A);
a learnable parameter)
(41)
= am
(c = 8 fixed scalar constant)
(42)
h; = a O h
(fe 2.).
(43)
Now consider the second branch of the recurrent block. followed by non-linear activation:
This performs linear transformation
gt = GeLU ya + bg).
(44)
To get the final output of the recurrent block, we multiply the components of the vectors computed from each branch within the recurrent block and then perform non-linear transformation:
h 9: C he:
(45)
O, Woh' + bo.
(46)
Gated MLP block. After passing through the recurrent block, we pass the output ot into gated MLP block. Again we have two branches, the first where we linearly transform the input to this block
e Wo + bey
(47)
and the second performs single layer MLP transformation
ft = GeLU(W, oo + by).
(48)
These are then combined through a Hadamard product and linear transformation as
e, E; C ft.
(49)
(50)
We then have that the vector m, acts as the output of the residual block given the input a₂.
Distinction between the Griffin and Hawk models. Hawk is the more simple of the two architectures proposed in (De et al., 2024). Here, residual blocks using the recurrent block described above simply stacked on top of each other to form the Hawk architecture. Griffin, on the other hand, mixes recurrent blocks and local attention. In particular, followed by one residual block using local MQA attention (Beltagy et al., 2020; Shazeer, 2019).
residual blocks with recurrent blocks
Simplifying Assumptions. We again follow the setup of replacing all Sigmoid and tanh activation functions with ReLU activations which we denote a. Furthermore, we assume for simplicity that din dnidden by choosing E 1. Moreover, the Hawk and Griffin architecture contains residual connections and normalising layers which we omit.* 4 We again assume compactness the input domain X' and denote vector of finite values kn. such that ka,i < (æx); for i E [din] and all 2t E X, just as before. Finally, we assume that d = T where is the maximum sequence length.
*We will force of our recurrent blocks to implement the identity function. So instead this, we could implement the function the recurrent block and use a residual connection between the residual block input and the output of the recurrent block to achieve the same identity function. However, for clarity we ignore residual connections in our derivations.
21

---

E.1 Representing the state update using a recurrent block
Starting with the input to the Hawk model, which we denote at, we define this to be function of the imput to the Gated RNN Dt as a₂
0
First, we set Wa so that aj dz. Next choose
zt
0
A 0
0
B
matrices A
and B
which then use, with convolutional window size
0
of deonv 7 form the convolutional kernel WMr B, AB, A²B, A.B, ...]. Setting the
convolutional bias as bconv
gives
2t
Σ
WM -i beonv
(51)
i=0
0 1
Ba,
ABa + + Ba +
(52)
8t
(53)
Now, we pass 2₁ through the RG-LRU. We set A = 0 so that a; = 0. We also define W; = 0 and b, so that 1. This gives h 2t, so that we pass the output of the one-dimensional convolution through he RG-LRU.
Next, let's focus the second branch. Making use of the lower bound ky the domain X, we set
W, = I and by = kio
so that
1
1
(1)
1
9: I +
(54)
kib
oxx
k(b)
Et
kib
where we used that (x, kib); 2 0 for every i. Combining the branches gives
h!
1
8t
9: kib
(55)
k1b
1
2t
0
We finally get the output of the recurrent block by defining W. = I and bo
so that
kib
81
O,
(56)
Z
E.2 Representing the identity function using a recurrent block
We now show that we can pass an input unchanged through a recurrent block. Assume that the input to the recurrent block a,
a1. a2.6
with W.
that af (a,- Then we define matrices À and
B- = which we then use to form the convolutional kernel WA = B AB, A"B. Á' A'B, Finally, setting the convolutional bias b 0 results z, a, From here, we can again set
A = 0, W. = O and b, = 1 so that hss Looking at the second branch and setting W, = 0 and b, 1 that h ht- Finally, we can simply output the input the recurrent block by setting
W, I and b, = 0 so that
h, which means that = a₂.
E.3 Representing each MLP layer as gated MLP block
We can represent the MLP layers of the networks ((st) and ((x)) as described in Eq. (4) using Gated MLP blocks. We again denote the i-th layer of Φ and y as o; and Assume that the input to the
gated MLP block is a. = [a..1
Then, on the first purely linear branch, let us define We = I and
a2,1
22

---

Wo,
0
b. so that e;
On the second non non-lineat branch, we can define W
and
0
W
bo by
b
This results in
W
a1.1 0221
ba. by
Φz ( m((a2,))
f: -
(57)
W
Due to our setting of e,, we get e/ ft. Further, defining Wm I and bm makes the output of the MLP block be
m:
(58)
Emulating the layers only one the two nefworks. Suppose without loss of generality (WLOG) that & has m layers and has n layers where m < Suppose also that our input to the MLP block
(xx))) 02
a, =
Again, on the first purely linear branch, us define We and b = 1 so that
0
e, 1. Now we modify the weights on the second non non-linear branch by defining W =
W
0
and by
This gives us
b
(66 0
(ax)) a2
0 by:
[o((b())))
d <b(x))) 7 (a(.))))
f-=o
W
+
=
(a(,,))
(59)
where we have used that since d(w;) is ReLU network whose final activation is a ReLU, we have that vb(x,) ()(x))). Hence, if our networks have different depths and have fully emulated one of the networks, we can continue to emulate the remaining layers of the other network while keeping the fully emulated network fixed and unchanged.
E.4 Representing the identify function using gated MLP block
In this section we show that we can represent an identity function using a gated MLP block. This can be simply done setting W, b, 1, W. = b. 0, W.. I and bm 0. This then
[a to the gated MLP block, the output of the gated MLP block is
gives that for an input a;
a2,t
m, = az. Thus, we pass the input through the gated MLP unchanged.
E.5
Representing multiplicative gating with a gated MLP block
The final thing we need to do is to compute element-wise product of two vectors in order to match the output Eq. (4). other words, to match the (xx.) @ ((s)) operation.
[I..]
Again, assume that the input to the gated MLP block is a,
Working with the first linear
02
0 0
branch, define W-
and b
so that
I 0
0
0
01,t
0
e,
I
0
02,1
a1.
(60)
Next, we define W, = I and b
0 so that
[a(a..))
f.
(61)
(aa).))
Setting W.. and bm = 0 gives the output of the gated MLP as
m. =
ar, ① (
(62)
23

---

E.6 Composing the operations to model single gated linear-RNN layer
Now that we have all the individual layers, we can combine them so that we can use Hawk model to emulate single Gated RNN layer.
First we start by taking the input of the form a₂ 0 [t
We use residual block that consists of
recurrent block computing the state update as descried in App. E.1 and then a gated MLP block that computes the identity function as demonstrated in App. E.4. This gives an output from this first
recurrent block as O,
8t
Next, we emulate the MLP layers of the networks & and y in parallel. Suppose WLOG that e and y have and MLP layers respectively, where S We stack m residual blocks using recurrent blocks that implement the identity function as described
App. E.2 followed by MLP blocks that apply the MLP layers of @ and y as described in App. E.3. Stacking such residual blocks results in the output m, = ((xx))
(8t) where we can fully emulate the shallower network ((x,).
Now, for the remaining k m layers for the network (a,,), stack residual blocks with recurrent blocks implementing the identity function as described in App. E.2 and MLP blocks that leave of ((x,) unchanged whilst applying the additional layers needed emulate y(s)) as described the end of App. E.3, After stacking k m additional residual layers in this fashion, the output of the final
(st)
residual block will now be m. = (0(z))
which fully reconstructs the MLP networks y and ø.
Finally, we utilise residual block with recurrent block that implements the identity function as described in App. F.2 followed by gated MLP block that applies multiplicative gating as described
in App. E.5. This then gives output of this final residual block m, (x)) O (
Since )(x;) is MLP network with the final activation function being a ReLU activation, we have that ar ())))) = db (.))), giving the required final output from the stacked block of residual blocks as
0
m (st)) O o(xz)
(63)
Hence, we have shown that single layer of a gated RNN as described by Eq. (5) can be represented using i- + 2 Hawk residual blocks where ls is the maximum depth of & and y. Once again, the two universal approximation programs in Lsts. 1 and 2 can also be applied to Hawk models they can represent Gated Linear RNNs. Therefore, Hawk models are also universal approximators in the sense described in Sec. 4.
Gated Linear-RNNs are Griffin models too. The above argument extends to the Griffin architecture which uses stacks of two residual blocks with recurrent blocks followed by residual block with attention. The only thing that changes is that for every third residual block, which in our argument will be used to compute the MLP layers of o and y in parallel, the recurrent block is now replaced with local MQA block.
We can set the key query and values matrices to implement the identity function which is to act input to the block. Hence, as a corollary of the above argument, we can also show that the universal approximation programs in Lsts. and 2 can also be implemented as Griffin models. Therefore, Griffin models can also be universal approximators in the sense described in Sec.
F
Definitions for some helper functions in LSRL
F1 f_not
This is convenience function that creates NOT function block. assumes that x is 1. Works with scalar and vector-valued inputs. With vector-valued inputs, it acts clement-wise.
not_x
24

---

F.2 f_and
This is convenience function that creates an AND function block. It assumees that and are 1. Works with scalar and vector-valued inputs. With vector-valued inputs, it acts element-wise. mu is the approximation parameter μ for f f_step as described in Sec. 3.
and_x_y
CeLucf_step(s, mu) f_atep(y, mu)
F.3 f_or
This is convenience function that creates OR function block. It assumes that and y are 1. Works with scalar and vector-valued inputs. With vector-valued inputs, acts element-wise. mu is the approximation parameter μ for f f_step as described in Sec. 3.
or_ryy (x y, mu=nu)
F,4 f_smaller
This is convenience function that less than comparison block. Works with scalar and vector-valued inputs. With
inputs, acts element-wiss wise. mu the approximation parameter for f_step described in Sec. 3.
snaller x.y
fystep numnu)
f_larger
This is convenience function that more than comparison block. Works with scalar and vector- valued inputs. With vector-valued inputs, it acts element-wise. mu approximation parameter ju for f f_step described in Sec. 3.
larger.x_y f_step(x y, mu*nu)
F.6 f_relu_identity
Identity operation using ReLUs. This is useful for debranching when some of the branches have Re but the other don't. We can add this as bypass for the ones that do not and can then merge the ReL Us together (see App. A for details).
positive_part ReLU(x)
negative_part ReLU(
Linear(
inputex, A=-1 di).)
b=zeros(x. din. 1).
both Concat([pasttive.cccc
negative_part.)
relu idu_identtty Linear(
11
input=hoth
12
A=hstack(eye(c din), eye(x.dim)),
b-zeros(x.dit 1),
E.7 f_modulo_counter
Computes the mod divisor where is counter starting from The idea is that rotate a unit vector so that makes a full revolution every divisor rotations. dummy_input can be any variable, we use it only to construct constant.
divisor R [[cos(angle), sangle),)),, (sin(angle), costangle))]
angle
unit_vector IE13, [A]] we first rotate, then cutout If vant the first output be 8 need to have the init_ one stop before that init_state () unit_vector
25

---

1/divisor
eycler LinState(
input-Gurny_Anputt
A=R,
B-zeros(2, ounry_inuut.din))
:nst_scate-intt_statt
12 )
13 t не nоs need to extract the position of the cycler
14 extractor_natrit
vstack(>[(R"1 unit_vecter)..t for range(divisor)))
Linear( 16 input=cycier.
17 i=extractor_nattre natrixx
18
brreros(civvier,
)
28 dot product with the row extractor-natris corresponding to the current position the eycler is
21 6 the dot product with the second highest is cas (angle)
22 thas,
hot encoding the current position the cycler threshold 1-cos(angle/2)
(indicatrn,,,, cos(angle integer 25 vod_value Linear(
layer
26 one_hot,
A=LLs for range(divisor)]).
b-zeros(1, 1)
29
26

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and the introduction clearly state all the contributions of the paper and clearly differentiate the theoretical results which hold in general and the empirical phenomena that we observe, which may not generalize to all settings.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
ii Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper discusses the limitations of the present work. There is dedi- cated Limitations section in Sec. 7 that addresses the fact that we only provide constructive existence results but not necessary and sufficient conditions for universal in-context approxi- mation to arise. We also highlight that our results might not hold models with structural constraints on their parameters. Moreover, we have dedicated section (Sec. 5) which addresses some the limitations constructing universal in-context approximators with fully recurrent architectures in practice. This section proposes solutions and demonstrates that they result more numerically stable models which are more likely to occur in practice. Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
27

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
i... Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
The paper has two main theoretical results: the constructions of universal in- context approximators for continuous and for discrete functions. Both results presented as LSRL programs which compile to the architectures considered in this work. Furthermore, these programs have been implemented in Python, their correctness has been tested and they are available in the supplementary materials.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
iv. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: There are two experimental aspects to this work. First, there is the implemen- tation of LSRL and the universal approximation programs in Lsts. and 2. The most critical aspect of implementing LSRL is the debranching algorithm which is described in detail in App. A. Additionally, the two programs are described in full in their correspond- ing listings. We also provide Python implementation for the LSRL compiler and the two programs.
Second, there is the study of how affected by parameter noise are the different implementa- tions of the conditional assignment operator f_ifelse which was presented in Sec. 6. The details of this experiment are described in Fig. 4 and we also provide the code with which did the experiment and our plots. Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
28

---

reproducible, they While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
V. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We providing code that includes an implementation of LSRL, the two universal in-context approximation programs in Lsts. and 2 and everything needed to reproduce the experiments in this work.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co Beaiii/gguiee///oeeeeeeeee
ssionPolicyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips. cc/pcbclplccccccccc quublice
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
vi. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [NA]
Justification: The experiments in our work are based on constructed models rather than trained models. Therefore, considerations such as dataset, optimizers and hyperparameters do not apply.
Guidelines:
The answer NA means that the paper does not include experiments.
29

---

The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
vii. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: For this work, uncertainty quantification could only make sense in the context of Fig. 4. However, the behaviour we observe, especially for the continuous case, is bimodal. As bimodal distributions cannot be properly captured with error bars we decided against using them. Furthermore, we are studying whether a phenomenon occurs, rather than quantifying Therefore, we decided instead use strip plot instead explicitly shows all our results unabridged, clearly indicates the bimodal nature of the results, and distinctly showeases the noise robustness trends of the different approaches we consider. Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
viii. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: Our experiments were ran on a single machine and using only CPU com-
pute. Therefore, the compute required is negligible for the contemporary machine learning standards.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
ix. Code Of Ethics
30

---

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips httpe/rrrp...pppoooo//ccccsstola
Answer: [Yes]
Justification: This is theoretical work with no human participants, datasets, potential societal impact or harmful consequences. Therefore, the present work has no moral or ethical relevance implications.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there special consid- eration due to laws or regulations in their jurisdiction).
x. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: As mentioned above, this theoretical work which establishes theoretical properties of mathematical objects that are already used in practice. However, we do discuss the implications of our findings, namely that if models are universal in-context approximators, then it might be difficult to ensure that they are aligned and cannot be misused. Nevertheless, we only show that this property already present in existing models, and hence our work does not introduce new attack misuse vectors. On the contrary, we hope that us highlighting this issues will help the community to develop safer and more secure generative AI systems. Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
xi. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification: We release no data or models.
Guidelines:
31

---

The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
xii. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: The paper does not use existing assets beyond common Python libraries. Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paper thcode. com/ datasatss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
xiii. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: The only new asset arising from this work the LSRL code base which we have ensured to be well-documented, accompanied by unit and integration tests, and with illustrative Jupyter notebooks.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
xiv.
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
32

---

Justification:
There were no human participants involved in any part of this work. Guidelines:
The answer NA means that the paper does not involve crowdsoureing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
xv. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This is a purely theoretical work and as such no IRB approval or equivalent was necessary or appropriate.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
33


Reviews:
[{"review_id": "dFMvjLp8QO", "reviewer": "Reviewer_uwk6", "summary": "The paper explores the potential of various recurrent neural network architectures (RNNs, LSTMs, GRUs, Linear RNNs, and gated architectures) to act as universal in-context approximators, crucial for zero-shot and in-context learning without fine-tuning. It introduces the LSRL programming language to facilitate this, highlighting the role of multiplicative gating in enhancing model stability for practical applications.", "questions": "See weakness.", "limitations": "N/A", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The concept of universal approximation in the context of in-context learning is both intriguing and innovative. This research presents a unique angle that adds valuable insights to the field.", "weaknesses": "- The paper is very dense and hard to read and follow. The construction is rather tricky. Do you have any high-level idea in the proof?\n\n- I don't understand the practical value of proving that linear RNNs are in-context universal approximators. Linear RNNs are known to be inferior at in-context learning and retrieval (see references [1], [2], [3]). How does demonstrating that linear RNNs are in-context universal approximators mitigate this issue in practice? Isn't it somewhat contradictory to claim that RNNs are in-context universal approximators but cannot perform in-context learning well?\n\n- Many theoretical papers claim that non-linear RNNs are universal approximators or Turing-complete, often under impractical assumptions like infinite precision and exponential hidden state sizes. What about the assumptions in this paper? Do you think all assumptions are practical?\n\n- The definition of gated (linear) RNNs in this paper is very unusual. Typically, gating refers to the recurrence itself being gated by forget gates or data-dependent decays, rather than gating the output. \n\n\n[1] In-Context Language Learning: Architectures and Algorithms https://arxiv.org/abs/2401.12973\n\n[2] Simple linear attention language models balance the recall-throughput tradeoff https://arxiv.org/abs/2402.18668 \n\n[3] RNNs are not Transformers (Yet): The Key Bottleneck on In-context Retrieval https://arxiv.org/abs/2402.18510", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for finding our work to be _intriguing and innovative_ and for _adding valuable insights to the field_. We would like to address their concerns.\n\n> The paper is very dense and hard to read and follow. The construction is rather tricky. Do you have any high-level idea in the proof?\n\nThe high-level ideas behind the two LSRL programs in Listing 1 and Listing 2 are rather succinct, although their implementation might indeed be a bit tricky. \n\nFor the _continous case_ (Listing 1, Section 4.1), the high-level idea is to discretize the domain into cells and to approximate the target function $f$ with a piece-wise constant function $g$ that is constant in each cell with a value equal to $f$ evaluated at the centre point of the cell. Then, evaluating this approximation at an input $x$ in an iterative fashion requires one to go through the cells one by one, find the cell containing $x$ (Lines 16-19) and return the value of $f$ at the centre of that cell (Lines 21-23). This is a variant of the classic universal approximation results that used similar discretization constructions.\n\nFor the _discrete case_ (Listing 2, Section 4.2), the high-level idea is to recognize that any discrete function can be represented as a map or a key-value dictionary. This dictionary can be provided in-context, iterating keys and values. Then, given a query, we can process the key-value pairs one by one until we find the key that matches the query (Lines 19-25). Then, we know that the following value is the one we should output, so we copy it into a state variable and output it repeatedly (Lines 26-30).\n\nWe understand that this is the trickiest part of the paper and will add the above explanations as a way to improve the presentation for the camera-ready version of the paper.\n\n\n> I don't understand the practical value of proving that linear RNNs are in-context universal approximators. Linear RNNs are known to be inferior at in-context learning and retrieval (see references [1], [2], [3]). How does demonstrating that linear RNNs are in-context universal approximators mitigate this issue in practice? Isn't it somewhat contradictory to claim that RNNs are in-context universal approximators but cannot perform in-context learning well?\n\nThe reviewer is raising some very interesting questions. \n\nIt is true that, especially in deep learning, there is often a tension between theory and practice.\nThe fact that we provide a proof by construction that something is possible does not mean that these are solutions that models will learn via gradient descent on real world data.\nHowever, our universal in-context approximation results might imply that perhaps the low empirical performance is an issue with how we train, prompt, or evaluate these models, rather than some fundamental limitations of the in-context learning abilities of Linear RNNs.\nThe very fact that models with good in-context learning abilities exist for these architectures (as shown in our results) might indicate that we might be doing something suboptimal with the training and inference of these models.\n\nFurthermore, there is no free lunch: for any given architecture, we can design problems that it won't perform well on. \nIn our view, the question should not be to figure out if Linear RNNs or Transformers are the \"better\" architecture.\nThey offer different trade-offs and are useful in different settings.\n\nWe would like to highlight that we have discussed the fundamental properties of such theoretical results in our Limitations section.\n\n> Many theoretical papers claim that non-linear RNNs are universal approximators or Turing-complete, often under impractical assumptions like infinite precision and exponential hidden state sizes. What about the assumptions in this paper? Do you think all assumptions are practical?\n\nWe make no assumptions about infinite precision and exponential hidden state sizes.\nAs we anyways use discretization, infinite precision is neither necessary nor helpful for our constructions. \nThe hidden state sizes of our constructions are also independent of the target precision $\\epsilon$ (in the continuous setting) or the key-value dictionary size (in the discrete setting).\nThe error bound on line 228 does assume Lipschitzness, but that is a standard condition.\nTherefore, we work with very few assumptions on the architecture and, hence, our assumptions are practical.\n\n> The definition of gated (linear) RNNs in this paper is very unusual. Typically, gating refers to the recurrence itself being gated by forget gates or data-dependent decays, rather than gating the output.\n\nThis paper focuses on Gated Linear RNNs and Gated RNNs. By definition, the gating of Gated Linear RNNs cannot be on the recurrence itself as the model will not be linear anymore. Similarly, the gating of Gated RNNs cannot be on the recurrence because the model would not be an RNN anymore (the standard RNN is a linear state update followed by an element-wise non-linearity). Perhaps the reviewer has LSTMs and GRUs in mind where the gating is indeed on the recurrence. However, as we show in Appendices B, C and D, our models subsume LSTMs, GRUs and the Hawk/Griffin architecture and hence they apply to these architectures as well.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "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": "XQstD4orAR", "reviewer": "Reviewer_Aa6c", "summary": "The paper shows that sequence-to-sequence networks like RNNs, LSTMs, and Mamba are universal in-context approximators, which was defined as the ability to compute any function with an appropriate prompt. The authors propose a programming language called LSRL which can express any language expressible by a linear RNN (and vice-versa). The authors discuss numerical stability issues in their proposed language and discuss possible solutions. Overall, the paper streamlines a theoretical framework to understand expressivity of recurrent models and will prove to be an interesting contribution to the wider community.", "questions": "Please check my questions above.", "limitations": "The authors discuss the limitations of their work in section 7, and clearly indicate the future directions that the community can pursue starting from their work.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The strength of the paper lies in its simplistic exposition of its motivation, proposed framework, and the discussions of various frailties and solutions. Sequence-to-sequence models have proven to be competitive to transformers, but the gaps between the two architectures are still under exploration. Expressivity for transformers has been studied through RASP (Weiss et al.'21), and this paper shows a similar programming language to understand linear RNNs. Through LSRL, the authors argue about the kind of languages linear RNNs can express and the impact of gating in RNNs. Overall, the paper takes an important step towards understanding sequence-to-sequence models.", "weaknesses": "As such, I don't see clear weakness with the work. However, I would like the authors to be more clearer on their contributions. \n How is the theoretical framework different from the expressivity studies on feedforward networks and RNNs [1, 2]? That is, as far as I understand, the constructions in section 4.1 and 4.2 still construct a RNN that can do a dictionary lookup for each possible query, very similar to ones constructed in previous works. If I understand correctly, the strength of the paper is more on creating a systematic programming language to compile languages into RNNs. The universal expressivity happens to be a by-product of this framework and is very similar to previous constructions. \n\nReferences:\n\n1: Andrew R Barron. 1993. Universal approximation bounds for superpositions of a sigmoidal function. IEEE Transactions on Information Theory, 39(3):930–945.\n\n2: Recurrent Neural Networks Are Universal Approximators. AM Schäfer, HG Zimmermann. In Artificial Neural Networks–ICANN 2006: 16th International Conference, Athens, Greece, September 10-14, 2006. Proceedings, Part I 16, pages 632–640. Springer.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for recognizing our work's contributions towards understanding sequence-to-sequence models. \n\n> As such, I don't see clear weakness with the work. However, I would like the authors to be more clearer on their contributions. How is the theoretical framework different from the expressivity studies on feedforward networks and RNNs [1, 2]? That is, as far as I understand, the constructions in section 4.1 and 4.2 still construct a RNN that can do a dictionary lookup for each possible query, very similar to ones constructed in previous works. If I understand correctly, the strength of the paper is more on creating a systematic programming language to compile languages into RNNs. The universal expressivity happens to be a by-product of this framework and is very similar to previous constructions.\n\nThere are two key differences between the classical results the reviewer has shared and this work. \nFirst, removing the non-linearity from the state update (going from an RNN (Eq.2) to a Linear RNN (Eq. 3)) is non-trivial and complicates the analysis significantly. \nIn fact, it was only last year that it was shown that Linear RNNs can be universal approximators (see Wang and Xue, 2023). \nSecond, and that's the key novelty of our work, universal _in-context_ approximation is a very different beast than the classical universal approximation. \nClassically, one can choose the model parameters conditional on the target function. \nWith universal _in-context_ approximation, however, the model parameters are _fixed_ and independent of the target function. \nIn other words, a single model needs to approximate all target functions.\nThat makes our proofs very different and, at least in our view, the results much more interesting as they extend to how we often currently interact with large pretrained models: via prompting, rather than via training.\n\nShida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. NeurIPS 2023", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "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": "uxOKxkci4O", "reviewer": "Reviewer_3Cyh", "summary": "This paper designs a new programming language LSRL that compiles to a fully recurrent structure and shows that multiple recurrent architectures including RNNs, LSTMs, and SSMs can serve as universal in-context approximators. They also show that multiplicative gating allows more numerically stable construction.", "questions": "1. If one considers a hybrid architecture that consists of both RNN and Transformer Layer, can one use a combination of LSRL and RaSP language to compile such a model?\n2. In the conclusion, it is mentioned that the compiled transition matrix is often diagonal. Are those matrices exactly diagonal here?", "limitations": "Limitation are adequately addressed.", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. This paper reveals interesting properties of multiple recurrent structures and shows they can perform universal in-context approximation.\n2. The discovery of the numerical instabilities when gating is missing is also interesting and provides intuitions on the role of gating.\n3. The programming language created by this paper may be of independent interest.", "weaknesses": "1. In the introduction, the paper mentions that RNNs can be prompted to act as any token-to-token function over a finite token sequence. However, in the actual construction, the prompt specifying the function must be put after the query. This claim, while necessary for RNNs, is not very standard. Clarity will be improved if this caveat is mentioned earlier in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's positive opinion of LSRL and our investigation into the numerical instabilities without gating.\n\n> In the introduction, the paper mentions that RNNs can be prompted to act as any token-to-token function over a finite token sequence. However, in the actual construction, the prompt specifying the function must be put after the query. This claim, while necessary for RNNs, is not very standard. Clarity will be improved if this caveat is mentioned earlier in the paper.\n\nWe have mentioned at the beginning of Section 2 (Line 74/footnote 1) why it is necessary to put the query before the prompt.\nHowever, we agree with the reviewer that mentioning this in the Introduction could further improve the readability of the paper.\n\n> If one considers a hybrid architecture that consists of both RNN and Transformer Layer, can one use a combination of LSRL and RASP language to compile such a model?\n\nThis is indeed an interesting proposition. We do not see any immediate reasons why that would not be possible. As long as one is willing to specify different layers in different languages, the layers could be compiled with the respective compilers and then fused together. Some care would have to be exercised to align how variables are represented in order to ensure interoperability but that should not be a problem.\n\n> In the conclusion, it is mentioned that the compiled transition matrix is often diagonal. Are those matrices exactly diagonal here?\n\nThe transition matrix depends on the program that is to be implemented. For the two universal in-context approximation programs we have in the paper the transition matrices are not diagonal but highly sparse with many blocks being diagonal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "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": "4zRwbnsIbN", "reviewer": "Reviewer_qz76", "summary": "In-context learning has emerged as one of the puzzling properties of language models at scale. While a lot of work has been done to understand the mechanics supporting this behavior in attention-based models, much less has been done on recurrent models. Given the renewed interest in these architectures (e.g. Mamba or Griffin), a better understanding of in-context learning in these models is needed.\nThe paper proposes to study these abilities through the notion of in-context approximation, that is whether a model can produce a certain function if it's given the right prompt. The authors show two universality results by introducing a low-level programming language that implements the canonical proof ideas for this kind of result within the neural dynamics.", "questions": "- Can the authors position their work compared to:\n    - [Zucchet et al. 2023](https://arxiv.org/abs/2309.01775): they study the link between gated RNNs (very similar to the one you consider) and attention, therefore connecting to the literature on how attention solves in-context learning. Additionally, they also highlight the importance of gating in practice.\n    - [Orvieto et al. 2023](https://arxiv.org/abs/2307.11888) and the references mentioned there on the universal approximation capabilities of linear RNNs. It would be valuable to compare with those results in more detail, c.f. next question.\n- I am wondering about how fundamentally different is in-context approximation compared to the universal approximation. Is one a subset of the other? How are they related / what are their difference? Given that one of the contributions of the paper is the introduction of this notion, answers to these questions would be particularly valuable.\n- What is the interest of introducing such a low-level programming language? It feels to me that LSLR is a tiny abstraction on top of the neural equations that is not necessarily worth it. Why not directly include the primitives used in e.g. Figure 3 in this language?\n\nSome more open questions:\n- The paper shows that deep linear RNNs can implement a variety of functions in context, somehow supporting the argument that linear recurrence + nonlinear instantaneous processing is quite powerful. Does the framework provided by the authors give hints about which kinds of functions are difficult to learn with such networks? The authors mention something around these lines in lines 165-167, but more detail would be appreciated.\n- Are there some links between the in-context learning approximation notion introduced here and the notion of controllability in control theory?", "limitations": "The authors have correctly addressed the limitation of the applicability of their construction to more practical settings as well as highlighted that better understanding in-context learning is critical in a world where LLMs are available through APIs.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper studies in-context learning from an interesting perspective: in-context approximation theory.\n- Abstracting neural dynamics in a programming language is an elegant way to prove the approximation results of the paper.\n- The paper is overall very well written and easy to follow.", "weaknesses": "- The review of the different recurrent architectures is inaccurate. For example, Mamba and Hawk do not have an A matrix that is constant: it is x-dependent. On top of that, this matrix is diagonal with real values, which makes the implementation of the rotation algorithm presented in 4.2 impossible. More clarity on these questions, starting in Section 2 is necessary.\n- Some synthetic numerical experiments involving the training networks to solve the task, e.g. comparing the solution found by gradient descent to the one of the construction or confirming the importance of gating to get more compact solutions, would be an interesting add-on. I am aware that it is a theoretical paper and I am not expecting such experiments to recommend acceptance. Yet, it would still nicely complement the paper.\n- Section 6 does not bring much to the story.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The review of the different recurrent architectures is inaccurate. For example, Mamba and Hawk do not have an A matrix that is constant: it is x-dependent. On top of that, this matrix is diagonal with real values, which makes the implementation of the rotation algorithm presented in 4.2 impossible.\n\nThe reviewer is right that Mamba and Hawk do have an $A$ matrix that is input-dependent and diagonal.\nHowever, the $A$ matrix of the Linear RNN does not directly map to the $A$ matrix of the Mamba/Hawk model.\nThe Mamba/Hawk has a number of channel-mixing operations which we can leverage to this end.\nThis is the basis of our proof in Appendix E that any Gated Linear RNN can be expressed as a Hawk or a Griffin model.\nOur construction respects their diagonal transition matrix structure, as can be seen in Eq. 43.\n\n> Can the authors position their work compared to: Zucchet et al. 2023: they study the link between gated RNNs [...] and attention, therefore connecting to the literature on how attention solves in-context learning. [...] \n\nThank you for bringing the work of Zucchet et al. 2023 to our attention. The connection is quite interesting and possibly opens a different route for proving the same properties via transformers. However, while it has been shown that transformers with softmax attention are universal in-context approximators (Petrov et al., 2024), to the best of our knowledge this hasn't been proved for linear attention yet. However, our intuition is that this is likely possible and a proof for the universal in-context approximation properties of linear attention transformers might be possible via our present work. Still, it is not clear what the prompt complexity in this case would be: whether it would be the most efficient regime for fully recurrent models from this paper or the less efficient regime for transformers from (Petrov et al., 2024).\n\n> I am wondering about how fundamentally different is in-context approximation compared to the universal approximation. Is one a subset of the other? How are they related / what are their difference? [...]\n\nUniversal _in-context_ approximation is a very different beast than the classical universal approximation. \nClassically, one can choose the model parameters conditional on the target function. \nWith universal _in-context_ approximation, however, the model parameters are _fixed_ and independent of the target function. \nIn other words, a single model needs to approximate all target functions (given a prompt that specifies the target function).\nIn our experience, proving _in-context_ properties has been more difficult because of the complex interactions between different prompt tokens.\nWhile it might appear that universal _in-context_ approximation is a more demanding property than universal approximation and hence should be a subset of it, the two hypothesis classes could be very different and hence we could not formally define a notion of an inclusion or subset. For example, _in-context_ approximation makes sense only for sequence-to-sequence models. However, intuitively, it does feel that if an architecture could be a universal _in-context_ approximatior, then it should also be able to be a universal approximator.\n\n> What is the interest of introducing such a low-level programming language? It feels to me that LSLR is a tiny abstraction on top of the neural equations that is not necessarily worth it. Why not directly include the primitives used in e.g. Figure 3 in this language?\n\nGreat question! We found the bookkeeping required to work directly with the layer equations to be very messy, error-prone and not particularly intellectually stimulating. LSRL helped us streamline the proof process significantly. In a way, you could also think of it as a proof assistant of sorts, where one specifies the high-level strategy and LSRL fills in the formal details.\n\n> The paper shows that deep linear RNNs can implement a variety of functions in context, somehow supporting the argument that linear recurrence + nonlinear instantaneous processing is quite powerful. Does the framework provided by the authors give hints about which kinds of functions are difficult to learn with such networks? [...]\n\nFor any neural network architecture, one can always construct functions that are difficult for it to learn.\nEspecially in the current \"war\" between Transformers and Linear RNNs, the literature is full of carefully constructed examples that trip one or the other.\nLinear RNNs will always be bad at recall within large context sizes due to their fixed-size hidden state.\nNo non-linear activations, gating, or any other fancy mechanism can fix this problem without breaking the full recurrence and bringing some flavour of attention.\nThat is why we put the query before the prompt, as that changes to problem into search rather than memorization and recall (See footnote 1).\nSimply flipping the order, i.e., putting the query _after_ the prompt, would be not possible for arbitrary precision and fixed hidden state.\n\n> Are there some links between the in-context learning approximation notion introduced here and the notion of controllability in control theory?\n\nAnother very good question! We did look into it but there is a key difference in the setup. With controllability, one seeks a control signal $u$ that can bring a (linear) system to a target state $x'$. However, $u$ depends on the target state $x'$. Instead, in the universal approximation setting, we are looking for a single control signal that is a map from _all possible inputs_ to their respective target states. Controllability would be more closely related to the ability of a model to be prompted to produce any desirable output $x'$. Note that our universal in-context approximation results subsume this prompt controllability setting as one can always define a constant function that produces $x'$.\n\nA. Petrov, P. HS Torr, and A. Bibi. Prompting a pretrained transformer can be a universal approximator. In ICML 2024.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "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": {}}]

Author rebuttal:

Meta-review: This work explores the ability for recurrent neural nets to be universal approximators through prompting.  A major question of interest is what potential benefits there are to transformer-based architectures vs. state space models like RNNs.  One potential explanation is a difference of approximation-theoretic: perhaps transformers are simply better function approximators.  This work shows this is not the case since both are capable of being universal approximators in-context.   The majority of the reviewers were in support of this work, and I concur, and recommend acceptance.
Decision: Accept (poster)