B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory

Thank you for your feedback and your insightful comments. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns.

Related work References. We are grateful for the references and associated insights, we have revised our manuscript to incorporate and discuss these works. We acknowledge that there are many relevant related works, our goal is to not propose yet another one, but rather a general framework that encompasses them using Stochastic Realization Theory. Indeed, the referenced hybrid models are special cases of B’MOJO, for example Block State Transformers can be obtained by removing B’MOJO’s the eidetic memory. Similarly, GSS can be obtained by removing the eidetic memory and some of the sliding window attention layers. We have incorporated more discussions in our related work section.

Unique Insights Gained From Our Framework. 1) Long term Eidetic Memory is unique to our work and provides crucial performance gains over models without it (see Figure 3 for perplexity gains at scale with eidetic memory, and Table A1 with results on NVIDIA’s RULER benchmark). 2) Our use of interleaved memory chunking, efficient innovation selection and input dependent SSMs (like Mamba) as opposed to slow long convolutions (like the ones used in Block-State Transformers) are specific innovations to ensure our method runs efficiently on modern hardware (see Section 3.3 and Appendix B2). These are important differentiators which we hope could inspire further developments in hybrid SSM architecture design.

Difference between B'MOJO/B’MOJO-F and Hybrid baseline. Succinctly, a vanilla hybrid layer can be written as “output = Attention(SSM(input_tokens))”, while our B’MOJO-F layers can be expressed as “output = Attention(Cat(Inputs, SSM(past inputs)))”. In particular, B’MOJO-F is not a simple stacking of SSM (Mamba) and Attention as in our baseline hybrid model and some recent hybrid models in the literature (like Jamba/Griffin). While vanilla hybrid models are limited to interleaved stacking of SSM and Attention layers, B’MOJO-F allows the attention module to attend to both the input tokens (to the SSM) and the output representations of the SSM (the fading memory). A mechanism that allow the Attention layer to merge information from "memory tokens" to the layers' input tokens. Differently from B’MOJO, B’MOJO-F does not use the innovation selection mechanism (no long term eidetic memory). Please see Figure 6 in the appendix for more details.

Motivation of the Innovation Selection Mechanism. The "state" of a stochastic realization is a sufficient statistic for (next token) prediction, meaning that the state is as good as the entire past for the purpose of minimizing the prediction error (i.e. making the prediction error unpredictable). Thus, the predictability of the prediction error, which is what the innovation test measures, is the natural test to know if the state has captured the past well enough. In practice, whenever the innovation test measures high residuals we know that the state must be refreshed with new information which in turn will help the predictor in reducing future prediction residuals.

Innovation Selection's efficient implementation. To have an efficient implementation we picked the squared loss and coupled it with a linear predictor which takes as input the most recent past states. However, having a learnable linear predictor would increase the number of parameters of the model and it would require to modify the training loss. To avoid this, we restrict the predictor to a moving average which we implement with short 1D grouped convolutions. To preserve efficiency, we leverage the same 1D convolution implementation used in the Mamba layers which uses a kernel size of length 4.

Experimental Results seem to be generally weak. When performing apples-to-apples comparisons, our largest B’MOJO model (1.4B) outperforms Mamba 1.4B in pre-training perplexity by 10% and by 1% on zero shot downstream evaluations (see Table 1, Figures 3, 5). Furthermore, as suggested by the reviewers we assessed the long context recall capabilities of B'MOJO on a more difficult needle-in-a-haystack benchmark (see the global comment). In these recall intensive tasks we found that B'MOJO significantly outperforms Mamba by 4% accuracy @2k context sizes and by 11% @4k. As well as outperforming our Transformer model on longer contexts sizes.

Synthetic MQAR experiments. We use a sequence length of 256 and an attention window length of 32 so that key-value pairs stay outside the sliding window context. We tried to strike a balance between detail and readability. We plan to make our work reproducible not just by listing numbers, but by releasing our source code upon completion of the review process. As we show in Figure 2, Panel 4, adding eidetic memory tokens (the difference between B’MOJO and B’MOJO-F) steadily increases performance.

Why is B'MOJO faster than Mamba and the Transformer baseline? B’MOJO is faster than the Transformer baseline because it replaces MLP layers with Mamba layers. In Table A3, we report the profiling results for different sequence lengths and show that MLP layers are slower than Mamba layers as the sequence length increases. Furthermore, B’MOJO is faster than the Mamba baseline since it replaces half of the Mamba layers with sliding window attention layers of length <1k. In Table A3 we report the profiling results for different sequence lengths (a similar observation can be found in Figure 8 of the original Mamba paper).

Table A3: Profiling time (in ms) of a forward call of basic blocks, measured on A100 40Gb

Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns.

Confused about the structure of the block, and the overall architecture. We have made revisions to improve the clarity of our work. Please see Figure 6 and 7 in the appendix for more details.

B'Mojo algorithm block (algorithm 1) refers to a single BMojo layer. Yes.

Odd that every layer in your architecture (even the first layer) is predicting the final output tokens. Each block is only predicting representations that are passed to the subsequent block, and not predicting the final tokens. Only the final block predicts the final tokens.

How are the spaces, x, y defined in a more formal sense. The spaces x and y are indeed \in \mathbb{R}^d, where d is the embedding dimension of the Mamba block.

Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns.

