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

Authors: luca  zancato
arjun  seshadri
yonatan  dukler
aditya  golatkar
yantao  shen
benjamin  bowman
matthew  trager
alessandro  achille
stefano  soatto
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 13:34:52

Abstract: 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 test whether models trained inductively on a-priori bounded sequences (up to 8K tokens) can still perform transductive inference on sequences many-fold longer. 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.

Main paper: 
 introduction: 
 In Machine Learning, data representations are parametric maps trained to re-present data either individually, by optimizing a reconstruction criterion, or collectively, by optimizing a classification criterion. A trained representation can be co-opted to map a previously unseen datum to a hypothesis, for instance a class label. Representation learning with deep neural networks has been at the core of recent progress in artificial intelligence (AI) during the past decade. This paper is instead concerned with data realizations, which are maps trained to realize, i.e., to make real, bring into existence, or generate data. Roughly speaking, realizations are "generative representations," trained by optimizing a prediction criterion, that can be used as sequential decision or prediction maps, or as generative models for sequence data. Large language models (LLMs) and other predictive models that involve randomness in the sequential generation process are special cases of stochastic realizations [29] .Representations are the backbone of inference from inductive learning, or induction for short. Induction refers to the process of mapping properties of a particular set of training data, through the parameters of the learned map, to properties of general test data. Successful induction, leading to generalization, hinges on an assumption of stationarity, namely the existence of some unknown distribution from which both past (training) data and present (inference) data are independently and identically drawn (IID). While uniform generalization bounds provide provable guarantees for any distribution, they do so under the assumption that the distribution exists and is the same for training and testing. This assumption is routinely violated in practice (e.g., language, climate, and business data), as manifest in so-called "out-of-distribution" effects or "distribution shift." These lead to apparent paradoxes involving generalization and memorization [57] .Realizations, on the other hand, are the backbone of transductive inference and generative AI (GenAI). Transduction refers to the process of inferring particular properties of test data by processing, at inference time, all given (particular) training data. 2 The boundary between induction and transduction is blurry: Induction can be viewed as a restricted form of transduction, where training data is accessible only through the learned weights. An over-parametrized representation, when overfit to training data, could in principle store the training set in the weights, thus making induction functionally identical to transduction. However, optimal induction aims to foster generalization by using various forms of regularization to prevent memorization, leading to sub-optimal transduction. Another form of sub-optimal transductive inference is termed "in-context learning" -which notably involves no learning if by learning one means to "improve by experience:" The same in-context task, presented multiple times, requires identical effort and leads to no improvement. In-context learning can be optimal transduction only if all the data of interest fits in the context (including the entire training set), and even then it has been proven optimal for Transformers only for simple tasks such as linear classification. In summary, memorization and specific inference computation at the core of transduction are in contrast with the biases fostered by inductive learning: Whereas inductive inference seeks to minimize memorization to avoid overfitting and to foster generalization to unseen data, transductive inference seeks to maximize memorization and forgo generalization in favor of sample-specific inference computation. 3Transduction does not require the train and test distribution to be the same. Indeed, the joint distribution from which both present and past data could have been jointly drawn can change with every sample. Therefore, optimality is not measured relative to one unknown distribution, as in generalization bounds, but rather relative to all possible distributions. If the data is generated by a physically realizable process, such distributions are computable. Optimal inference measured on average over all possible computable distributions, weighted by the Universal Prior, has been described by Solomonoff [42] . Solomonoff-style inference can be thought of as the limit of transduction, and similarly involves no learning. Instead, it consists of cycling through all computable programs using a Universal Turing Machine, which requires infinite time, memory, and compute resources, which in turn renders such inference unattainable.Nonetheless, this Solomonoff limit points to two directions for improving inference: (a) efficient memorization, ideally by losslessly encoding all past data, and (b) efficient test-time computation through hardware and model co-design. Ideally, the resulting realizations would be such that, if memory and compute resources were extrapolated to infinity, inference would approach the Solomonoff limit. In reality, inference is always resource bound, so if something has to grow it would have to be external to the core inference engine. Accordingly, our goal in this work is to design and analyze families of architectures that (a) natively incorporate retrieval from a growing external memory (retrieval-augmented generation, or RAG), and (b) can scale to perform efficient inference.In order to design families of architectures that efficiently memorize and scale inference computation, we must exploit the structure of the data they realize. When past and present data are known to be independent and identically distributed (IID) one can encode inference computation through a fixed "stateless" map, or model, which is a representation that can be trained inductively regardless of the inference query. When past and future data are not IID but generated by a Markov process of known order, there exist finite-dimensional statistics, called states, that summarize all past data for the purpose of prediction [26, 23, 29] . But even if the underlying process is Markovian of bounded order, unless such an order is known a-priori optimal realization is generally not possible with constant complexity [33] . To perform optimal inference, memory has to grow, and if the data generation mechanism has finite complexity at some point an efficient encoding of past data into memory will stop growing, but it is not possible to know when [1] . Therefore, a suitable architecture has to always allow the possibility of adding new storage.In the seventies, Stochastic Realization Theory [11] studied State Space Models (SSMs) under the known-order Markov assumption, since realizations with growing memory were unmanageable then. 4 Today, Foundation Models can ingest a large portion of the growing volume of data accessible through the Internet and make it available for inference. Current AI systems are typically hardconstrained by inference time and compute resources, but not by storage-one can always add more disks. To extend Stochastic Realization beyond the IID or known-Markov cases, Foundation Models need scalable architectures that comprise short-term memory updated synchronously within the given computational constraints and long-term memory updated and accessed sparingly and asynchronously. The former includes both eidetic (lossless) and fading (lossy) memory for efficient computation, while the latter is akin to an integrated form of "retrieval-augmented" inference. Such architectures would seamlessly manage short-term eidetic memory "in-context", fading memory "in-state", long-term structural memory "in-weights" and long-term eidetic memory 'in-storage' [13] .Existing architectures such as Transformers and SSMs fall short of encompassing these criteria. Transformer-based architectures use eidetic memory restricted to a finite span, "context length" [6, 24] , while recent SSM-based architectures [19, 16, 51] use only fading memory in their state. In both cases, scaling requires allowing the context (and the key-value cache) or recurrent state to grow unbounded. Recent work on hybrid combinations of Transformer and State Space layers [28, 39, 10] show promise in striking a balance between eidetic memory, fading memory and compute.In this work, we describe a class of models that encompasses both recent SSMs [19] , Transformers [24] , and hybrid architectures [28, 10] as special cases, which we call B'MOJO. This model family simultaneously renders the high expressivity and recall of Transformers, and the high compute efficiency of SSMs. And, rather than assigning tokens to the attention mechanism by recency, an asynchronous selection mechanism assigns tokens based on unpredictability. That is, whenever the model processes a new token that cannot be well-explained it will append it to the registers that implement B'MOJO's eidetic memory-a process we call Innovation Selection.We demonstrate through synthetic tasks that B'MOJO outperforms existing SSM and hybrid model architectures in transductive inference. Empirically, we show that our implementation is 15% faster than similarly-sized Transformers and SSMs while achieving comparable perplexity on language model tasks up to the 1.4B scale. Finally, we show that B'MOJO can operate effectively at inference time on sequences far longer (tested up to 4×) than those used for training. Specifically, experiments with a B'MOJO architecture trained with sequences of at most 8K tokens show consistent length generalization on test sequences of 32K tokens. 
