Universal In-Context Approximation By Prompting Fully Recurrent Models

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.

The reviewer is right that Mamba and Hawk do have an $A$ matrix that is input-dependent and diagonal. However, the $A$ matrix of the Linear RNN does not directly map to the $A$ matrix of the Mamba/Hawk model. The Mamba/Hawk has a number of channel-mixing operations which we can leverage to this end. This is the basis of our proof in Appendix E that any Gated Linear RNN can be expressed as a Hawk or a Griffin model. Our construction respects their diagonal transition matrix structure, as can be seen in Eq. 43.

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. [...]

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

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? [...]

Universal in-context approximation is a very different beast than the classical universal approximation. Classically, one can choose the model parameters conditional on the target function. With universal in-context approximation, however, the model parameters are fixed and independent of the target function. In other words, a single model needs to approximate all target functions (given a prompt that specifies the target function). In our experience, proving in-context properties has been more difficult because of the complex interactions between different prompt tokens. While 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.

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?

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

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? [...]

For any neural network architecture, one can always construct functions that are difficult for it to learn. Especially in the current "war" between Transformers and Linear RNNs, the literature is full of carefully constructed examples that trip one or the other. Linear RNNs will always be bad at recall within large context sizes due to their fixed-size hidden state. No non-linear activations, gating, or any other fancy mechanism can fix this problem without breaking the full recurrence and bringing some flavour of attention. That is why we put the query before the prompt, as that changes to problem into search rather than memorization and recall (See footnote 1). Simply flipping the order, i.e., putting the query after the prompt, would be not possible for arbitrary precision and fixed hidden state.

Are there some links between the in-context learning approximation notion introduced here and the notion of controllability in control theory?

Another 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'$.

A. Petrov, P. HS Torr, and A. Bibi. Prompting a pretrained transformer can be a universal approximator. In ICML 2024.

We appreciate the reviewer's positive opinion of LSRL and our investigation into the numerical instabilities without gating.

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.

We have mentioned at the beginning of Section 2 (Line 74/footnote 1) why it is necessary to put the query before the prompt. However, we agree with the reviewer that mentioning this in the Introduction could further improve the readability of the paper.

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?

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

In the conclusion, it is mentioned that the compiled transition matrix is often diagonal. Are those matrices exactly diagonal here?

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

We would like to thank the reviewer for recognizing our work's contributions towards understanding sequence-to-sequence models.

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.

There are two key differences between the classical results the reviewer has shared and this work. First, 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. In fact, it was only last year that it was shown that Linear RNNs can be universal approximators (see Wang and Xue, 2023). Second, and that's the key novelty of our work, universal in-context approximation is a very different beast than the classical universal approximation. Classically, one can choose the model parameters conditional on the target function. With universal in-context approximation, however, the model parameters are fixed and independent of the target function. In other words, a single model needs to approximate all target functions. That 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.

Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. NeurIPS 2023

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.

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?

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

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

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

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

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?

The reviewer is raising some very interesting questions.

It is true that, especially in deep learning, there is often a tension between theory and practice. The 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. However, 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. The 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.

Furthermore, there is no free lunch: for any given architecture, we can design problems that it won't perform well on. In our view, the question should not be to figure out if Linear RNNs or Transformers are the "better" architecture. They offer different trade-offs and are useful in different settings.

We would like to highlight that we have discussed the fundamental properties of such theoretical results in our Limitations section.

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?

We make no assumptions about infinite precision and exponential hidden state sizes. As we anyways use discretization, infinite precision is neither necessary nor helpful for our constructions. The 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). The error bound on line 228 does assume Lipschitzness, but that is a standard condition. Therefore, we work with very few assumptions on the architecture and, hence, our assumptions are practical.

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.

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