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

We describe a family of architectures to support transductive inference by allowing memory to grow to a finite but a-priori unknown bound while making efficient use of finite resources for inference. Current architectures use such resources to represent data either eidetically over a finite span ("context"in Transformers), or fading over an infinite span (in State Space Models, or SSMs). Recent hybrid architectures have combined eidetic and fading memory, but with limitations that do not allow the designer or the learning process to seamlessly modulate the two, nor to extend the eidetic memory span. We leverage ideas from Stochastic Realization Theory to develop a class of models called B'MOJO to seamlessly combine eidetic and fading memory within an elementary composable module. The overall architecture can be used to implement models that can access short-term eidetic memory"in-context,"permanent structural memory"in-weights,"fading memory"in-state,"and long-term eidetic memory"in-storage"by natively incorporating retrieval from an asynchronously updated memory. We show that Transformers, existing SSMs such as Mamba, and hybrid architectures such as Jamba are special cases of B'MOJO and describe a basic implementation, to be open sourced, that can be stacked and scaled efficiently in hardware. We test B'MOJO on transductive inference tasks, such as associative recall, where it outperforms existing SSMs and Hybrid models; as a baseline, we test ordinary language modeling where B'MOJO achieves perplexity comparable to similarly-sized Transformers and SSMs up to 1.4B parameters, while being up to 10% faster to train. Finally, we show that B'MOJO's ability to modulate eidetic and fading memory results in better inference on longer sequences tested up to 32K tokens, four-fold the length of the longest sequences seen during training.

Paper

References (58)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer MJ985/10 · confidence 4/52024-06-23

Summary

Introduces a class of models called B'MOJO which combines sliding window attention with SSMs. The main version of B'MOJO leverages an error function to decide what gets added to the long term sliding KV cache, inspired by ideas from Stochastic Realization Theory. Experiments are performed on multi-query associative recall (MQAR) and well as language modeling.

Strengths

- Considers the important and relevant problem of balancing memory and efficiency in modern sequence models - Provides a nice, high level overview of stochastic realization theory and its connections with modern sequence models - The idea of using an error function inspired by the idea of an innovation process to determine a storage/eviction policy in sliding window attention appears to be novel

Weaknesses

- There are several related lines of work that are not discussed well enough and/or should be cited: - Block state transformers: https://arxiv.org/abs/2306.09539 and SeqBOAT: https://arxiv.org/abs/2306.11197, considers using an SSM to contextualize information that is input to a block-windowed/sliding window attention mechanism, quite similar to the setup in this paper - GSS (https://arxiv.org/abs/2206.13947), H3 (cited in paper but not in the context of hybrids) both propose combining attention with SSMs to try and balance eidetic and fading memory, while hybrid attention methods such as Longformer (https://arxiv.org/abs/2004.05150) are also relevant - Works on adaptive kv caches and optimal eviction policies are also relevant, a few but incomplete list of examples: https://arxiv.org/abs/2310.01801,https://arxiv.org/abs/2401.06104, https://arxiv.org/abs/2402.18096, https://arxiv.org/abs/2402.06262 - There is a brief discussion of input dependency in recent SSMs in Line 135, but https://arxiv.org/abs/2209.12951 should also be cited here as it was the first of this recent wave of linear RNN/SSM papers to propose efficient input-varying systems. https://arxiv.org/abs/2311.04823 is also relevant here. More generally, it could be helpful to point the reader to the SSM papers that led to the more recent variants mentioned in line 135, e.g. S4 (https://arxiv.org/abs/2111.00396), S5 (https://arxiv.org/abs/2208.04933), LRU (https://arxiv.org/abs/2303.06349) and their variants to better position this work in its proper context. - While the paper does a nice job framing the modern sequence models within the ideas of Stochastic Realization Theory and presenting its history, in my opinion, it fails to really justify what insights are being gained from this framework. It discusses fading memory (as in SSMs) and eidetic memory (as in attention) and then proposes a way to combine these, but the result is not that different from the hybrids mentioned above, which were motivated by the same thing. - I think the opportunity for this was in the Innovation Selection process, however the actual version discussed in lines 219-223 does not seem to be described very well. See Questions below. - I have other questions below which may help me better appreciate the value of this framework. - The experimental results seem to generally be weak - MQAR: - Too many details are missing to judge the usefulness of this experiment. E.g. What task sequence length vs sliding window lengths were used? - It does seem the B MOJO models provide an edge compared to the other efficient baselines, but I am skeptical since the BMOJO-F appears to do almost as well as B-MOJO. Why is this? information about sequence lengths and sliding window lengths would help to better understand this. What happens as the task is made even more difficult with more key value pairs and longer sequences (compared to the sliding window length)? - Language task: - What dataset was used? I do not see this information anywhere. - It would have been nice to see a hybrid model with a few full attention layers (not just sliding attention), to better asses the potential performance vs compute advantages that are being claimed. - The Downstream tables are hard to read since bolds and underlines are not used to denote top scores, but in general B'MOJO's performance underperforms Transformers and also seems to not significantly outperform Mamba. - It can be difficult to asses how much actual exact recall vs parametric knowledge is used to perform the long context tasks considered. Perhaps using some previously proposed recall stress tests for pretrained models such as passkey/phonebook retrieval (e.g. as in https://arxiv.org/abs/2402.19427, https://arxiv.org/abs/2402.01032 ) or some of the more difficult needle in a haystask tasks proposed in https://arxiv.org/abs/2404.06654 would help to better assess B'MOJO's long context recall ability.