background and related work: 
 We start with a general overview of models for sequence data, whether of text tokens (Large Language Models), images (e.g., Vision Language Models), video, or other physical sensory data (World Models). Any predictive model inferred from a sequence is called a realization. 5 If one assumes that there exists one "true" model, coming from a known model class, then the problem of realization reduces to System Identification [31] , which is to estimate the true model parameters from data. If the model is linear and driven by Gaussian noise, the Cayley-Hamilton theorem [25] implies that the model's "true" order can be inferred with a variety of model selection criteria, including greedy selection. Otherwise, the rank of the non-linear analog of the observability matrix, the Observability Codistribution, can grow in fits and starts, making greedy model selection undecidable [22] . If not only the model class, but the model itself are known, then the problem further reduces to filtering or prediction, which for the linear-Gaussian case can be implemented as a closed-form iterative update [26] . When building a (generally non-unique) realization, all one is given is the data, which leaves complete freedom of choice of model class and order, or number of free parameters. Stochastic Realization. A "sequence model" is a mathematical representation of sequential data capable of predicting the next datum in the sequence given previous ones. The Wiener Filter [48] was among the earliest to be deployed, superseded by the Kalman Filter [26] , the first State Space Model (SSM) with an explicit "state" updated recursively. The ensuing decades saw extensions to more general models leading to Stochastic Realization Theory [29] . The general problem of stochastic realization is, given a (potentially infinite) sequence {. . . , u t-1 , u t } . = u ≤t observed up to time t, infer (i.e., "learn") some parametric model, a function ϕ with parameters θ, ϕ θ (u ≤t ) such that the prediction ût+1 = ϕ θ (u ≤t ) yields a residual ϵ t+1 = u t+1 -ût+1 ("innovation process") that is as close as possible to independent and identically distributed samples (IID). In a nutshell, an optimal predictor is one that makes the prediction error unpredictable.The state of a model is a statistic that makes the future independent of the past. In particular, such a statistic (function of past data) ξ(u ≤t ) yields to the following conditional entropy equality H(u t+1 |ξ(u ≤t ), u ≤t ) = H(u t+1 |ξ(u ≤t )) relative to the joint distribution of all past and present data [31, 30] . Trivially, the data itself fits the definition of state with ξ(u ≤t ) = u ≤t , but it grows unbounded over time. Instead, one seeks a bounded complexity state given which all past data can be ignored with no loss of information. Building such a state is the core goal of stochastic realization.Elementary model classes: LTI, LTV and LIV. Any finite sequence can be realized by a linear time-invariant (LTI) system driven by IID Gaussian noise [29] . It would therefore appear that this model class is sufficient for any practical purpose, and we need models no more expressive than linear time-invariant ones driven by white zero-mean Gaussian noise:6 x t+1 = Ax t + Bu t y t = Cx t + v t LTI x t+1 = A(u t )x t + B(u t )u t y t = C(u t )x t + v t LIVGiven observations y 0:t , stochastic realization deals with the problem of inferring an equivalence class of model parameters A, B, C [44] and a state x t along with a covariance matrix P t of [u t , v t ] that can be propagated deterministically to approximate trajectories produced by the underlying data generation mechanism, y 0:t [49] . However, arbitrarily long and complex time series would require a growing state, therefore making the model no longer time-invariant. More expressive model classes, such as time-varying [23] and input-varying [27] afford more efficient representation by considering A t , B t , C t known functions of time (Linear-Time Varying) or of their input (Linear-Input Varying, LIV). As we describe next, a special case of LIV model class where the dependency on the input is linear, resulting in a so-called bilinear realization, is gathering considerable attention [19, 10, 51] .Modern realizations. Input-dependent (bilinear) models first gained popularity half a century ago starting with Brockett [5] , when [27] showed that "every nonlinear system with controls entering linearly is locally almost bilinear"; [11] developed a complete realization theory for this class of systems including minimal realizations, and showed that they not only minimize the dimension of the state, but also minimize the number of operations necessary for the update. Most recently, special cases of bilinear realizations are being used as building blocks in stacked architectures like Mamba [19, 46] , Jamba [28] and Griffin [10] ; they refer to input-dependency as "selectivity" [19, 20, 37] and combine Attention mechanisms and other techniques to scale up to 52B parameters. Similarly, but at a smaller scale, Block State Transformers [14] , GSS [34] and H3 [16] consider using State Space Models to contextualize information that is later aggregated with Attention. All these models differ by their choices of A, B, C, each with its own advantages and limitations which we describe next. along with with the most recent past. B'MOJO's eidetic memory stores tokens selected from the past using an innovation test on the SSM's state and appends them to the current sliding window. The innovation test measures how difficult it is to predict the next token using the SSM's state. If a token is difficult to predict, we store it in the eidetic memory and pass it to the attention module together with the state, a compressed summary of the past, and the most recent tokens.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We have a Limitations section that explicity describes multiple limitations of the model and our experiments.Guidelines:• The answer NA means that the paper has no 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 a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to 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 on a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
b'mojo: 
 We now introduce B'MOJO, a general family of architectures based on a stackable module, designed to foster transductive inference. We represent fading memory using the state of a dynamical model whose size is fixed a-priori based on hardware constraints. Since the state of a dynamical model is a lossy memory of past data, we implement a complementary eidetic memory with shifting registers that directly encode past data as new information is processed. Although adding storage is simple and cheap, peering back into the growing history with every new query is not. Approximate retrieval [35, 52] also grows as more tokens are processed, potentially sub-linearly if optimized off-line. We propose to access information arbitrarily far in the past at a fixed compute cost by keeping only the most unpredictable tokens according to an innovation test. While we respect hard constraints on the amount of memory processed at inference-time, we impose no constraint on its time span.We begin by describing the elementary components of B'MOJO, which we show to encompass existing models. We first show that the Attention mechanism implements a non-linear nil-potent system whose state grows as more samples are processed. Much like a Moving-Average (MA) system with non-linear read-out, the attention state is required to increase as its span is enlarged. Then, we show that Mamba [19] has a fixed-dimensional Auto-Regressive (AR) (fading) state and hence cannot perform exact recall. Finally, we describe a model that modulates the two forms of memory, B'MOJO-Fading (B'MOJO-F), effectively realizing a non-linear ARMA model [31, 54, 55] .Transformers use Attention to map input to output sequences according toy t = t i=1 exp(qtk T i )vi t i=1 exp(qtk T i ), where q t , k t , v t are the query, key and value vectors and are all computed directly from the input tokens u t . We write this equation as a nil-potent dynamical system with a softmax read-out function ρ as follows:EQUATIONwhereA(u t ) = A ATT =     0 I . . . I 0     ∈ R 2V N ×2V N and B(u t ) = b ATT (u t ) =     0 . . . 0 b(u t )     with b(u t ) := k t v t ∈ R 2V, where 2V is the embedding dimension of the KV cache and N is the length of the Attention window. A Transformer has only short-term eidetic memory that is deadbeat in N steps: what slides "out of context" is permanently removed.Mamba is complementary in that it only has fading memory [19] with decoupled (diagonal) dynamics. Specifically,A(u t ) = A Mamba (u t ) =    a 1 (u t ) . . . a N (u t )    , b(u t ) = b Mamba (u t ) =    b 1 (u t ) . . . b N (u t )   where N is picked a-priori and the output is y t = C(u t )x t + Du t . A Mamba model is obtained by stacking diagonal Mamba layers, hence insufficient to realize even a simple oscillator without resorting to (numerically slow) complex algebra. A marginally stable pole (a i = 1) can be used as permanent memory but at an unbounded cost of encumbering the state and re-processing it at each time step. [19] emphasizes the use of "selective state space," which is simply a bilinear realization [12, 11] . However, crucial to its implementation are interleaved convolutional layers that retrieve some of the eidetic functionality lost with the diagonal dynamics. In Sect. E we derive in detail the form above from the description of the paper, forgoing the unnecessary continuous time narrative, and comment on the model in relation to the actual implementation.B'MOJO-F bypasses the limitations of Transformer and Mamba layers by using a Controllable Canonical Form [25] to realize Equation (1) (see appendix C), that is:A B'MOJO (u t ) =     0 I . . . I a 1 (u t ) a 2 (u t ) . . . a N (u t )     , b B'MOJO (u t ) =     0 . . . 0 b(u t )    In Appendix C we show that B'MOJO-F is a minimal realization and, when the state dimension is fixed, it generalizes both Mamba and Attention modules. B'MOJO-F uses a fixed computational budget to capture long-range dependencies far beyond the span of Attention thanks to the last rows in the state transition matrix. Specifically, the order N fixes the number of tokens (the most recent ones) that are processed by a local attention window ρ thanks to the upper diagonal identities in A B'MOJO (u t ). On the other hand, the last rows of A B'MOJO (u t ) aggregate information from the most recent past data, akin to "attention sinks" [50] .While B'MOJO-F generalizes both Transformers and Mamba models, it can only access information outside the current window with lossy fading memory. Data that becomes relevant only after a long time span would be ignored. We therefore modify our model class to detect and incorporate such tokens into the span of the local attention eidetically. We do so with a simple method that we call Innovation Selection: whenever B'MOJO processes a new token that cannot be explained using the lossy fading memory, we affix it to the eidetic memory. Innovation Selection operates similarly to the mechanism behind the Lempel-Ziv-Welch (LZW) compression algorithm [58, 47] . 7 Like with fading memory, the tokens in the new eidetic memory are processed by the same sliding window attention we use in B'MOJO-F. Algorithm 1 captures the steps we perform from input to output of each layer of B'MOJO architecture, which we discuss below.Data: Input data u t , inner recurrent state x t-1 , fading output y t-1 , eidetic memory M t-1 , window size w, B'MOJO's state update matrix A(u t ), input to state vector B(u t ), state to output vector C(u t ), a predictor function ŷ(•) over a span of length k.Result: Output out t u t-w:t ← u t-w-1:t-1 ∪ u t ;// Short-term memory windowx t ← A(u t )x t-1 + b(u t ) ; y t ← C(u t )x t ; // Long-term fading memory ϵ t ← error(ŷ(y t-1:t-k ), y t ) ; // Innovation computation M t ← M t-1 ∪ {u t , ϵ t } if ϵ t > min ϵ∈Mt-1 (ϵ) M t-1 otherwise ; // Long-term eidetic memory out t ← Attn(u t-w:t , y t-w , M t ); // Model outputIn Algorithm 1, the first line updates the short-term memory (collection of the last w seen tokens) by ejecting the oldest token and appending the new token. The second line updates the SSM fading memory, which stores new information into the state x t . Then, we augment fading memory with our eidetic memory M t , a set of tokens that the model decides to keep based on their unpredictability: tokens that are difficult to predict given the past, as measured by error(ŷ t (y t-1:t-k ), y t ), could be valuable far in the future when their memory has faded away, and are hence curated in the eidetic memory. The amount of information stored in M t can grow unbounded over time up to hardware limitations, in Section 3.3 we further discuss how to efficiently implement the predictor ŷ.Finally, we weave fading and eidetic memory together to predict a new token through a sliding attention mechanism, which aggregates relevant information from short-term memory u t-w:t , the fading memory y t and the long term eidetic memory M t .Differences between B'MOJO, B'MOJO-F and vanilla Hybrid models. Differently from vanilla hybrid models that stack SSM and Attention layers, B'MOJO-F and B'MOJO allows the Attention module to attend to both the input and output tokens of the SSM, allowing the Attention layer to merge information from fading "memory" tokens (output of the SSM) with "eidetic" tokens (the layers' inputs). Differently from B'MOJO, B'MOJO-F does not use the innovation selection mechanism and therefore it does not implement long term eidetic memory.We efficiently implement B'MOJO's tiered memory hierarchy at scale. The modelling choices that lead to B'MOJO closely follow ideas from Stochastic Realization and have an obvious recurrent efficient implementation. On the other hand, during training, one is more interested in a parallel formulation that can allow processing of multiple tokens at the same time and in a single forward pass. In the following, we describe how we use chunking to develop B'MOJO's parallel form.Efficient sliding window and memory chunking. Computing fading memory and selecting the most unpredictable tokens to store before feeding them into an attention layer is a sequential process which is hard to parallelize. To solve this we use chunks of length w and use fading and eidetic memory to summarize the whole past before the beginning of each chunk. Then, to aggregate information from the input and memory tokens, we use a sliding window over the interleaved concatenation of input chunks and their respective memory tokens as shown in Figure 7 (see Appendix B.2). Our modular and interleaved approach in Figure 7 enables us to leverage optimized kernels for both SSM [19] and sliding window attention [9] , enabling fast training and inference beyond the 1B scale. In Figure 4 we show that our efficient implementation is faster than both Mamba [19] and Mistral [24] .Efficient Innovation Selection. The Innovation Selection process we describe in Algorithm 1 requires the predictor ŷ(y t-1:t-k ). While building a new parametric predictor with learnable parameters is possible, in practice, this requires modifying the training loss. In our implementation, we consider a fixed predictor so no extra weights need to be learned. We fix ŷ as a running weighted average and implement it using short 1D grouped causal convolutions.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Justification: We prove our theoretical claims about transformers's eidetic memory and SSMs's fading memory constructively inline by showing the corresponding state matrix for each model in the main text of the paper.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 crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
experimental results: 
 To evaluate the scalability and efficiency of B'MOJO, we compare it with state-of-the-art model classes (SSMs, Transformers and hybrid variants) on both synthetic and language modeling tasks.Our experimental results are of two main types, (i) language modeling scaling laws (Section 4.2) with zero-shot evaluation on short context text benchmarks (Section 4.3), and (ii) specific long context/recall-based tasks where finite-context models are ill-fit. We wish to emphasize that in this setting, and for the results of type (i), Transformers are a paragon, not a baseline, since most tasks are answerable within the context. Therefore in the results of type (ii) we leverage specific benchmarks, like synthetic tasks (Section 4.1), long context evaluation and length generalization to assess our models (Section 4.3). The goal of our novel model class is to cover the entire spectrum, i.e. perform comparably to the paragon on finite contexts while preserving higher performance than Transformers whenever the data relevant to solve the inference tasks fall outside the context window.All experiments compare against three baselines: (1) Transformers, represented by a downscaled Mistral-7B [24] architecture re-trained from scratch (2) SSMs, represented by Mamba and (3) Hybrid models, implemented by stacking Mamba with a sliding window attention [10, 28] . For a fair apples-to-apples comparison all our models are trained from scratch using the same pre-training data, The plot also exhibits a non-saturating scaling law, showing that increasing the amount of resources leads to increasingly better B'MOJO models.tokenizer [24] , and context length. And, in order to abalte the contribution of eidetic memory, we consider B'MOJO-F in addition to B'MOJO in all our experiments.We use synthetic tasks to test B'MOJO's ability to recall exact information from beyond the attention span [2, 51] . We do so with Multi-Query Associative Recall (MQAR) data in the main text, and consider other tasks such as Induction Heads [36] , Selective Copying [19] , Fuzzy MQAR [39] , and Noisy MQAR [39] in Appendix D.1.B'MOJO's Memory Efficiency on Associative Recall Tasks. The MQAR task [2] has been shown to correlate well with language modeling performance [2, 39] . Compared to its peers (e.g. Induction Heads [36] ) MQAR is considerably more difficult and requires strong recall capabilities. In Figure 2 , we display accuracy on MQAR as we vary the size of the recurrent state for 2-layer instances of our models. Panels 1-3 consider varying numbers of key-value pairs to illustrate increasing complexity.Here, a Transformer (Mistral), with its sequence-length-sized KV cache serves as the paragon, always achieving a 100% accuracy. Should its window size be restricted not to include the KV pairs, its accuracy would drop to that of a random guess. Our results show that while every model class improves in recall accuracy as its size increases, B'MOJO and B'MOJO-F do so more reliably and faster than others. We explain these findings as follows. Mamba possesses only fading memory to propagate information to the future, and therefore has a limited recall capacity [2] for a given size of its recurrent state. Hybrid models leverage a sliding window attention to mitigate this issue, however they require a lengthy window to increase their span and recall the reference pairs. The strong performance of B'MOJO-F showcases the value of fading memory over a simple hybrid configuration, and the even stronger performance of B'MOJO, most evident in the panel on the right, highlights the added contribution of an eidetic memory for precise recall.We next demonstrate B'MOJO's favourable scaling laws on mid-size language modeling, baselining against the same set of state-of-the art-model classes as the previous section. We report the training setting and hyper-parameters in Appendix B.1.Results. In Figure 3 , we report perplexity at different scales against (left) the number of parameters and (right) the wall-clock training time. B'MOJO is faster than Mamba and Mistral that use efficient CUDA kernels (Flash attention and Selective Scan), outperforms our pre-trained Mamba baseline at all scales, and is comparable with our Mistral Transformer model, in Appendix D.2 we further comment our results in relation to other recent hybrid models like Griffin [10] and Zamba [18] .We additionally report the scaling behavior of B'MOJO-F as an ablation of the eidetic memory, as well as a vanilla hybrid architecture composed of interleaved SSM and attention layers as an ablation of both fading and eidetic memory. Despite using only short context sizes (2k) in Figure 3 , our results show that adding fading memory strictly improves pre-training perplexity over the baseline hybrid model at all scales. Performance is further improved when eidetic memory is added despite incurring a slightly higher training time. Our results suggest that B'MOJO exhibits a non saturating scaling law: increasing the amount of resources (parameters/FLOPs) leads to increasingly better models.We catalog the performance of our pre-trained models on an assortment of short and long context zero-shot evaluation tasks. While perplexity captures the models' ability to predict language, these evaluations characterize their generalization capabilities to unseen tasks. We use the EleutherAI LLM Harness [17] to conduct all evaluations.In Table 1 we report results on common-sense reasoning and questionanswering tasks that require processing both short [19] and medium-sized contexts [2, 35] . Since these language tasks do not require long range modeling we would not expect B'MOJO to meaningfully outperform our baselines. Moreover, B'MOJO uses a smaller sliding window (512 tokens) than Mistral (1024 tokens), placing the former on an uneven footing. Despite these caveats, we find that B'MOJO still bests Mamba at the 1.4B scale and performs comparably to the Mistral model.Long Context Evaluation. We next investigate the ability of B'MOJO to process long contexts using the PG-19 dataset [17] and more recall-intensive natural language tasks using the SWDE, Scrolls [2, 41] benchmarks. Our results show that B'MOJO outperforms Mamba and our hybrid baseline on PG-19 and SWDE, while B'MOJO-F outperforms Mamba on the Scrolls datasets, in line with our previous findings on the synthetic tasks in Figure 2 . These long context tasks showcase-in a practical setting-B'MOJO's efficacy in recalling information from beyond the attention span. We evaluate the ability of B'MOJO to improve its predictions with longer contexts than ones seen during training, an attribute termed length generalization [10, 51, 3] . Whereas length generalization in Transformers is limited by positional encodings and memory constraints [10] , for SSMs and B'MOJO, it is instead limited by the capacity of the recurrent state. In Figure 5 we report perplexity on PG-19 as the model processes contexts larger than pre-training contexts. We observe that B'MOJO 1.4B and Mamba 1.4B are capable of reducing and maintaining lower perplexity levels than Mistral at long context sizes. Curiously, we find that on models trained on longer sequences (8k), length generalization still holds and allows the model to continuously reduce perplexity as more tokens are processed, up to 4× the pre-training sequence length. Additionally, the pre-training perplexity of B'MOJO models trained on longer contexts is lower than identical ones trained on the same amount of tokens but on shorter contexts, showcasing that our model can properly leverage the eidetic and fading memory to process long sequences. In Appendix B.2, we report our implementation of backpropagation through time that we use to efficiently train our models on even longer contexts (beyond 16k). Our results on smaller scales (up to 370M) show that a model trained on 16k tokens can length generalize up to contexts of length 64k.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We use the supplement to add a plethora of details to reproduce our empirical findings.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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 be possible for other researchers to have some path to reproducing or verifying the results. 
conclusions and limitations: 
 Although our experiments have been conducted up to a 1.4B scale, scaling B'MOJO further requires non-trivial engineering and compute. Therefore, despite B'MOJO's promising scaling laws, it is difficult to ascertain whether it could scale to even larger models and datasets, and do so competitively. Scaling our work to even larger models could result in positive societal benefits such as ease of access to information. However, these models could be used also to spread misinformation, so novel algorithms and research is important to improve controllability and reduce hallucination [15] . Despite our promising results on length generalization, we observed that Mamba checkpoints trained on more compute tend not to length generalize that well. Since B'MOJO leverages a Mamba-like module to implement fading memory, we cannot exclude the possibility that it will be less effective in length generalization as we scale more. However, exploring simple time normalization techniques as mitigations is a promising area for future work [32] .Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [No] Justification: We do not provide code at submission time.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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/public/guides/CodeSubmissionPolicy) 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 a subset of experiments are reproducible, they 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) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We have a section titled Training Details in the appendix that covers precisely the hyperparameters, and type of optimizer and other crucial training details.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
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: We do not report error bars in our paper, but clarify the sample sizes of our evaluation sets and other choices of hyperparameters so that the reader recognizes the significance of our results.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, confidence intervals, or 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, or 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).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not 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 in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: We clearly state the compute hardware we used and the quantity we used in Training defaullts.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 of compute required for each of the individual experimental runs as well as 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 it into the paper). Guidelines:• The answer NA means that the authors have not 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., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss how if our elementary realizations were scaled to the size of foundation models there could be both positive and negative societal impacts in the limitations section.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Justification: The paper does not release new assets at this time, but will thoroughly document its assets when realeased at camera ready.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. 
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]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as 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. 
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: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Example A.1 (Biology). We note that biological agents have no option but to operate inductively, due to (a) hard memory bounds, and (b) evolutionary pressure towards minimizing inference latency: When faced with a threat, a biological agent is better served by a quick suboptimal decision than by reasoning over all past experience. AI built on silicon has no such limitations: Memory can grow unbounded and test-time computation can be distributed and improved by hardware design. Nonetheless, any practical realization involves some kind of constraint on inference time or compute resources. Therefore, resource-constrained optimal inference hinges on how to best use the available resources against a growing memory.Example A.2 (CNN Classifiers, VAEs and GANs). A trained representation can be co-opted to generate data. For example, a CNN can be used to classify random data until one is labeled with the desired class, and the resulting sample considered as being "generated" by the CNN. Similarly, one could generate random data indirectly by feeding noise to an encoder, as done in Generative Adversarial Networks (GANs), again co-opting a representation for generating data. In a Variational Autoencoder (VAE), data is generated by perturbing the latent representation of a map trained to re-construct the dataset.Example A.3 (Diffusion Models). Diffusion Models are representations, trained to re-construct the original data, but the mechanics used to reconstruct the data during training are sequential, using an artificial "time" variable, akin to a realization. This makes their use as "generative representation" natural since the reconstruction process is already a stochastic realization. 8Example A.4 (The Sage and the Savant). Picture the Library of Alexandria in 40 BCE, with the two best known experts, Sage and Savant. Sage had spent years reading the entire library and distilled its content down to maxims, proverbs, and various nuggets of wisdom. When asked a question, Sage would quickly return a pithy answer, although for the occasional unusual question, a generic answer. Savant was the librarian, with a preternatural ability to rapidly find content to assemble answers to any question. Savant did not have ready answers but, when asked, Savant would scour the library to retrieve relevant sources, speed-read through them, and assemble an answer. If asked the same question by the next customer, Savant would repeat the process anew, to the dismay of customers who saw Savant re-read the same material over and over, seemingly without understanding or learning anything. When voices spread that Savant produced more accurate answers, the enraged Sage burned down the library, putting Savant out of work. Sage regained the status of preeminent source of consultation, who could generalize wisdom from sources long gone, until two millennia later, when a new Savant was created from bits.Sage personifies inductive learning, favoring thoughtful and time-consuming learning to enable quick inference and rapid answers to questions. Savant represents transductive inference, which requires access to memory and efficient computation that is specific and tailored to the question, at the cost of having to repeat it all over.We train all the models from scratch on a common language dataset at scales from 130M to 1.4B using instances with 8 40GB NVIDIA A100 GPUs. For our 1.4B experiments, we use 8 instances simultaneously to train our model and perform the remaining experiments on a single instance. Models are trained using AdamW on 20x the number of tokens as parameters [21] . We use a batch size of 0.5M tokens and a cosine learning rate schedule with 5% warmup and a minimum learning rate of 1e-5. Finally, we use a weight decay of 0.1 and gradient clipping of 1 (see the Appendix for additional details). B'MOJO and B'MOJO-F do not use positional encodings, and both, along with Hybrid use a sliding window of 512 tokens. Furthermore, we found that the learning dynamics of our hybrid models can be improved by having two different learning rates parameters for the hidden SSM and the sliding window. We follow Mamba's learning rates [19] and GPT3's [6] for the sliding window.Sliding Window Attentionx t-K -∞Summary of the remote past Recent past u t-K+1 , . . . , u t x t-K -∞ SSM u -∞ , . . . , u t-K SSM Most informative remote past Summary of the remote past Recent past x t-K -∞ x t-K -∞ u tWe efficiently implement B'MOJO's tiered memory hierarchy at scale. The modelling choices that lead to B'MOJO closely follow ideas from Stochastic Realization and have an obvious recurrent efficient implementation. On the other hand, during training one is more interested in a parallel formulation that can allow to process multiple tokens at the same time in a single sequence. In the following we describe how we use chunking to develop B'MOJO's parallel form.Efficient sliding window. Computing fading memory and selecting the most unpredictable tokens to store before feeding them into an attention layer is a sequential process which is hard to parallelize.To solve this we use chunks of length w and use fading and eidetic memory to summarize the whole past before the beginning of each chunk. Then, to efficiently aggregate information from the input and memory tokens we use a sliding window over the interleaved concatenation of input chunks and their respective memory tokens as shown in Figure 7 . To make sure that every token is computed uniformly we pick the size of the sliding window to be K := w + m f + m e , where m f is the number of fading memory tokens, and m e is the number of eidetic tokens. Note that the number of eidetic tokens is not known a priori however, in our experiments, we fixed it to an upper-bound. The modular and interleaved approach Figure 7 enables us to leverage optimized kernels for both the SSM [19] and the sliding window attention [9] components enabling fast training and inference beyond the 1B scale. Furthermore, one can accelerate the Attention components in B'MOJO using standard techniques like GQA, which, for example, can make B'MOJO 1.4B with 2k context length up to 7% faster than a vanilla implementation.Efficient Innovation Selection. The Innovation Selection process we describe in Algorithm Line 1 requires a ŷ(y t-1:t-k ) whose task is to predict the next output y t of the SSM. Whenever, it is difficult to predict y t using samples from the past y t-1:t-k the consider the most recent input token u t highly informative and store it. While building a new parametric predictor with learnable parameters is possible, in practice, this will require to modify the training loss. In our experiments we consider a fixed predictor so no extra weights need to be learned. We fix ŷ as a running weighted average and implement it using short 1D grouped causal convolutions.BPTT: Back-propagation Through Time In this section we discuss our technique for backpropagation through time which enables training with arbitrary context length, potentially infinite. This enables us to treat the training data as batch of samples or as one large string. The trick to enable such training lies in splitting the data into multiple chunks, performing computations on each chunk, caching the statistics like the hidden states for SSMs or the tokens in the previous sliding window/eidetic memory, and then using them for the next chunk of data. One way to efficiently implement this is to write a custom CUDA kernel for this task which defines a Mamba layer which can process a non-zero initial hidden state (as the current implementations only support a zero initial hidden state). However, we would like to use the existing implementations and instead modify the inputs/layer weights such that we can load the cached hidden state before that start of every chunk. Let x t ∈ R b×d×h be the hidden state of an SSM layer (Mamba), where b, d, h are the batch-size, model hidden dimension, state hidden dimension (channels for each state dimension) respectively, and u t ∈ R b×d is the input to the layer. The state update equations in Mamba evolve asx t+1 = A(u t ) * x t +B(u t ) * u t , where A(u t ) ∈ R b×d×h , B(u t ) ∈ R b×d×h , u t is u t stacked h×,and * is an element-wise multiplication operation. In the traditional implementations x 0 = 0, however, we would like to use a non-zero initial state x ′ 0 . The trick we use here is to break x ′ 0 ∈ R b×d×h into h tokens of size b × d (which is the same as the size of u t ) and pass them as additional prompts to the input. In practice the dimension h is usually small (like 16) and as a result we can afford to break the hidden state into tokens and pass them as inputs. Let T be the size of each chunk (input), then to enable processing non-zero initial hidden states (or caching) we need to process chunks of size T + h where T ≫ h. We perform h initial steps of the SSM to load the previous hidden state, before we start processing the current chunk. To dynamics for the first h timesteps are governed by the following equation for t ∈ [0, h]: x t+1 = x t + B ′ t * x ′ 0,t , where x ′ 0,t chooses elements from the t column in the final dimension, B ′ t is one-hot with 1 in dimension t, and acts a shifting operation which ensures that the tokenized previous hidden state is correctly loaded. Thus after h time-steps the hidden state x h is loaded with the final hidden state from the previous chunk, which can be used with gradient accumulation techniques to compute gradients on arbitrarily sized input sequences. Note that B'MOJO has additional components like the eidetic memory, however, that can be efficiently cached from the previous chunk in constant time/memory cost.Lets start from the usual expression of the attention, 9 given an input sequence {u i } ∈ R n ch , we create the keys, values and queries vectors as follows,k i := {W K u i } ∈ R n , q i := {W Q u i } ∈ R n , v i := {W V u i } ∈ R n . Nowthe output of a causal attention layer is given by:EQUATIONC.1 Linear attention It is possible to approximate the exponential in the numerator using a kernel representation which makes the attention a linear operator (the ratio of two linear operators) onto the augmented feature space. The basic idea is to write exp(q T t k i ) ≈ ϕ(q t ) T ϕ(k i ), we therefore get that the linear attention can be computed as: 9 For the sake of simplicity we do not use Multi-Head Attention, however the generalization is straightforward. Now consider the numerator, the denominator can be obtained by fixing the input v i = 1 ∀i. It is trivial to show that we can represent the linear attention using a Finite-Impulse Response dynamical system as follows:y t ≈ y lin t = t i=1 ϕ(q t ) T ϕ(k i )v i t i=1 ϕ(q t ) T ϕ(k i ) (3)EQUATIONThis is a linear dynamical system that evolves over time and returns the final values of the linear attention at each time instant t.Remark: Note that we can modify the previous equations to represent the attention mechanism by simply using a non-linear read out function (the exp) in place of the identity. However, this is not very insightful since this simple FIR system is simply a shifting register over time.It is worth noticing that the state of this system is t, and it always increases over time, usually up to a design parameter specified by the system design which is dictated by the available compute and memory (typical values are set to 2k).A simple way to measure the state of the system is to write its state space realization, whose dimension directly informs us on the expressivity of the system. In our case we shall assume that the attention is computed on a sliding window of size K. A canonical realization of the FIR described above is:EQUATIONIn this section, we describe the connection with State Space models, in particular Mamba (input dependent state space model) and characterize how it approximates a linear attention mechanism.First we state the Mamba equations:   ūt = w 0 u t + w 1 u t-1 + w 2 u t-2 + w 3 u t-3 x t = a(ū t )x t-1 + b(ū t )ū t y t = c(ū t )x t + dū t (6)For the sake of simplicity we shall now study the input to state and the read out equation. It is easy to show that this is not a strictly causal realization of a dynamical system (since the state at time t is updated with the input at the same time).Remark: Every single channel in a Mamba block is independent from each other and unnormalized, Mamba reduces the variability across channels with a coupling in the lower dimensional projections of the gating parameter ∆.A natural way to prevent the system matrices to grow ever larger is to approximate the FIR above using an AR component, which would essentially keep a running average of the keys for each token seen so far. This will allow the model to keep some higher level statistics of past data which would be used to summarize past information into a single dimension of the dynamical system.EQUATIONwhich can be written as:EQUATIONand canonically realized as:EQUATIONWhen studying dynamical systems it comes particularly helpful to study canonical forms. They are particularly well suited to assess properties of dynamical systems and to realize state space models that realize a desired input-output behavior.In particular, given a transfer function of a LTI systemEQUATIONwe can write its state space canonical controllable realization asEQUATIONThe first simple result to show is that any causal (but non-strictly) dynamical system in the following formEQUATIONcan be rewritten in its canonical form asEQUATIONIn fact, starting from the transfer function of the first systemEQUATIONWe getEQUATIONNote that the terms C and B only appear as the product CB, which then we can rename as Ĉ.Remark: Extending the previous canonical from to Time-Varying Systems is tedious but straightforward.The controllable canonical form we introduced in Equation ( 12) can be easily extended to Time-Varying Systems and Input-Varying as well since it encodes the algebraic properties of the relationship between "positional" variables at time time instants (or of different inputs). Hence, it is straightforward to see that when setting all the coefficients on the last row of the state transition matrix to zero we get back the same nilpotent dynamical system that represents the attention mechanism (note that, differently from the linear attention case, the read-out function is non-linear).On the other hand, we can use this canonical form to represent a Mamba (diagonal) model (which is non-canonical) by simply picking the coefficients of the characteristic polynomial such that its poles are the same as the diagonal entries of the Mamba block. Note however, that Mamba being non-minimal, could have some zero-pole cancellations depending on the specific values of the input matrix B(u t ) and C(u t ), in such cases the input-output behaviour associated with the cancellation does not appear in the output of Mamba and, so long as it is stable (which is always the case thanks to Mamba's parametrization), it is guaranteed to remain bounded and decay to zero exponentially fast.Hence, both Mamba and the Attention mechanism can be implemented by B'MOJO for any fixed state N .In Table 3 we report results on synthetic tasks other than Multi-Query Associative Recall (MQAR). The table below expands the range of synthetic tasks from MQAR to four more synthetic tasks and compares performance across multiple scales.In Figure 3 , we report perplexity at different scales against (left) the number of parameters and (right) the wall-clock training time. B'MOJO is faster than Mamba and Mistral that use efficient CUDA kernels (Flash attention and Selective Scan), outperforms our pre-trained Mamba baseline at all scales, and is comparable with our Mistral Transformer model.While prior works like Jamba [28] (and the recently released Zamba [18] ) 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 [10] also use sliding windows, but manage to slightly outperform Transformer leveraging much longer sliding windows sizes than ours.Table 3 : BMOJO demonstrates strong performance on a wide range of synthetic tasks, and both small and medium scale. We compare the performance of various models at the 2 layer and the 130M scale on 4 different synthetic tasks, (1) Selective Copying, a task involving recall of a specific sequence of tokens with random spacing (2) Induction Heads, the recall of a specific token amongst noisy tokens (3) Noisy MQAR, an associative recall task retrieving keys in a noisy environment and (4) Fuzzy Recall, an associative recall task involving keys and values that are multiple tokens each. We find that B'MOJO models consistently outperform or match all existing baselines. In this section we derive the form of MAMBA reported above from the original source, combining the published paper and software implementation provided by the authors [19] .Every Mamba layer contains a State Space Model which maps sequences (B, L, d in ) to (B, L, d in ).Call the input of the state space models in Mamba is (B, L, d in ) is {u i } L i=0 . The input sequence is used to generate the discretization time-step ∆, the input matrix B and the output matrix C, using the following projection matrices: P ∆ down ∈ R ∆ down ×din , P ∆up ∈ R din×∆ down and P B ∈ R N ×din and P C ∈ R N ×din . Overall, we haveEQUATIONNow, given a state representation for all d in dimensions x t ∈ R din×N and the state update matrix A ∈ R din×N , We can now write the mamba update rule as:∆(u t ) = softplus(P ∆up P ∆ down u t ) ∈ R din x t+1 = exp(Diag(∆(u t ))A) * x t + Diag(∆(u t ))u t u T t P T B y t = x t P T C u t + Dwhere * denotes the element wise product. The elementwise product makes clear that Mamba's dynamics are diagonal. Revisiting the description of Mamba's diagonal dynamics in the main paper,A(u t ) = A Mamba (u t ) =    a 1 (u t ) . . . a N (u t )    , b(u t ) = b Mamba (u t ) =    b 1 (u t ) . . . b N (u t )   it is natural to see that the values of a i (u t ) are elements of the matrix exp(Diag(∆(u t ))A) = exp(Diag(softplus(P ∆up P ∆ down u t ))A), and b i (u t ) are elements of Diag(∆(u t ))u t u T t P T B . Our nonlinear readout function ρ is simply bilinear in x t and u t .With some approximations, we can further show that Mamba can be written as a bilinear system on an augmented input space. Using the first order Taylor Series, we havex t+1 = e Diag(∆t)A * x t + Diag(∆ t )u t u T t P B = (Diag(∆ t )A + 11 T ) * x t + Diag(∆ t )u t u T t P B = x t + Diag(∆ t )(A * x t + u t u T t P B ) = x t + Diag(log(1 + exp(P ∆ u t )))(A * x t + u t u T t P B ) = x t + Diag([P ∆ u t ] + )(A * x t + u t u T t P B )where the second line follows from the application of the first order Taylor Series, the third from rearranging terms while applying the definition of ∆ t , and the forth line follows from replacing notation for the soft hinge loss gating with [•] + . Now, we carry out some algebra as follows:x t+1 = x t + Diag([u t ] + )(A * x t + P † ∆ u t u T t P T † ∆ P B ) = x t + Diag([u t ] + )A * x t + Diag([u t ] + )P † ∆ u t u T t P B = x t + Diag([u t ] + )A * x t + (u T t ⊗ Diag([u t ] + ))vec(P † ∆ )vec(P B ) T (I ⊗ u T t )T where vec vectorizes a matrix into a column vector, and ⊗ is the Kronecker Product between two matrices. Now we can look at the vectorized evolution of h, denoting vec(x t ) as xt ∈ R dN :xt+1 = xt + vec(Diag([u t ] + )A) * xt + vec((u T t ⊗ Diag([u t ] + ))vec(P † ∆ )vec(P B ) T (I ⊗ u T t ) T ) = xt + (A T ⊙ I) P arameters [u t ] + F eatures * x t + ((I ⊗ u T t ) ⊗ (u T t ⊗ Diag([u t ] + ))) F eatures (vec(P B ) ⊗ vec(P † ∆ )) P arameters, where ⊙ is the Khatri-Rao product. The purpose of this highly messy derivation is simply to demonstrate that there exists a feature map of u t that we will call z t , and sparse parameter sets Ã, PB , PC such thatEQUATIONand the system can be described as a discrete bilinear system. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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 as it is clear that these goals are not attained by the paper.Sometimes ambiguously termed 'test-time computation'[7,45], which is confusing since inductive inference also involves test-time computation. Another term sometimes used is 'test-time training'[43,53], which is also confusing since 'training' and 'testing' are disjoint and complementary phases of inference from inductive learning. All these are different forms of transduction.We report more examples on the differences between representations/realizations and induction/transduction in Appendix A.Under the stationary assumption, one stochastic realization of an SSM is a convolution. This input-output map is equivalent to an SSM and standard transformations can be applied from one to the other.Diffusion Models, a special case, are treated as realizations during training but used as representations during inference since all but the final element of the denoised sequence are discarded.Note that Stochastic Realization Theory does not assume that the data is Gaussian, zero-mean and white, it only requires the innovation process to be so. That is the condition that the prediction residual (innovation) be unpredictable, regardless of whatever distribution the data is drawn from.In brief, LZW scans an input sequence to find the shortest sequence that is currently unknown, appends this sequence to a dictionary while returning indices for known sub-sequences. Similarly, Innovation Selection scans the input sequence to find the shortest sequence that has high prediction error, adds to the eidetic memory unpredictable inputs, while returning outputs that leverage known sub-sequences.It is curious that the reverse diffusion equation was first derived in the context of Stochastic RealizationTheory by[29] to argue that a stochastic realization is a model, distinct from the physical system it realizes: Time is reversible in the former, but not in the latter. 


Decision: The submission studies an interesting combination of recurrent architecture motifs (such as Mamba) with Transformers. Three out of four reviewers were satisfied with the authors’ responses and raised their scores. The sole negative reviewer did not participate in the discussion. However, it looks like their main concerns about not having ablation studies and discussion of acceleration methods (such as GQA) has been addressed by the authors. The ablations are already present in the paper, and I suggest that the authors make it clearer in the camera ready version. The authors should also include a more comprehensive discussion of related work on modern hybrid architectures (attention + SSM), see comments by reviewer MJ98.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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?
limitations: limitations: Please check the above.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: questions: See weaknesses section.
limitations: limitations: The authors discuss the limitation of their work in Section 5.
paper_summary: paper_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.
strength_weakness: strength_weakness: -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: 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.
limitations: limitations: The limitations seem to be fairly discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: Adequately discussed.

Comments: comment: 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.
Thank you for your extremely thoughtful feedback and remarks. We will continue to think hard about how we can refine our claims about long context inference.
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.
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 |
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?