Mistral 7B seems to perform comparably or even better. As we clarify in general comments, in a true apples-to-apples comparison in short contexts a Transformer represents a paragon, not a baseline. The primary strength of the B’MOJO model relative to Transformers is in the long context regime. However, performance in long context tasks is sensitive to scale, especially for hybrid state space models like B’MOJO, and in Table 2 we provided experiments only at a scale that is reproducible in academic environments. However, in order to address the reservations raised in this review, we also ran larger-scale experiments up to 1.4B scale (Table A1) featuring NVIDIA’s RULER benchmark using our larger models: B’MOJO-F 1.4B (with only fading memory) and B’MOJO 1.4B (with fading + eidetic memory), where we show that both B’MOJO and B’MOJO-F outperforms Transformer models on long context tasks as expected.

Eidetic memory can in principle be unbounded, is there a criterion to select a specific subset? In the proposed implementation we bound the number of tokens to a maximum size (to keep the cost of attention manageable). However, the token span can be arbitrarily large. We are yet to explore a selection mechanism (e.g. Landmark attention).

Clearly state the dimension of each matrix/vector. We apologize for the confusion and have revised our draft to include dimensions of each matrix and vector.

Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns.

Lack of sufficient ablation study. Throughout our work (see Table 1, 2, and Figures 2, 3, 4, and 5), we compare B’MOJO to BMOJO-F (B’MOJO without Long-Term Eidetic Memory) in order to ablate the contribution of Long-Term Eidetic Memory, we further compare B’MOJO-F to a vanilla Hybrid Baseline (B’MOJO-F without Fading Memory) in order to ablate the contribution of Fading Memory, and we ablate Hybrid Baseline with Mamba to ablate the contribution of Short Term Eidetic Memory. The contribution of permanent structural memory is trivially ablated by the fact that trained weights outperform untrained ones. This covers all needed ablations: Overall, we find that B’MOJO > B’MOJO-F > Hybrid baseline > Mamba, ablating the role of each individual component.

Impact of short-term eidetic memory. Short term eidetic memory in-context refers to the most recent set of tokens (sliding window) processed by the model, analogous to the context of a standard transformer model. For typical language modeling it is impactful --- it is used by the model to process and recall information from the most recent past (akin to a n-gram model). However it has no impact on long range recall capabilities beyond the context length. In both Tables 1 and 2, we refer to long-term eidetic memory, not short-term.

Adding eidetic memory does not improve performance. Our scaling law results, see Figure 3 show that the eidetic memory has a positive impact on B’MOJO, which decreases perplexity as we scale the model size. To further isolate the eidetic memory gains we followed reviewer’s suggestions and report in Table A1 results on NVIDIA’s RULER benchmark using our larger models: B’MOJO-F 1.4B (with only fading memory) and B’MOJO 1.4B (with fading + eidetic memory). We show that B’MOJO-F is strictly weaker and the gap is higher the longer the required context length. Complementing the perplexity results, this showcases that our innovation selection mechanism indeed helps the model to recall specific information from the past. Furthermore, this is in line with our synthetic experiments in Figure 2 albeit at larger scale. Note that in Figure 2 panel 4 we show that increasing the number of eidetic memory tokens produces an increase in recall performance prior to saturation.

The main advantage of the proposed architecture Our proposed model’s main advantage over Transformers is that it simultaneously 1) introduces linear inference complexity with respect to context length and constant KV cache while 2) preserving Transformers’ performance without incurring in the quadratic complexity cost of Attention.

Acceleration methods. Both GQA and Kernel Fusion speed up Transformers by efficiently moving data in memory, but neither change the fundamental quadratic inference complexity of Transformers. Even for moderate context lengths these acceleration methods cannot close the linear vs quadratic gap between B’MOJO and Transformers. We make this point empirically in Table A2, where we report profiling results on both our Transformer baseline and B’MOJO with and w/o GQA at the 1.4B/3B scale for different context lengths (1k/2k/4k). Note GQA can be applied to all attention layers in B’MOJO for further acceleration. In all cases, B’MOJO is still faster than Transformers w/ and w/o GQA. GQA does help Transformers more than B’MOJO, and the relative speed gap for, e.g. at the 1.4B scale, with 2k context length case is reduced from ~10% to 7% . As expected, B’MOJO’s improves more as longer sequences are used due to the quadratic vs linear scaling of attention on the Transformer baseline.

Empirical results B’MOJO vs Mamba vs Transformer. When performing apples-to-apples comparisons (i), our largest B’MOJO model (1.4B) outperforms Mamba 1.4B (see Table 1, Figures 3, 5 and the new experiments with the needle-in-a-haystack results in Table A1, in the global comment), a trend observed in our synthetic experiments too (see Figure 2). While prior works like Jamba (and the recently released Zamba) slightly outperform pure Transformer models they used full attention, and thus retain the quadratic dependence of Transformers. In contrast, our work only uses a small 512 token sliding window attention to produce a model with linear dependency on the sequence length with a constant KV cache size. Other works like Griffin and Samba (which has been published after the submission of this manuscript) also use sliding windows, but manage to slightly outperform Transformer leveraging much longer sliding windows sizes.

Impact of the window length and profiling of the KV Cache. The window length measures short-term eidetic memory and is similar to the impact of changing the window size on Transformer models like Mistral, the larger the better the results but at a higher FLOPs count, all our experiments leverage a fixed window size of 512. B'MOJO's forward time is constant with respect to the sequence length, and its KV cache is comprised of the sliding window cache, the fading memory and eidetic memory tokens (see Figure 7 for more details), whose number is fixed by the user.

Table A2: Profiling time (in ms) of a forward call for various model sizes, batch size = 1, measured on A100 40Gb