Questions

Many of my major questions are listed in the weaknesses. Here are a few others. Major: - Could you please make the difference between B'MOJO-F and B'MOJO more explicit? This is not explained well. I think B'MOJO-F is a simple combination of SSM and Sliding sindow attention? What is the difference between this and the Hybrid baseline in the experiments? - Could you please explicitly explain what is the actual innovation selection mechanism used and what is its motivation? The paper has a nice buildup of ideas from stochastic realization theory, but then in lines 219-223 just says a short convolution is used. How long of a convolution? What error function? What is the motivation for why this could work? It would improve the paper it the connection with the motivation from stochastic realization theory was better developed. - It is not explained well why B'MOJO is so much faster than Mamba or the Transformer baseline, since I think it uses the Mamba and Flash attention kernels. Is this related to the sequence lengths considered or the chunking? Why would it be faster than Mamba? Other: - I wouldn't refer to the Transfomer baseline as "Mistral" in the figures or Tables since to many readers this will sound like the models released by the company, as opposed to an architecture that is similar to their implementation that you trained.

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

Adequately discussed.

Reviewer PLgm6/10 · confidence 2/52024-07-12

Summary

I appreciate the clarification about the method being specified for a single block as well as the new longer context results. I have decided to increase my score. --- This seems like a very interesting paper which proposes a new recurrent architecture that seeks to combine advantages of transformers and other recurrent architectures such as Mamba. The results in this paper are quite nice, and seem to systematically outperform Mamba. In particular, the results on OOD length generalization are great. There's also an innovative idea of selecting the highest error tokens to add to the set of tokens to be attended to. This is where I have some concerns with the paper. I found it difficult to follow the distinction between the individual BMojo blocks and the stack of BMojo blocks. I think this could easily be addressed in the algorithm block, or even better, by adding some pseudocode to make the computation easier to follow. It's possible that it's my own fault for not understanding this, but if it could be clarified well, I'm open to raising my score. notes from reading the paper: -New recurrent architecture called BMojo. -Adds eidetic and fading memory, aiming to outperform Mamba. -Substantially improved OOD length generalization. -Transductive inference for sample-specific inference. -"Unpredictable" tokens are added to a sliding window that is attended over. -Mamba has only fading memory.

Strengths

-The introduction and exposition are both well written. -The improved OOD generalization results are impressive. -The idea of adding the most surprising tokens to be attended seems like an interesting idea.

Weaknesses

-While reading this paper, I got confused about the structure of the block, and the overall architecture. This could be my fault, but it's possible that other readers will also get confused.

Questions

