Two Heads are Better than One: Simulating Large Transformers with Small Ones

We thank the reviewer for the comments. Here are our responses.

Empirical validation: We agree that our paper gives rise to many interesting questions which should be the study of follow-up empirical work. However, there are two main reasons we chose not to include them here.

First, we feel that the theory results presented in our paper stand on their own. The goal of our work here is to perform a complete theoretical analysis of how small transformers can be used to simulate big ones. It is important to nail down these precise dependencies before implementing the approach in practice, and that is exactly what we do here - we provide a complete and systematic analysis, including optimal and yet simple simulation algorithms. Most notably, some of our results show that surprisingly efficient simulations are possible (for example, Theorem 1.1 and Lemma 3.3 show that we can make use of all the layers of an oracle simultaneously), and we’re also able to rigorously pin down when techniques like attention sinks, sliding window, or assumptions on the data can yield real improvements.

We also want to highlight that our results indicate that quadratic/linear amounts of small transformers have at least as much representational strength as a constant amount of large transformers in different settings. Such general representation results can only be proved theoretically.

Second, an empirical validation of our results here that is truly meaningful would need to relate to real LLMs that are trained on large corpa of text. The implementation and parameter tuning can now build on our results, but this would be a substantial undertaking, and likely depends quite a bit on the setting like the reviewer suggests. Our simulation algorithms are simple and even “differentiable”, and therefore could be implemented effectively in practice.

For these reasons, we feel that careful experimental work should comprise a second paper (or perhaps even more than one follow-up paper for different applications).

Writing is redundant; theorems are informal: Although Definitions 2.3-2.5 are similar, we believe that it is important to distinguish causal masking, sliding window and attention sinks as they are all different in practice. In addition, these differences will affect our proof details, so we want to make sure that they are stated as clearly as possible; we need to formally define all the different terms and models, even if some are similar to each other, so that we can formally state and prove all our results. As for Theorem 3.4 and 3.5, we want to emphasize Theorem 3.4 as our main result whose proof is based on the lemmas in section 3, while Theorem 3.5 uses a different set of lemmas that are presented in appendix (supplementary material). The statements may appear similar, but they require substantially different work to prove.

Our main result (Theorem 1.1) is stated formally in mathematical language. We use informal theorem (Theorem 1.2) for average case results because of the page limit, and we believe that this makes our results easier to understand in the introduction. All the precise, formally-stated theorems are presented and proved in the supplementary material.

Typos and formatting: Thank you for the suggestions. We will adjust them accordingly.

Multiplicative error bound for single layer: Our main result (Theorem 1.1) provides an exact algorithm if we assume infinite (full) bit precision. The reason we define simulating as an approximation up to $O(1/2^{N})$ multiplicative error is that, as in most theoretical works in this area, we assume only $O(\log n)$ bit precision to more accurately model realistic hardware, which means that all parameters are represented by $O(\log N)$ bits. In this case, it is impossible to achieve perfect accuracy because the $exp$ function in the attention mechanism will have to be approximated. In other words: this approximation error is unavoidable in limited-precision hardware. Our results provide the best possible approximation in this scenario.

Cannot be improved in the worst case: Our main theorem (Theorem 1.1) states that $O((N/M)^2)$ small transformers suffice, and we discuss in line 109-113 that this is also optimal under standard complexity-theoretic assumptions. So yes we have proved it, but we did not make it a formal theorem because the argument only requires a few sentences.

We hope our responses make sense, and we sincerely hope that the reviewer can reevaluate our work as we believe that our work has its important theoretical values. Please let us know if there are any further questions!

We thank the reviewer for the detailed comments. We address the questions and concerns below.

Empirical validation (and acknowledgement): We agree that our paper gives rise to many interesting questions which should be the study of follow-up empirical work. However, there are two main reasons we chose not to include them here.

First, we feel that the theory results presented in our paper stand on their own. The goal of our work here is to perform a complete theoretical analysis of how small transformers can be used to simulate big ones. It is important to nail down these precise dependencies before implementing the approach in practice, and that is exactly what we do here - we provide a complete and systematic analysis, including optimal and yet simple simulation algorithms. Most notably, some of our results show that surprisingly efficient simulations are possible (for example, Theorem 1.1 and Lemma 3.3 show that we can make use of all the layers of an oracle simultaneously), and we’re also able to rigorously pin down when techniques like attention sinks, sliding window, or assumptions on the data can yield real improvements.

We also want to highlight that our results indicate that quadratic/linear amounts of small transformers have at least as much representational strength as a constant amount of large transformers in different settings. Such general representation results can only be proved theoretically.

Second, an empirical validation of our results here that is truly meaningful would need to relate to real LLMs that are trained on large corpa of text. The implementation and parameter tuning can now build on our results, but this would be a substantial undertaking, and likely depends quite a bit on the setting (data, hardware etc like the reviewer suggests). Our simulation algorithms are simple and even “differentiable”, and therefore could be implemented effectively in practice.

For these reasons, we feel that careful experimental work should comprise a second paper (or perhaps even more than one follow-up paper for different applications). We will modify the limitation section - thank you for pointing it out.

Whether modern GPUs can do faster transformer inference needs to be shown experimentally: Our main argument here is that when processing an extremely long input sequence, it is possible to partition the input sequence into multiple segments and deal with them separately while still obtaining the correct inference output. For example, we prove that a single GPU that can only take input length of at most $M$ can still perform inference on sequences of $N\gg M$ by doing inference multiple times. It is not clear what a straightforward/trivial time (wall-clock, FLOPs, etc) to compare to even is, because it is initially unclear how inference can even be performed with such limited resources.

In addition, the fact that transformer-specific GPUs allow faster inference on short-to-moderate input length has been experimentally verified and is a core idea for many startups, as we mentioned inline 31-32. For example, Etched’s Sohu ASIC claims to be 20x times faster than NVIDIA H100 on Llama 70B (we are not allowed to add links here due to Neurips rules, but the details can be found on Etched’s website).

Number of calls to different models could still be very large: We are not sure if we fully understand what “different models” here mean. We only consider transformers, so as long as all the models are based on the transformer architecture, our analysis will hold. Our assumption here is that the small transformers have all the intricacies of the large ones: if the large transformer has rotary embedding, skip connection, layer normalization for example, the small ones can simulate it if they also have the same components.

Output approximation and bit precision: Our main result (Theorem 1.1) provides an exact algorithm if we assume infinite (full) bit precision. The reason we define simulating as an approximation up to $O(1/2^{N})$ multiplicative error is that, as in most theoretical works in this area, we assume only $O(\log n)$ bit precision to more accurately model realistic hardware, which means that all parameters are represented by $O(\log N)$ bits. In this case, it is impossible to achieve perfect accuracy because the $exp$ function in the attention mechanism will have to be approximated. In other words: this approximation error is unavoidable in limited-precision hardware. Our results provide the best possible approximation in this scenario.

More insights on model size vs gains: Our main results require around $2N^2/M^2$ small transformers (or calls) to simulate a large one, so in theory the larger $M$ is, the fewer inferences we need to perform. We agree that tuning $M$ (and other model parameters) and finding out the regime where we can obtain most gains is an interesting empirical problem to study in the future.

Is the core of the results Lemma 3.1 and comparison to Flash Attention: This is a great question! Yes, Lemma 3.1 is one core of our proof (the other core being Lemma 3.3 where we deal with multiple heads and layers) and it shares some insights with Flash Attention, although they are still different.

1. Similarities: Both Flash Attention and our algorithm deal with a large attention matrix by decomposing it into smaller blocks.

2. Differences: The goals are different: the major concern of Flash Attention is IO complexity, and here we study representational strength. In addition, the main technical challenge is different: in Flash Attention, you load query, key, matrix value blocks into SRAM and do whatever computation you can do on SRAM, while here we can only perform what small transformers can perform. Therefore, our model is more limited.

More insights on our algorithms: We present proof sketches in section 3 for our main results. Although the algorithms for average-case inputs and attention sinks are slightly different, they share similar ideas. Proof sketch of Lemma 3.1 (starting from line 273) best summarizes our ideas. We will make the sketch clearer in the final version of this paper.

We hope that our responses make sense, and make the reviewer believe that our work has its important value even though it is purely theoretical. Please let us know if there are any more questions!

We thank the reviewer for the insightful comments, and answer the questions below.

Empirical validation and parameter tuning: We agree that our paper gives rise to many interesting questions which should be the study of follow-up empirical work. However, there are two main reasons we chose not to include them here.

First, we feel that the theory results presented in our paper stand on their own. The goal of our work here is to perform a complete theoretical analysis of how small transformers can be used to simulate big ones. It is important to nail down these precise dependencies before implementing the approach in practice, and that is exactly what we do here - we provide a complete and systematic analysis, including optimal and yet simple simulation algorithms. Most notably, some of our results show that surprisingly efficient simulations are possible (for example, Theorem 1.1 and Lemma 3.3 show that we can make use of all the layers of an oracle simultaneously), and we’re also able to rigorously pin down when techniques like attention sinks, sliding window, or assumptions on the data can yield real improvements.

We also want to highlight that our results indicate that quadratic/linear amounts of small transformers have at least as much representational strength as a constant amount of large transformers in different settings. Such general representation results can only be proved mathematically/theoretically.

Second, an empirical validation of our results here that is truly meaningful would need to relate to real LLMs that are trained on large corpa of text. The implementation and parameter tuning can now build on our results, but this would be a substantial undertaking, and likely depends quite a bit on the setting (data, hardware etc). Our simulation algorithms are simple and even “differentiable”, and therefore could be implemented effectively in practice.

For these reasons, we feel that careful experimental work should comprise a second paper (or perhaps even more than one follow-up paper for different applications).

Causal masking: We state in our main Theorem (Theorem 1.1, line 107-108) that indeed, a quadratic amount of small transformers are able to simulate a large one, assuming they both have causal masking. A similar argument to the one in line 109-113 will show that this is optimal as well.

KV cache: This is a great question! Our algorithms are indeed compatible with KV cache because it can be used in a standard way for small transformer inference. The trivial space needed will be $O(N^2/M^2)\cdot O(Md) = O(N^2d/M)$ for the main simulation. However, in our actual construction, we only need to store what needs to be stored in a large transformer inference (all keys and values) plus a constant amount of fixed vectors such as the all one vector. As a result, the space needed is $O(Nd)$, which is exactly the same as normal KV cache takes.

We hope our responses make sense, and please let us know if there are any more questions!

We thank the reviewer for the thoughtful comments, and answer the questions below.

Transformer model simplification: This is a great question! We simplify our transformer model for the sake of simpler theoretical analysis, but in general these additional components will not affect our conclusion, and the reason is that whatever component we add to the large transformer, we can also add it to the small transformers in the same way to ensure that our analysis still holds.

1. Rotary embedding: When a large transformer has rotary embedding, each query and key is “rotated” by some fixed rotation function in each attention head. The same rotation function can be applied to each query and key in attention heads in small transformers as well.
2. Skip connection: One can usually think of skip connection as if we have an additional attention head (let’s call it skip attention head) in each layer whose output is identical to input. This additional attention head can be easily simulated by $N/M$ smaller skip attention head because we can feed the input of length $N$ to $N/M$ many smaller skip attention heads. Therefore, if both large and small transformers have skip connections, a large one will still be simulated by small ones (number needed is still the same).
3. Layer normalization: Layer normalization will not affect our analysis because whatever normalization that is performed in large transformers, we can do the exact same in small transformers. Therefore, if both large and small transformers have layer normalization, a large one will still be simulated by small ones (number needed is still the same). We will make sure we mention this in the final version of the paper.

Abstract is limited to one paragraph: Thank you for pointing this out, and we will make sure we fix it in our final version.

We hope that our answers make sense, and make the reviewer more confident about the value of our work. Please let us know if there are any more questions!

Upon careful thoughts on the authors' rebuttal and discussions in the fellow reviewers, I decide to raise my score to Accept and my confidence to 2. Here are my primary considerations:

1. Small models are significant in LLM research and application. The alignment of small and large models is popular, which gives rise to many frontier topics, such as model distillation, speculative decoding, etc. This paper reveals a very interesting aspect of alignment, and the authors successfully provide theoretical guarantee for this. This approach may offer new solutions to many cutting-edge problems, not only accelerating inference.
2. The author successfully identify the primary and secondary components within the complex transformer architecture and simplify it accordingly. This simplification is effective for other components as well, such as rotary embedding, skip connections, layer normalization, and KV-cache. This indicates that the assumption about the problem is not detached from reality.
3. The method proposed by the authors can save computation. However, during the model inference, things are relatively complex, and the inference process is influenced by various factors. I still encourage the authors to provide some numerical results and wall clock speedup, although these may go beyond the domain of this paper and is not considered a huge drawback for the theoretical analysis.