-In the B'Mojo algorithm block (algorithm 1), I read this as referring to a single BMojo layer, which is then stacked multiple times to yield the final architecture. Is this correct? If so, it's a bit odd that every layer in your architecture (even the first layer) is predicting the final output tokens (y_t)? If so, this strikes me as fairly unusual? I think it would help to have something like the block in Figure 9 to clarify the architecture better. -Could you say how the spaces $x$, $y$ are defined in a more formal sense (e.g. something $x \in \mathcal{R}^d$? I think this wouldn't take much space and it would benefit readability.

Rating

6

Confidence

2

Soundness

3

Presentation

2

Contribution

4

Limitations

The limitations seem to be fairly discussed.

Reviewer Wgie7/10 · confidence 4/52024-07-12

Summary

The paper presents B'MOJO, a novel building block combining the strengths of Transformers and modern SSMs. The main motivation of the paper is to develop a system capable of combining an eidetic memory, responsible for performing transductive inference via in-context learning, and a fading memory capable of storing information from the past into a finite-size (hence lossy) state. To do so, the paper shows that Transformers and Mamba can be described as different parametrizations of simple dynamical systems, the former requiring to increase its state as its span gets larger, the latter having a fixed state and hence implementing a lossy compression of the past preventing exact recall. B'MOJO results from a natural combination of the state updates of the two aforementioned models, bypassing the limitations of both. The final model ultimately consists of a sliding window attention operating on three different memory sources: the recent past (most recent input tokens), a lossy compression of the distant past (computed via an SSM update) and the most informative tokens from the distant past. The latter memory source is obtained via "Innovative Selection": the tokens that are difficult to predict given the past are stored in a - possibly unbounded up to hardware limitations - external memory. The model is compared to Transformers, Mamba and hybrid models on synthetic tasks for associative recall, language modelling scaling laws, zero-shot evaluation on short and long range tasks and length generalisation. The proposed model performs favourably compared to Mamba and comparably with Mistral on most tasks.

Strengths

* The paper is well-written and generally pleasant to read. * The idea of combining eidetic and fading memory, while not new, is very interesting and B'MOJO represents an original and elegant way to do so. * B'MOJO is shown to integrate the benefits of transformers and modern state-space models. The model compares favourably with Mamba on most tasks. While inferior compared to Mistral7B on zero-shot tasks, B'MOJO inherits its length generalization performance from SSMs, outperforming Mistral7B in this latter case. The model is shown to be marginally better than baselines in terms of training time.

Weaknesses

* Generally Mistral 7B seems to perform comparably or even better than the proposed model, with the exception of the length generalization task. In particular, Mistral is often significantly better on the zero-shot evaluation tasks, also on those involving long contexts. Do the authors have an explanation for this? * It is not clear how the information from the various sources of memory are aggregated in the sliding-window attention mechanism. In particular, given that the size of the eidetic memory $M$ can in principle be unbounded (up to hardware limitations), it is not clear to me how such a large memory can be efficiently processed by attention. By looking at the Appendix, it appears that the number of tokens taken from $M$ is bounded. How are tokens selected from this memory? Is there a criterion to select a specific subset of them? * Some parts of the writeup could be made clearer: for example, in the "Transformer" part in section 3.1, the dimension $V$ is not introduced before in the text. Generally it would be helpful to clearly state the dimension of each matrix/vector for the sake of clarity.

Questions

See weaknesses section.

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors discuss the limitation of their work in Section 5.

Reviewer majr3/10 · confidence 4/52024-07-15

Summary

This work presents a new module by combining eidetic memory, fading memory, and long-term eidetic memory (through an innovative selection operation). The proposed new module has strong sequence modeling capacity and high inference efficiency. The proposed new architecture achieves perplexity comparable to Transformers and SSMs with promising long-sequence ability.

Strengths

* This work provides analysis for both attention and Mamba and proposes "B'MOJO-F" by combining sliding window attention and Mamba. * In order to further improve the memorization/recall ability, the Innovation Selection is proposed to compensate for the lossy memory in fading memory and add important token information to the long-term eidetic memory. The Long-term eidetic memory can increase as the input sequence length increases but could be way more efficient than standard attention. * The resulting architecture shows better memory efficiency on Associative Recall tasks than Mamba, although it still underperforms Transformers.

Weaknesses

* Lack of sufficient ablation study. This work claims that it leverages 4 kinds of memory: short-term eidetic memory 'in-context,' permanent structural memory 'in-weights,' fading memory 'in-state,' and long-term eidetic memory 'in-storage'. Specifically, how does the short-term eidetic memory impact the capacity for recall and language modeling? Additionally, in tables 1 & 2, for "BMoJo (Fading + Eidetic)", which eidetic memory is this referring to? It seems that adding this eidetic memory does not improve performance. * It seems that the new architecture achieves similar benchmark average accuracy and perplexity compared to pure Transformer. Although the authors claim that it is 10% faster than Mistral, there are many acceleration methods for pure Transformer models like GQA and kernel fusion which may mitigate the gap. What is the main advantage of the proposed architecture against Transformer?

Questions

* For the hybrid model baseline, what is the architecture (activation, layer ordering, FFN size, and attention-Mamba ratio)? * In previous work, usually hybrid models (e.g., Griffin, Jamba, Samba, Zamba) can slightly outperform pure Transformer or pure Mamba. Why does it perform worse than Transformer or Mamba in this work? * How is the proposed model's performance on real-world retrieval tasks like phonebook lookup and needle-in-a-haystack? * What is the impact of window length for the sliding window attention? * How is the profiling of the "KV Cache" of the proposed method in terms of sequence length?

Rating

3

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

Please check the above.

Reviewer MJ982024-08-09

Thank you for your comments and clarifications. - On providing the sliding window length and the comment _"We tried to strike a balance between detail and readability."_: Please make the sliding window length clear throughout the paper (e.g. in figure captions, experiment descriptions etc). Understanding the sliding window length is crucial to assess the performance of the method on the recall intensive tasks, since recall tasks that fall within the sliding window length should be solved easily, so being able to solve tasks that extend beyond the sliding window length is what is interesting. - Regarding the note above and the new Ruler experiments: Thank you for including these. Can you please clarify again the sliding window length used for each of these experiments? Is it also possible to include the results for the trained hybrid (sliding window + attn) method? Could you also include a random guessing baseline for each task and average? - It is worth noting that for the 2048 context length on the new Ruler tasks, there is generally a severe dropoff in performance from the full attention method to the SSM and sliding window hybrids. - This dropoff is concerning, because often in practice, methods will not be deployed in zero shot extrapolation settings and the broader concern is high quality within the context length it was trained. This limitation is of interest to the community and should be thoroughly discussed to strengthen the paper, not ignored. - In addition, it is not entirely clear how significant the performance boost the proposed methods have over the Mamba method? Is a 1-3% boost for the 2048 context or the 2-11% boost for the 4096 sequence significant? Especially in the context of the ~28% boost the 2048 context Transformer has over the other 2048 context methods. What is the random guessing baseline?

Authorsrebuttal2024-08-10

Thank you for your comments, we hope the following helps adding more clarity. **Sliding Window Length**. Our RULER experiments used 1.4B pre-trained models. The sliding window sizes are: Transformer Baseline — 2048 tokens; B’MOJO, B’MOJO-F and the hybrid baseline — 512 tokens. B’MOJO’s modules never see a sliding window longer than the Transformer’s context length. **Hybrid baseline**. We have added results in Table B1 in this comment (to complement Table A1 which we copy here). Overall, we find the hybrid model (w/512 token sliding window attention) improves over mamba by a relative 1% @2k and 3% @4k. However, it is slightly weaker than our BMOJO-F (512 token window) and significantly weaker than our full B’MOJO model (512 token window), with a relative performance gap of 7% @2k and 55% @4k. **Random baseline**. Each task in the RULER benchmark requires generating some specific subset of tokens mentioned in the context, e.g. a 5 digit number. A typical Needle-in-a-Haystack (NIAH) example follows this template: "Some special magic numbers are hidden within the following text. Make sure to memorize it. I will quiz you about the numbers afterwards. \n{context}\n What are all the special magic numbers mentioned in the provided text? The special magic numbers mentioned in the provided text are" (please see Table 2 in the RULER paper). A random baseline (in this case) has to correctly guess 5 numbers out of the vocabulary size, the probability of a correct guess is 1/(10)^5. Harder cases include multiple words, or uuids, which have even lower success probabilities — in practice we measure 0%. **Concerning Drop w.r.t Transformers**. The drop in performance from using full attention on the 2k context tokens is not concerning, but expected when using a sliding window approach that only leverages 512 tokens. Indeed, as you note above “being able to solve tasks that extend beyond the sliding window length is what is interesting.”; we agree, the Transformer baseline is the paragon in the 2k setting. To further show this, we also evaluate our models on smaller context sizes 512 and 1024, see results in the table below. At size 512, the gap with full attention is indeed null and the gap increases only slightly at size 1024. However, longer contexts set B’MOJO apart from a Transformer model: the latter’s recall performance goes to zero if tested on a context length longer than its attention span, while our models still can recall information from contexts that are up 8x longer than the attention span. **Extrapolation not often deployed in practice**. Although it is true that often in academic benchmarks, the information supporting the query fits in context, this is not true in many business applications, where the relevant context can be thousands to millions of documents, lines of code, metrics, tables, datasets, and other data that would most definitely not fit in 2048 tokens. With B’MOJO, we are developing a class of models that can cover this long tail of tasks, since the Transformer does not. If in a particular application, 2048 tokens capture the majority of use cases, we would recommend that B’MOJO’s sliding window be set to that value. This way, a practitioner attains the best of both worlds. **Performance boost**. See [Random baseline] above, and [Concerning Drop w.r.t Transformers]. For Mamba, perhaps looking at relative percentage performance is more revealing. Our B’MOJO model improves over Mamba by 8.5% @2k and 140% @4k relative performance and decreases over transformer by 55% @2k and achieves ~20% accuracy on NIAH at 4k where the transformer model cannot solve the task. ### Table B1: Long context evaluation with RULER (needle in a haystack) | Context Length | Model | S-NIHA | MK-NIAH | MV-NIAH | MQ-NIAH | Average | |----------------|--------------|--------|---------|---------|---------|----------| | 512| Transformer | 100 | 100 | 100 | 100 | 100 | | | Mamba | 100 | 67 | 78 | 53 | 75 | | | Hybrid | 100 | 100 | 100 | 100 | 100 | | | BMOJO-F | 100 | 100 | 100 | 100 | 100 | | | BMOJO | 100 | 100 | 100 | 100 | 100 | | 1024| Transformer | 100 | 97 | 63 | 100 | 90 | | | Mamba | 100 | 44 | 34 | 48 | 57 | | Hybrid| 100 | 53 | 42 | 89 | 71 | | | BMOJO-F | 100 | 59 | 48 | 98 | 76 | | | BMOJO | 100 | 81 | 59 | 100 | 85 | | 2048 | Transformer | 100 | 95 | 62 | 61 | 79 | | | Mamba | 100 | 32 | 29 | 28 | 47 | | | Hybrid| 90 | 35 | 34 | 31 | 47.5 | | | BMOJO-F | 90 | 36 | 35 | 31 | 48 | | | BMOJO | 90 | 45 | 37 | 33 | 51 | | 4096 | Transformer | 0 | 0 | 0 | 0 | 0 | | | Mamba | 9 | 12 | 5 | 7 | 8 | | | Hybrid | 9 | 13 | 5 | 8 | 8.75 | | | BMOJO-F | 10 | 16 | 5 | 8 | 10 | | | BMOJO | 22 | 21 | 17 | 17 | 19 |

Reviewer MJ982024-08-10

Thank you for the clarifications and adding they hybrid baseline to the tables. I personally think the importance and benefits of the extrapolation performance is overstated. To be clear on my previous comment “being able to solve tasks that extend beyond the sliding window length is what is interesting.”, my hypothesis based on other results in the literature mentioned above is that if you pretrained the Transformer on a 4096 context and pretrained the sliding window hybrid and BMOJO methods on 4096 context with 512 sliding window, the gap between the Transformer performance and the sliding window methods would be even greater than the gap between the Transformer and sliding window methods when trained on 2048 context with 512 sliding window. Similarly for training on 32K context (while maintaining a much shorter sliding window). This issue of a performance gap when in-distribution seems more practically important right now then counting on slightly improved (yet still poor performance) extrapolation abilities, since in many cases the practitioner will be better off training the Transformer on the longer context to ensure stronger performance. But this is just an opinion and hypothesis, and to be clear I am not asking the authors to rerun this experiment on longer context due to the resources required. I do think the rebuttal addresses most of my main concerns and I am increasing my score.

Reviewer Wgie2024-08-14

I would like to thank the authors for their response to my concerns and for the additional experiments. The rebuttal addressed most of my concerns and I decided to increase my score accordingly.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC