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 2024 poster (2024)
Keywords: Sequence Models
Language Models
State Space Models
Hybrid Architectures
Primary area: deep_learning_architectures
TL;DR: We introduce a novel hybrid architecture (State Space Models + Attention)  that efficiently processes past information using both fading and eidetic memory.

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

B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory
Luca Zancato*
Arjun Seshadri
Yonatan Dukler
Aditya Golatkar
Yantao Shen
Benjamin Bowman
Matthew Trager
A Alssannrro Achille
Stefano Soatto
AWS AI Labs
Abstract
We describe a family of architectures to support transductive inference by allowing memory to grow 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 finite span ("context" in Transformers), or fading over infinite span (in State Space Models, SSMs). Recent hybrid architectures have combined eidetic and fading memory, but with limitations that do not allow the designer the learning process to seamlessly modulate the two, nor to extend the eidetic memory span. We leverage ideas from Stochastic Realization Theory to develop class of models called B B'MOJO to seamlessly combine eidetic and fading memory within an elementary composable module. be used to implement models that can access short-
The overall architecture 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 Jamba are special cases of B' MOJO and describe a basic implementation that can be stacked and scaled efficiently in hardware. We B' MOJO 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 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' 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.
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 map previously unseen datum 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 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 prediction criterion, that can be used as sequential decision or prediction maps, or generative
*Correspondence to: zancato@amazon.cco amazon.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

models for sequence data. Large language models (LLMs) and other predictive models that involve randommess in the sequential generation process are special cases of stochastic realizations [28].
Representations are the backbone of inference from inductive learning, or induction for short. In- duction 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 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), manifest in so-called Nout-of-distribution effects or "distribution shift." These lead to apparent paradoxes involving generalization and memorization (55).
Realizations, on the other hand, are the backbone of transductive inference and generative AI (GenAl). Transduction refers to the process of inferring particular properties of test data by processing, at inference time, all given (particular) training data [ The boundary between induction and transduction is blurry: Induction can be viewed as restricted form of transduction, where training data is accessible only through the learned weights. An over-parametrinn 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 unseen data, transductive inference seeks to maximize memorization and forgo generalization in favor of sample-specific inference computation P
representation, when overfit to training
Transduction 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 [1] Solomonoff-styy 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, 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
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
which is confusing since inductive inference *Sometimes ambiguously termed 'test-time computation' 7 44
also involves test- time computation. Another term sometimes used is 'test-time training' 42] 51]. 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 represenaationn/raaaizaiooi and induccion/transduution in Appendix A

---

be independent and identically distributed (IID) one can encode inference computation through fixed "stateless" map. or model, which is 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 mnite-imenniioall statistics, called states, that summarize all past data for the purpose of prediction [25 22 28). 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 [2]. To perform optimal inference, memory has to grow, and if the data generation mechanism has finite complexity at some point an efficient encoding past data into memory will stop growing, but it is not possible to know when [I] Therefore, 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 large portion of the growing volume of data accessible through the Internet and make it available for inference. Current AI systems are typically hard- constrained 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- terg-termmie memory updated and accessed sparingly and asynchronously. The former includes both eidetic (lossless) and fading (lossy) memory for efficient computation, while inference. Such architectures would
the latter is akin to an integrated form of 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"
11
Existing architectures such as Transformers and SSMs fall short of encompassing these criteria. Transformer-bases architectures use eidetic memory restricted to finite span, "context length" 6,223, while recent SSM-based architectures [T8] T6 49 use only fading memory in their state. In both cases, scaling requires allowing the context (and the key-value cache) recurrent state to grow unbounded. Recent work on hybrid combinations of Transformer and State Space layers [27 [38] [0] show promise in striking a balance between eidetic memory, fading memory and compute.
1.1 Contributions
In this work, we describe a class of models that encompasses both recent SSMs [1]. Transformers [23]. and hybrid architectures [27] [0] 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' D'MJOO's eidetic memory-a process we call Innovation Selection.
We demonstrate through synthetic tasks that B B'MOJO outperforms existing SSM and hybrid model architectures in transductive inference. Empirically, we show that our implementation is 15% faster than similarly-sian model tasks up to the 1.4B scale. Finally, we show that B' MOJO can operate effectively at inference time sequences far longer (tested up to 4x) 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.
Transformers and SSMs while achieving comparable perplexity on language
2 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 sequence is called 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 (30], which is to estimate the true model parameters from data. If the model is linear and driven by Gaussian noise, the Haaley-Hamilton theorem 24 implies that
'Under the stationary assumption, one stochastic realization of SSM is convolution. This onput map is equivalent to an SSM and standard transformations can be applied from one to the other. DDiffusion 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 discarded.
3

---

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 [21]. If not only the model class, but the model itself are known, then the problem further reduces to filtering or prediction, which for the Thear-Gaussian case can be implemented as a closed-form iterative update [25]. 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 mathematical representation of sequential data capable of predicting the next datum the sequence given previous ones. The Wiener Filter [46] was among the earliest to be deployed, superseded by the Kalman Filter (25). 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 (28). The general problem of stochastic realization is, given a (potentially infinite) sequence (-... U4-1, 4 = Ugge observed up to time t, infer (i.e. "learn") some parametric model, function Φ with parameters A, (ugge) such that the prediction ut+1 = &o (ug)) yields a residual c++1 = u++1 ut+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 umpredictable.
State. The state of a model is statistic that makes the future independent of the past. In particular, such a statistic (function of past data) E(uge) yields to the following conditional entropy equality H (u (((u))))))))))))))))(()) uuze)= ) = H (u++1 (((gg)))))) relative to the joint distribution of all past and present data [30 22]] Trivially, the data itself fits the definition of state with {(u<i) ust, but it grows unbounded over time. Instead, one seeks bounded complexity state given which all past data can be ignored with no loss of information. Building such state is the core goal of stochastic realization.
Elementary model classes: LTI, LTV and LIV. Any finite sequence can be realized by linear (LTI) system driven by IID Gaussian noise (28). It would therefore appear that this model class is sufficient for any practical purpose, and we need models no more expressive than linear
ones driven by white zero-mean Gaussian noise
Tz+1 Ax, + But
x A(u)) + B(u))ue
LTI
LIV
y/ = VC= + et.
8/1 = C )xr + 1t.
Given observations your, stochastic realization deals with the problem of inferring an equivalence class of model parameters A, B, C (43) and a state Te along with covariance matrix Pe of [w,, ug] that can be propagated deterministically to approximate trajectories produced by the underlying data generation mechanism, forr [47). 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 [22] and imput-varying 26 afford more efficient representation by considering A,, B,, C, known functions of time (Linear- Time Varying) or of their inpur (Linear-Inpu Varying, LIV). As we describe next, a special case of LIV model class where the dependency on the input is linear, resulting in so-called bilinear realization, is gathering considerable attention [18] I0 49
Modern realizations. Input-dependent (bilinear) models first gained popularity half a century ago starting with Brockett (5) when (26) showed that "every nonlinear system with controls entering linearly is locally almost bilinear" "; 11 developed complete realization theory for this class of systems including minimal realizations, and showed that they not only minimize the dimension 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 [8]. Jamba (27) and Griffin [10]: they refer to input-dependency as "selectivity" [18] 19 (36 and combine Attention mechanisms and other techniques to scale up to 52B parameters. Similarly, but at smaller scale, Block State Transformers [14,, GSS [33] 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.
"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.

---

Bldding Window Attention
Iew mehere
SCAA
Skidng Windoww Attention
Iroet takre
M4-K
Innovation
Selection
Musi informaiio renicl
Surnmny the
Rocart
Figure 1: B'MOJO's memory management. (Left) Illustration of the B'MOJO layer. (Right) B'MOJO's Realization. B' MOJO's fading memory is computed by a SSM that represents long- range dependencies through its state fixed-dimensional representation) which is later aggregated 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 is to predict the next token using the SSM's 's state. If a token difficult to predict, 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.
3 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 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, implement 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 34 (50] also grows as more tokens are processed, potentially sub-linearly if optimized off-line. We propose to access information arbitrarily far in the past 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 inference-time, we impose no constraint its time span.
3.1 Fading Attention
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 non-linear nil nil-potent system whose state grows as more samples are processed. Much like a Moving- with non-linear read-out, the attention state is required increase as its span enlarged. Then, we show that Mamba [18] has 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-FFadee Fadd-Fading (B MOJO-F), effectively realizing a non-linear ARMA model [30 52 [33].
(MA) system
Transformers Attention to map input to output sequences according 3/1 Σi expla Ag)); Σi x
where lt. k, U/ are the query, key and value vectors and are all computed directly from the input tokens 14. We write this equation nil-potent dynamical system with softmax read-out function
as follows:
E++1 = A(ug)a + B(ue);
V: p(ue, v..)
where
0 I
A(u)) AATT
E 2V2VV and B(uz) barT(Ue)
b(u))
[kg
with b(u,) = E 2V where 2V is the embedding dimension of the KV cache and N is the 1i
length of the Attention window. Transformer has only short-term eidetic memory that is deadbeat in N steps: what slides "out of context" is permanently removed.
5

---

// Long-term fading memory Mamba is complementary in that it only has fading memory 18 with decoupled (diagonal) dynamics. Specifically,
(164
(1(we)
A(u)) = AManba(u))
b(ue) = bManba (u+)
=
a (we)
by (ue)
where N is picked a-priori and the output is yr = C(ue)ar Du. 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; = 1) can be used as permanent memory but at an unbounded cost of encumbering the state and re-processing it at each time step. [8] emphasizes the use of "selecive state space," which is simply bilinear realization 12 However, crucial to its implementation are interleaved convolutional layers that retrieve some of the eidetic functionality lost with the diagonal dynamics. In Sect. F we derive in detail the form above from the description of the paper, forgoing the unnecessary continuous time narrative, and comment the model in relation the actual implementation.
B' S'MOJO-F bypasses the limitations of Transformer and Mamba layers by using a Controllable Canonical Form [24] to realize Equation (I) (sce appendix Cp. that is:
I
AB Agmovo(us) =
an (we) a2(u)) aN(u))
b(u,)
In Appendix C we show that B'MOJO-F a minimal realization and, when the state dimension fixed, it generalizes both Mamba and Attention modules. B' B"MOJO-F uses 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 P thanks to the upper diagonal identities in Ag (ut). ). On the other hand, the last rows of Agoo (u,) aggregate information from the most recent past data, akin to "attention sinks" [48).
3.2 B'MOJO's complete functional form
While B D'OJOO-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 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 new token that cannot be explained using the lossy fading memory, we aftix it to the eidetic memory. Innovation Selection operates similarly to the mechanism behind the Zempel-Ziv-Wellee (LZW) compression algorithm [56 45) Like with fading memory, the tokens in the new eidetic memory are processed by the same sliding window attention we use in MOJO-F. Algorithm [ captures the steps we perform from input to output of each layer of B' MOJO architecture, which we discuss below.
Algorithm 1: The MOJO Mechanism
Data: Input data ue inner recurrent state 2;-1, fading output yt-1, eidetic memory M,-1, window size w, B' MOJO's 's state update matrix A(ue), input to state vector B(u), state to output vector C(u.), a predictor function 5(-) over span of length k.
Result: Output out,
Ut wit ---w-11 1 U
// Short-term memory window
It f Alue) + b(u)) Yt + C(u))xe
// Innovation computation
M, < [M U
i > min (c)
// LLongterm eidetic memory // Model output
MM-1
otherwise
out, + Attn(u,-oa Mt- M,):
II 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 sequences..

---

In Algorithm| | the first line updates the short- tertt 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 xt. Then, we augment fading memory with our eidetic memory M,, 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 (Vt--ikk))))) Ve ) 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 Mc can grow unbounded over time up to hardware limitations, in Section 3.3 we further discuss how to efficiently implement the predictor 9.
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 ut-wet the fading memory Yt and the long term eidetic memory M,.
Differences between B' 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 merge information from fading "memory" tokens (output of the SSM) with "eidetic" tokens (the layers' inputs). Differently from B' MOJO, B' M'MOOO-F does not use the innovation selection mechanism and therefore it does not implement long term eidetic memory.
3.3 B B'MOJO's efficient implementation
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 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 sequential process which 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 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 ennables us to leverage optimized kernels for both SSM [88 and sliding window attention [9] enabling fast training and inference beyond the 1B scale. In Figure we show that our efficient implementation faster than both Mamba [8] and Mistral 23
Efficient Innovation Selection. The Innovation Selection process we describe in Algorithm requires the predictor (())-1:t<)) 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 running weighted average and implement it using short ID grouped causal convolutions.
Jrequirsserruuirres
4 Experimental Results
To evaluate the scalability and efficiency of B' MOMOO, we compare it with state-of-the-arr 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 (42) with zero-shot evaluation on short context text benchmarks (Section 4.3), and (ii) specific long context/recall-basee 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 D. long context evaluation and length generalization to assess models (Section 4.3) The goal of our novel model class is to cover the entire spectrum, i.c. perform comparably to the paragon 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 2 architecture re-trained from scratch (2) SSMs, represented by Mamba and (3) Hybrid
models, implemented by stacking Mamba with a sliding window attention 10, 27 For a fair apples-to-apples comparison all our models are trained from scratch using the same pre-training data,
7

---

Number of Key Value Palrs: 6
Number Say Value Pains 12
Number of Key Value Paing.
Adding Eddtk Memory
120
200
100
300
50
Harrba
Narba
Mamta
Matrel
Kibra
Mtarrl
hsbaid
Mybri
fazid
22
20
Tocets
#NOjo
awoj
5'00
Tokens
9.20
315 Model Dimension
128
125
120
Dimension
Modal Dimension
Nodal Dimension
Figure 2: (Panels 1-3) B'MOJO has high memory efficiency on Associative Recall Tasks (se- quence length is 256 and attention window 32). For various models, we plot accuracy on the Multi-Query Associative Recall (MQAR) task as a function of the model dimension (totaling the SSM state, eidetic memory and KV cache where applicable). The transformer paragon attains 100% accuracy because it operates on the full context. While all models benefit strongly from increased memory, B' MOJO and B' MOJO-F consistently achieve the best accuracies for given memory budget. Panels 1-3 report MQAR tasks of increasing difficulty, on which the performance gap between B' MOJO and other models increases, showcasing the value of cidetic memory. (Panel 4) Increases eidetic memory size corresponds to gains in recall. We ablate the effects of eidetic memory by growing the number of eidetic memory tokens in B B'MOJ. Each added token contributes to an increase in recall accuracy until performance is saturated.
22
22
Memba
Hanba
20
Mub
20
Matry
frand
18
18
Perplexity
Perplexity
16-
16
14-
14
12
12
10
10
130
370
790
1400
10a
105
Parameters (Millions)
Training Time (seconds)
Figure 3: B'MOJO language modeling scaling laws. We plot the perplexity reached by models at different scales against the number of parameters and the wall-clock training time. B' B'MOJO is faster than Mamba and Mistral at training time while achieving better perplexity than Mamba and comparable perplexity with Mistral. The plot also exhibits non-saturating scaling law, showing that increasing the amount of resources leads to increasingly better B MOJO models.
tokenizer 23
and context length. And, in order to abalte the contribution of eidetic memory, we consider B' B'MOJO-F in addition to B D'MOJO in all our experiments.
4.1 Synthetic Tasks
We synthetic tasks to test B'MOJO's ability to recall exact information from beyond the attention span 12. (4]. We do so with Multi-Query Associative Recall (MQAR) data in the main text, and consider other tasks such as Induction Heads (35), Selective Copying [8] Fuzzy MQAR 38). and
Noisy MQAR
38 in Appendix D.I
B' B'MOJO's Memory Efficiency Associative Recall Tasks. The MQAR task 2 has been shown to correlate well with language modeling performance 2 (38. Compared to its peers (e.g. Induction Heads (35) 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 illustrate increasing complexity. Here, a Transformer (Mistral), with its sequence-length-eee 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 random guess. Our results show that while every model class improves recall accuracy as its size increases, B B'MOJO and B' MOJO-- 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 given size of its recurrent state. Hybrid models leverage sliding window attention mitigate this issue,
8

---

Table 1: B'MOJO's performance on downstream tasks. We compare different architectures on several zero-shot downstream tasks used to test common- sense reasoning and question-answernn answering on relatively small contexts. These tasks, however, do not require strong recall capabilities because the input text is typically very short (results on longer contexts are reported in Table 2 On pre- training perplexity B' MOJO performs on par with our pre-trained Mistral model and outperforms our pre-trained Mamba models at the largest scale we test 1.4B. However, accuracy metrics, while MOJO still outperforms Mamba, its gap with the Mistral model increases.
Pre training Perplexity
Short Contest
LAMBADA 37
HellaSSog
X
PIQA
[
ARC-E E
ARC-C WinoGnmde
Avg.
Mistral (Pull-Anzation)
31.6
25.5
41.4
370M
Manha (SSM)
31.4
22.3
41.2
(Sliding
SSM)
22.4
39.3
(Fading)
29.5
23.0
518
40.7
BMido (Fading. Edetic)
2.67
28.6
63.9
44.3
22.1
50.7
40.5
Mistra (Rull. Ariemton)
$1.9
Mamba (SSM)
2.37
13.9
45.0
70.3
52.4
28.0
51.9
48.6
Hybrid (Sliding Attestion + SSM)
2.42
37.6
38.8
66.1
48.4
25.4
52.6
44.8
RMado (Fading)
2.27
45.4
46.0
70.0
52.3
26.6
53.3
48.9
BMolo (Fading Eidetic)
2.26
44.8
46.5
09.9
54.7
26.5
52.1
49.1
however they require a lengthy window to increase their span and recall the reference pairs. The strong performance of B'MOJO-F showeases 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.
4.2 Language Modeling Scaling laws
We next demonstrate B' MOJO's favourable scaling laws on mid-size language modeling, baselining against the same set state-of-the art-model classes as the previous section. We report the training setting and hyper- parameters in Appendix B.T.
we report perplexity at different scales against (left) the number of parameters and (right) the wall-clock training time. B ''MJOO is faster than Mamba and Mistral that use efficient CUDA kernels (Flash attention and Selective Scan), outperforms our pre-trained Mamba baseline all scales, and is comparable with Mistral Transformer model, in Appendix [D.2 we further comment our results in relation to other recent hybrid models like Griffin (TO).
Results. In Figure 3)
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 our results show that adding fading memory strictly improves pre-training perplexity over the baseline hybrid model at all scales. Performance further improved when eidetic memory is added despite incurring slightly higher training time. Our results suggest that B' B'MOJO exhibits a non saturating scaling law: increasing the amount of resources (paramcters/LLOP)) leads to increasingly better models.
4.3 Zero Shot Evaluation
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 EleutherAl LLM Harness to conduct all evaluations.
Short Context Evaluation. In Table I we report results common-sense reasoning and question- answering tasks that require processing both short [18 and medium-sizze contexts [4]. 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 B' MOJO still bests Mamba at the .4B scale and performs comparably to the Mistral model.
uneven footing. Despite these caveats, we find that
Long Context Evaluation. We next investigate the ability of B' B'MOJO to process long contexts using the PG-19 dataset [17] and more -caall-nttensie natural language tasks using the SWDE, Scrolls 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 Figure These long context tasks showcase-in a practical setting- B' MOJO's efficacy in recalling information from beyond the attention span.

---

Table Long range downstream tasks. We compare different architectures on several zero-shot long range recall-inte -itall-nttensvve downstream tasks [2]. Our B' MOJO variants outperform Mamba since they have stronger recall capabilities.
Archtaadare
Matill
Tong pp PG-19
[por]] Scnti-QAape ScoilaSarrQA
SWDE
Mtaral Ateaton)
02.16
(SSM)
Rileid/Stsling Miciinn SSNO
313
8.55
3.77
5.35
(Madd
3,59
15.56
11.39
Figure 4: Time in ms to process 2k se- quences. B'MOJO is faster than other ef- ficient implementations of Mamba [18] and Transformers [8] at all scales.
Miaral (Pal- (ic-io))
2733
512
14.00
12.20
17.37
Tlyeen (tiding Attaaion SSM)
1.06
IMdd (Paiing)
2.47
13.8
BMale (Fading Bidsid)
19
23
Mamba
Mamba
18-
B'MOJO-F
22
B'MOJO-F
Perplexity
27
B'MOJO
21
B'MOJO
Mistral
20
16
jerple
19
18
17
14
16
512
2048
8192
32768
512
2048
8192
32768
Evaluated Context Size
Evaluated Context Size
Figure 5: Length generalization. (Left) We pre-train B' MOJO 1.4B and Mamba 1.4B on 2k context lengths and a 1.4B Transformer baseline on 1k. (Right) We pre-train B'MOJO 790M and Mamba 790M on 8k context length and compare models evaluating perplexity on longer sequences up to 32K tokens on PG-19. Transformers cannot length generalize (a known failure mode), on the other hand B' B'MOJO preserves/improees in perplexity better than Mamba even longer sequences.
4.4 Length Generalization
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 49 3.. Whereas length generalization in Transformers is limited by positional encodings and memory constraints [10], for SSMs and MOJO, it is instead limited by the capacity of the recurrent state. In Figure 3 we report perplexity on PG-19 the model processes contexts larger than pre-training contexts. We observe that B MOJO 1.4B and Mamha 1.4B are capable of reducing and maintaining lower perplexity levels than Mistral at long context sizes. Curiously, we find that 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 4x 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.
5
Conclusions and Limitations
Although our experiments have been conducted up to a 1.4B scale, scaling B' BOMO further requires non-trivial engincering and compute. Therefore, despite B BOMO''s promising scaling laws, 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 ease of access to information. However, these models could be used also to spread misinformation, so novel algorithms and research is important improve controllability and reduce hallucination ([5). 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' D'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 [31]
10

---

References
[1] Alessandro Achille and Stefano Soatto. On the learnability physical concepts: Can a neural network understand what's real? arXiv preprint arXiv:2207.12186, 2022.
[2] Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the cccall-hrooghhut tradcoff. arXiv preprint arXiv:2442.18668, 2024.
[3] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517. 2024.
[4] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piga: Reasoning about phys- ical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432-7439, 2020.
[5] Brockett. On the algebraic structure of bilinear systems. theory and applica-tions of variable structure systems, 1972.
[6] Tom Brown, Benjamin Mann, Nick Ryder, Melamie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[7] Olivier Chapelle, Vladimir Vapnik, and Jason Weston. Transductive inference for estimating values of functions. Advances in Neural Information Processing Systems, 12, 1999.
[8] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint 1rXiv:1803.053 2018.
[9] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 16359, 2022.
[10] Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murary Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
[11] Paolo D' Alessaadro, Alberto Isidori, and Antonio Ruberti. Realization and structure theory of bilinear dynamical systems. SIAM Journal on Control, 12(3):517-535 1974.
[12] David LeRoy Elliott. Bilinear control systems: matrices in action, volume 169. Springer, 2009.
[13] B. Bowman et al. Inductive learning of transductive inference: ArXiv [forthcoming). 2021.
[14] Mahan Fathi, Jonathan Pilault, Pierre-Luc Bacon, Christopher Joseph Pal, Orhan Firat, and Ross Goroshin. Block-state transformers. In Neural Information Processing Systems, 2023.
[15] Alessandro Favero, Luca Zancato, Matthew Trager, Siddharth Choudhary, Pramuditha Perera, Alessandro Achille, Ashwin Swaminathan, and Stefano Soatto. Multi-modal hallucination control by visual information grounding. arXiv preprint arXiv:2403. 2024.
[16] Daniel Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. /4052, 2022.
[17] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h. Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for shot language model evaluation, 12 2023.

---

[18] Albert Gu and Tri Dao. Mamba: Linear-tine sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
[19] Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209, /2951, 2022.
[20] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[21] Alberto Isidori. Nonlinear control systems: an introduction. Springer, 1985.
[22] Andrew H Jazwinski. Stochastic processes and filtering theory. Courier Corporation, 2007.
[23] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulmier, et al. Mistral 7b. arXiv preprint arXiv:2310.06225 2023.
[24] Thomas Kailath. Linear systems, volume 156. Prentice-Hal Hall Englewood Cliffs, NJ, 1980.
[25] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. 1960.
[26] Arthur Krener. Bilinear and nonlinear realizations of input-output maps. SIAM Journal on Control, 13(4):827-834, 1975.
[27] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai transformer- mamaor language model. arXiv preprint arXiv:2403.19887, 2024.
et al. Jamba: A hybrid
[28] Anders Lindquist and Giorgio Picci. On the stochastic realization problem. SIAM Journal Control and Optimization, 17(3):365-389, 1979.
[29] Greta M I.jung and George EP Box. On measure of lack of fit in time series models. Biometrika, 65(2):297-303, 1978.
[30] Lennart Ljung. System identification toolbox: User's guide. Citeseer, 1995.
[31] Xuezhe Ma, Xiaomeng Yang. Wenhan Xiong, Beidi Chen, Lili Yu. Hao Zhang, Jonathan May. Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient I1m pretraining and inference with unlimited context length. arXiv preprint arXiv:240.008801 2024.
[32] Mireille Chaleyat Maurel and Dominique Michel. Results on the non existence of finite dimen- sional filters. Stochastics; An International Journal of Probability and Stochastic Processes,
13(--2):33-002 1984.
[33] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint 1oXXi:2206.13944 2022.
[34] Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305. 16300, 2023.
[35] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zae Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. https//transffrmmr--iit www.ciitssuu//222/inmmoooo wowtext-learning-aaeeee..-e induucioon and-inuctiinnnii heads/index.hne
[36] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference Machine Learning, pages 266698 PMLR, 2023.
12

---

[37] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint arXiv: 06031, 2016.
[38] Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844 2024.
[39] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge scale. Communications of the ACM, 64(9):99-106,
2021.
[40] Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 12007-12021, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics.
[41] RJ Solmonoff. A formal theory of inductive inference. Information and Control, 7:224-254, 1964.
[42] Yu Sun, Xiaolong Wang. Zhuang Liu, John Miller, Alexei Efros, and Moritz Hardt. Test-time training with self-supervision for generalization under distribution shifts. In International conference on machine learning, pages 9229-9248. PMLR, 2020.
[43] Peter Van Overschee and Bart De Moor. Subspace identification for linear systems: The- ory citlemeennationnAppicccttte
Springer Science & Business Media, 2012.
[44] Vladimir Vapnik. Transductive inference and semi-supervised learning. 2006.
[45] Terry A. Welch. A technique for high-performance data compression. Computer, 17(06):8-19, 1984.
[46] Norbert Wiener. Extrapolation, interpolation, and smoothing stationary time series: with engineering applications. The MIT press, 1949.
[47] William Wonham. On a matrix riccati equation of stochastic control. SIAM Journal on Control, 6(4):681-697, 1968.
[48] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2023.
[49] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:23/2.06635, 2023.
[50] Lili Yu, Dániel Simig, Colin Flaherty, Armen Aghajanyan, Luke Zettlemoyer, and Mike Lewis. Megabyte: Predicting million-byte sequences with multiscale transformers. Advances in Neural Information Processing Systems, 36, 2024.
[51] Luca Zancato, Alessandro Achille, Tian Yu Liu, Matthew Trager, Pramuditha Perera, and Stefano Soatto. Train/test-time adaptation with retrieval. 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 1 2023.
[52] Luca Zancato, Alessandro Achille, Giovanni Paolini, Alessandro Chiuso, and Stefano Soatto. STRIC: Stacked residuals of interpretable components for time series anomaly detection. 2021.
[53] Luca Zancato and Alessandro Chiuso. A novel deep neural network architecture for non-linear system identification. PAMC-PapesSOn..ine, System Identification SYSID 2021.
machine really finish your sentence? arXiv preprint 18-191, 2021. 19th IFAC Symposium on
[54] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can 2019.
13

---

[55] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning (still) requires rethinking generalization. Communications of the ACM, 64(3):107- 115, 2021.
[56] Jacob Ziv and Abraham Lempel. Compression of individual sequences via variable-rate coding. IEEE transactions on Information Theory, 24(5):530-55) 1978.
14

---

A Induction and Transduction
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 threat, a biological agent is better served by quick suboptimal decision than by reasoning over all past experience. Al 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-cone constraiie optimal inference hinges 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 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 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 realization. This makes their use as "generative representation" natural since the reconstruction process is already a stochastic realization
Example 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 pithy answer, although for the occasional unusual question, a generic answer. Savant was the librarian, with 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 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.
B B'MOJO implementation details
B.1
Training details
We train all the models from scratch on a common language dataset at scales from 130M to 1.4B using instances with 40GB NVIDIA A100 GPUs For our 1.4B experiments, we use 8 instances simultancously to train our model and perform the remaining experiments on a single instance. Models are trained using AdamW on 20x the number of tokens parameters (20). We use a batch size of 0.5M tokens and a cosine learning rate schedule with 5% warmup and a minimum learning rate of le-5. Finally, we use a weight decay of 0.1 and gradient clipping of (see the Appendix for additional details). B' MOJO and B' MOJO-F do not use positional encodings, and both, along with Hybrid use 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 [18 and GPT3's 6 for the sliding window.
*It is curious that the reverse diffusion equation was first derived in the context of Stochastic Realization Theory by (28) to argue that stochastic realization is a model, distinct from the physical system it realizes: Time is reversible in the former, but not in the latter.
15

---

f-
8BM
Bliding Window Attention
SSM
Slidng Window Attention
Hageeste My-K
Innovation Salection
Noe Nomintorstion www.s//eeeeeeee........
Summary the remote past Recent past
phal
Figure 6: B'MOJO's Memory management. (Left) Fading Memory B B'MOJO fading memory is computed by SSM that represents long-range dependencies through a fixed-dimensional representa- tion which is later aggregated on the current tokens along with with the most recent past. (Right) Eidetic + Fading Memory Fading memory handled as in the left panel while tokens from the past selected using innovation test over the SSM output and appended the current sliding window. The innovation test measures how difficult a new tokens to predict using the state of the SSM, if a tokens is difficult to predict from the state store it in the eidetic memory and pass it to the attention module.
Siding Window Attention
Heroy
Iunny
Skiding Window: Langth
Figure 7: (Left) B'MOJO's implementation. We report the basic layer we use to implement MOJO and its memory hierachy (fading and eidetic). (Right) Efficient interleaved implementa- tion. We show to efficiently implement sliding window attention over chunks of input, fading and eidetic tokens.
B.2 B'MOJO's efficient implementation
We efficiently implement B D'MJOO's ticred memory hierarchy at scale. The modelling choices that lead to B' 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 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' MOOO's parallel form.
Efficient sliding window, Computing fading memory and selecting the most unpredictable tokens 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 sliding window over the interleaved concatenation 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 K == w + m, + Me. where my the number of fading memory tokens, and m. is the number of eidetic tokens. Note that the number of eidetic tokens not known a priori however, our experiments, we fixed to an upper-bound. The modular and interleaved approach Figure enables us to leverage optimized kernels for both the SSM [18 and the sliding window attention
9 components enabling fast training and inference beyond IB scale. Furthermore, can accelerate the Attention components in B P'MIOO using standard
16

---

techniques like GQA, which, for example, can make B' B'MOJO 1.4B with 2k context length up to 7% faster than vanilla implementation.
Efficient Innovation Selection. The Innovation Selection process we describe in Algorithm Line T
requires a (( whose task is predict the next output y: of the SSM. Whenever, it is difficult to predict Ye using samples from the past 1t-1 the consider the most recent input token uy highly informative and store it. While building 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 be learned. We fix ÿ as a running weighted average and implement it using short ID grouped causal convolutions.
BPTT: prok-pprong Through Time In this section we discuss our technique for back- propagation through time which enables training with arbitrary context length, potentially infinite. This enables treat the training data batch of samples 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 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 24 Rxxdxh 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 E R is the input to the layer. The state update equations in Mamba evolve as ft-1 = A(ux) *xx + B(u,)* + Uh, where A(w)) E R B(ue) E RRxddxh, ût is ue stacked hx, and is an element- wise multiplication operation.
the traditional implementations To = 0. however, we would like to use non-zero initial state z/g. The trick we use here is to break z/g- E Rx into tokens of size x d (which is the same as the size of u,) and pass them as additional prompts to the input. In practice the dimension h is usually small (like 16) and
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 initial steps of the SSM to load the previous hidden state, before start processing the current chunk. To dynamics for the first i timesteps are governed by the following equation for E [0, hl: F2+1 = x + B₁ =/ where d',e chooses elements from the t column in the final dimension, B3 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 ht time-steps the hidden state In is loaded with the final hidden state from the previous chunk, which can be used with gradient accumulation techniques to compute gradients 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.
C State Space realization Attention
Lets start from the usual expression of the attention given an input sequence {u;} E Re we create the keys, values and queries vectors as follows, k; == [WK [WKue] u;) E R", 9: == (Wouk) E R".
2 = [Wyu} E R". Now the output of causal attention layer is given by:
El=1 exp(1 k.)
yt E exp(at k:)
(2)
C.1 Linear attention
It is possible to approximate the exponential in the numerator using 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(g] ki) x @(g:)" (kk)), we therefore get that the linear attention can be computed as:
Dinola (kk)) (()))
Vt Rd yi
E T (k:)
For the sake of simplicity do not Multi-Head Attention, however the generalization is straightforward.
17

---

Training Context Size=
31.5
Training Context Size=16k
Perplexity
30.5
29.5
28.5
2048
4096
8192
16384
32768 65536
Evaluated Context
Figure B'MOJO's long context training using BPTT Appendix B.2 We train B'MOJO with two different context sizes, 2k and 16k (BPTT) respectively and evaluate on long context task (PG-19). We show that our model trained with 2k context size is able to extrapolate for context size upto 65536 (with marginal increase in the perplexity), while model trained with 16k context size can handle long context much more effectively which can be seen by the lower perplexity values as the context size increases. Remarkably, previously observed in [10), B' B'MOJO models trained on long contexts underperform models trained on shorter ones if evaluated on fewer tokens (i.e., when the inference context size is much smaller than the training context size).
Now consider the numerator, the denominator can be obtained by fixing the input vs = Vi. is trivial to show that we can represent the linear attention using Finite-Impulse Response dynamical system as follows:
xt Ei ((k))
Vt = (g))x
This is linear dynamical system that evolves over time and returns the final values of the linear attention at each time instant I.
Note that can modify the previous equations to represent the attention mechanism by simply using a non-linear read out function (the exp) in place the identity. However, this is not very insightful since this simple FIR system is simply a shifting register over time.
Remark:
is worth noticing that the state of this system is t, and it always increases over time, usually up to 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 the expressivity of the system. our shall assume that the attention is computed on a sliding window of size K. A canonical realization of the FIR described above is:
0
0
K+1
T2l-K
K
0
Ze+1 ==
+
o(ke)v
gt+1
1
(5)
1+1
Z, = 1] Z, + (kk))
y!in (b()))
C.2 Connection with State Space Models
In this section, we describe the connection with State Space models, in particular Mamba (input dependent state space model) and characterize how it approximates linear attention mechanism.
18

---

First we state the Mamba equations:
üt Woll + w144-1 + W244-2 + W334-3
a( xz-1 b(ū,)
(6)
1/t c(ū,) )ar + dũ,
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 not strictly causal realization of a dynamical system (since the state time updated with the input at the same time).
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.
Remark:
C.2.1 Local Global factorization of the attention mechanism
natural way to prevent the system matrices to grow ever larger is approximate the FIR above using an AR component, which would essentially keep running average of the keys for each token seen far. This will allow the model to keep some higher level statisties of past data which would be used to summarize past information into a single dimension of the dynamical system.
t-K d(kizs)vrr E o(k.)us)
V/1 Σ )lu))" (o ((ki)o Aaa)) Σ
K
)(a)) (kkk () + Zz-k)
(8)
which can be written as:
It Tt-K + ExT o o(ke-i)ur- Ut (q:))'xx
(9)
and canonically realized as:
--- K+17
7+11= =
+
(kk))z
at+1
(10)
[+1
7, [1
Z, + (kr))o
1/: oa:)
C.3 Elements of Realization Theory
When studying dynamical systems it comes particularly helpful to study canonical forms. They are particularly well suited assess properties of dynamical systems and to realize state space models that realize a desired onput-output behavior.
In particular, given transfer function of LTI system
3n - 8iz + + 8n-1z" 12
W(=)
(11)
Og 01 ... an-122- + nn
we can write its state space canonical controllable realization as
X =
X +
Ub
(12)
-00 -01
1/1
Bo
Bn-1]
Xt + Du,
19

---

The first simple result to show is that any causal (but non non-strictly) dynamical system in the following form
Ft Ar 1 + Bu+
(13)
y/ = Czi + Dut
can be rewritten in its canonical form
i A + ut-1
(14)
y: C.. + Duu
In fact, starting from the transfer function of the first system
B
2B
AB
W(2)
B
(15)
A
2-
2 A
We get
xz41 Ad; + u1-1
fit+ = A&t + 4n-1 = CABI, + (CB + D)u
3t = A A..; + But
(16)
1/1 = Cxt + Du
Note 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- Vime-narying Systems is tedious but straightfor- ward.
The controllable canonical form we introduced in Equation 12 can be easily extended to Time- Varying Systems and Varying as well since it encodes the algebraic properties of the relationship between "positional" variables time time instants (or of different inputs). Hence, it straightforward to see that when setting all the coefficients the last of the state transition matrix to zero 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 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 (w,) and C(u)), ). in such cases the input-output behaviour associated with the cancellation does not appear in the output of Mamba and, so long as stable (which is always the case thanks to Mamba's '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.
D Further empirical results
D.1 Synthetic Tasks Beyond Associative Recall
In Table 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.
D.2 Scaling laws and zero-shot evaluations
In Figure 3 we report perplexity at different scales against (left) the number parameters and (right) the wall-clock training time. B' 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 Mistral Transformer model.
While prior works like Jamba [27] slightly outperform pure Transformer models they used full attention, and thus retain the quadratic dependence of Transformers. contrast, our work only uses a small 512 token sliding window attention produce a model with linear dependency on the sequence length with a constant KV cache size. Other works like Griffin [T0 also use sliding windows, but manage to slightly outperform Transformer leveraging much longer sliding windows sizes than ours.
20

---

Table BMOJO demonstrates strong performance on wide range of synthetic tasks, and both small and medium scale. We compare the performance of various models at the layer and the 130M scale on different synthetic tasks, (1) Selective Copying, task involving recall of a specific sequence of tokens with random spacing (2) Induction Heads, the recall of specific token amongst. noisy tokens (3) Noisy MQAR, an associative recall task retrieving keys in noisy environment and (4) Fuzzy Recall, an associative recall task involving keys and values that multiple tokens each. find that B' MOJO models consistently outperform or match all existing baselines.
Selective Copying
Induction Heads
Noisy MQAR
Fuzzy MQAR
Model
layers
130M
layers
130M
layers
130M
layers
130M
Full context
GPT2
93.56
97.28
100
99.92
99.97
49.17
98.79
Mistral
94.67
100
99.99
98.23
Pythia 160m
99.9
100
98.99
SSM
Mamba
94.42
99.74
100
7.6
99.99
88.04
60.04
Reduced context (smaller window)
Hybrid
93.82
97.69
7.64
99.99
99.98
90.47
98.56
B' D'MOJO-F
93.92
98.58
100
99.97
99.99
90.62
98.35
B'MOJO
94.04
99.85
99.94
99.99
99.99
90.38
96.87
E Strip MAMBA
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 [8].
Every Mamba layer contains State Space Model which maps sequences (B, L, din) (B, din). Call the input of the state space models in Mamba is (B, L, d..) is {u,)
The input sequence is used to generate the discretization time-step A, the input matrix B and the output matrix C, using the following projection matrices: Pa E 12 Adown Xdin, Paop Rdd.. Adown and PB E Nddi and Pc E Nxdi..
Overall, have
B(u) = PBu E Rdi
C(ue) Pou E Rdd..
(17)
A A(u,) softplus PlAn,Pannee Pa Ui) Rdin
Now, given a state representation for all din dimensions #a E X and the state update matrix A E IRd x. We can now write the mamba update rule
(uu) = softplus (PA., Pa down uc) E 1dd
F++1 exp(Diag( (()))
Diag( (oag(A(u)))unu PT
3t = xr PT.
where * denotes the element wise product. The elementwise product makes clear that Mamba's dynamics diagonal. Revisiting the description of Mamba's 's diagonal dynamics in the main paper,
01 (4))
bi(ue)
A(u)) AManba (u;)
b(u)) 1Mammb
N (u.)
by (ue)
it is natural to see that the values of az az(u,) are elements of the matrix exp A (u: (A)))A) = exp(Diaa(softtluu (softplus( PA.p PAgor 44)) A), and by(u,) are elements of Diag( ((4)))())))))))))) pT Our non- linear readout function is simply bilinear in I; and ue.
21

---

E.1
Bilinear Mamba
With some approximations, we can further show that Mamba can be written as a bilinear system on augmented input space. Using the first order Taylor Series, have
2++1 Ding( (A,)A 24 Diag( A.) u PB
= (Diag [4) A 11") 2; + Diag (A, PB
+ Diag( A ) A 22 + iiii PB)
21 + Diag( exp (A)
ue un PB)
21 + (iag([Paa)]))
ueu! PB)
where the second line follows from the application of the first order Taylor Series, the third from rearranging terms while applying the definition of 4 and the forth line follows from replacing notation for the soft hinge loss gating with [+]+. Now, we carry out some algebra as follows:
It+1 = It + Diag [) za + PA u PA PB)
4 A + Diag [u, ++ P un Ps = Iz + Diag([uu]+)))))
+ Diag [u,
Fz + (u, @ Diag([u: (oc()
where vec vectorizes 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:) (2) E. RdN;
ie41 81 vec(Diag((uu) * Rt vec((u) @ Diag((u)++))vec(]]]) PB) " )
-1 + OI [1+++
*õ, + ((I 2
u/ @ (ul @ (aag((u+++)) (vec (PB) @ vec(P
Parameters Features
Features
Parameters
where ① is the Khatri-Rao product. The purpose of this highly messy derivation is simply to demonstrate that there exists feature map of un that will call 2,, and sparse parameter sets A, Pn, Po such that
. dz 3t + 2, PB 1/+ 7. P. D,
(18)
and the system can be described as a discrete bilinear system.
Mamba Model
(B. L, vocab size)
Mamba Block
Output (B, L, n,a))
Head
Out Proj.
PMS Norm
(B. L, dny
(B, I., R.
5S0M
N
RMB Norm
ConviD (k-d)
(B. n..)
(B, L, d
(B, L, d.
Tazio
Proj.
Text tokens (B.L)
Input (B, L,, Pl.)
Figure 9: An illustration of the Mamba architecture and main block.
22

---

SSM
Selective Soon
[R. [R.A.A.)
850
Urting
N)
015.
(KAL))
SAL-A
Figure 10: An illustration of the Mamba SSM block and Selective Scan operation.
23

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Our abstract and introduction proposes a novel model class leveraging both eideic and fading memory for sequence modeling. In our paper we propose the B'MOJO model class, we motivate it through stochastic realization theory, and we conduct experi- ments demonstrating its main capabilities, all reflect what is described in the abstract and introduction.
Justification:
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 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 is clear that these goals are not attained by the paper.
2. Limitations
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 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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 faimess.

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, 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 impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
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 '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 cross- referenced.
All assumptions should be clearly stated or referenced 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 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.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We the supplement to add plethora of details to reproduce 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.
the contribution is 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 novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
25

---

While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g.. large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how construct the dataset).
(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. the case of closed-source models, it may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described 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 NeurlPS code and data submission guidelines (https://nips. cc/ htte://uideee///eSbbmeeeeee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurlPS code and data submission guidelines (https mmnpp/.cc/public/guudee///////// wwwp:iidde//oeeSummiiiieeeeee htt.ssummmiiiiooooiooiooooo 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 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 the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
26

---

We have section titled Training Details in the appendix that covers precisely the hyperparameters, and type of optimizer and other crucial training details.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: We do not report error bars in 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, confi- dence 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful 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 reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We clearly state the compute hardware used and the quantity used in Training defaulits.
Guidelines:
The answer NA means that the paper does not include experiments.
27

---

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 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 't make it into the paper).
9. Code Of Ethics
Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpge///hoorrrs..ccco/a///////p
Question:
Answer: [Yes]
Justification: Upon review of the NeurlPS Code of Ethics we conform.
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 deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
10. 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 foun- dation models there could be both positive and negative societal impacts the limitations section.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA 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 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 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 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 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
28

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly credit the creators and original owners of assets used in the paper by citing them according to their request. explicitly mention the license our codebase for the paper where the creators and owners have requested we do so.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
29

---

Justification:
The paper does not release new assets at this time, but will thoroughly document its assets when realeased camera ready.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. 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 tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
supplemental material is fine, but if the main contribu-
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.
15. 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 equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing 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 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.
30


Reviews:
[{"review_id": "j8KZ4lhPLT", "reviewer": "Reviewer_majr", "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.", "questions": "* For the hybrid model baseline, what is the architecture (activation, layer ordering, FFN size, and attention-Mamba ratio)?\n* 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?\n* How is the proposed model's performance on real-world retrieval tasks like phonebook lookup and needle-in-a-haystack?\n* What is the impact of window length for the sliding window attention?\n* How is the profiling of the \"KV Cache\" of the proposed method in terms of sequence length?", "limitations": "Please check the above.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "* This work provides analysis for both attention and Mamba and proposes \"B'MOJO-F\" by combining sliding window attention and Mamba.\n* In order to further improve the memorization/recall ability, the Innovation Selection is proposed to compensate for the lossy memory in fading memory and add important token information to the long-term eidetic memory. The Long-term eidetic memory can increase as the input sequence length increases but could be way more efficient than standard attention.\n* The resulting architecture shows better memory efficiency on Associative Recall tasks than Mamba, although it still underperforms Transformers.", "weaknesses": "* Lack of sufficient ablation study. This work claims that it leverages 4 kinds of memory: short-term eidetic memory 'in-context,' permanent structural memory 'in-weights,' fading memory 'in-state,' and long-term eidetic memory 'in-storage'. Specifically, how does the short-term eidetic memory impact the capacity for recall and language modeling? Additionally, in tables 1 & 2, for \"BMoJo (Fading + Eidetic)\", which eidetic memory is this referring to? It seems that adding this eidetic memory does not improve performance.\n* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns. \n\n**Lack of sufficient ablation study.** Throughout our work (see Table 1, 2, and Figures 2, 3, 4, and 5), we compare B’MOJO to BMOJO-F (B’MOJO without Long-Term Eidetic Memory) in order to ablate the contribution of Long-Term Eidetic Memory, we further compare B’MOJO-F to a vanilla Hybrid Baseline (B’MOJO-F without Fading Memory) in order to ablate the contribution of Fading Memory, and we ablate Hybrid Baseline with Mamba to ablate the contribution of Short Term Eidetic Memory. The contribution of permanent structural memory is trivially ablated by the fact that trained weights outperform untrained ones. This covers all needed ablations: Overall, we find that B’MOJO > B’MOJO-F > Hybrid baseline > Mamba, ablating the role of each individual component. \n\n**Impact of short-term eidetic memory.** Short term eidetic memory in-context refers to the most recent set of tokens (sliding window) processed by the model, analogous to the context of a standard transformer model. For typical language modeling it is impactful --- it is used by the model to process and recall information from the most recent past (akin to a n-gram model). However it has no impact on long range recall capabilities beyond the context length. In both Tables 1 and 2, we refer to long-term eidetic memory, not short-term.\n\n**Adding eidetic memory does not improve performance.** Our scaling law results, see Figure 3 show that the eidetic memory has a positive impact on B’MOJO, which decreases perplexity as we scale the model size. To further isolate the eidetic memory gains we followed reviewer’s suggestions and report in Table A1 results on NVIDIA’s RULER benchmark using our larger models: B’MOJO-F 1.4B (with only fading memory) and B’MOJO 1.4B (with fading + eidetic memory). We show that B’MOJO-F is strictly weaker and the gap is higher the longer the required context length. Complementing the perplexity results, this showcases that our innovation selection mechanism indeed helps the model to recall specific information from the past. Furthermore, this is in line with our synthetic experiments in Figure 2 albeit at larger scale. Note that in Figure 2 panel 4 we show that increasing the number of eidetic memory tokens produces an increase in recall performance prior to saturation. \n\n**The main advantage of the proposed architecture** Our proposed model’s main advantage over Transformers is that it simultaneously 1) introduces linear inference complexity with respect to context length and constant KV cache while 2) preserving Transformers’ performance without incurring in the quadratic complexity cost of Attention.  \n\n**Acceleration methods.** Both GQA and Kernel Fusion speed up Transformers by efficiently moving data in memory, but neither change the fundamental quadratic inference complexity of Transformers. Even for moderate context lengths these acceleration methods cannot close the linear vs quadratic gap between B’MOJO and Transformers. We make this point empirically in Table A2, where we report profiling results on both our Transformer baseline and B’MOJO with and w/o GQA at the 1.4B/3B scale for different context lengths (1k/2k/4k). Note GQA can be applied to all attention layers in B’MOJO for further acceleration. In all cases, B’MOJO is still faster than Transformers w/ and w/o GQA. GQA does help Transformers more than B’MOJO, and the relative speed gap for, e.g. at the 1.4B scale, with 2k context length case is reduced from ~10% to 7% . As expected, B’MOJO’s improves more as longer sequences are used due to the quadratic vs linear scaling of attention on the Transformer baseline. \n\n**Empirical results B’MOJO vs Mamba vs Transformer.** When performing apples-to-apples comparisons (i), our largest B’MOJO model (1.4B) outperforms Mamba 1.4B (see Table 1, Figures 3, 5 and the new experiments with the needle-in-a-haystack results in Table A1, in the global comment), a trend observed in our synthetic experiments too (see Figure 2). While prior works like Jamba (and the recently released Zamba) slightly outperform pure Transformer models they used full attention, and thus retain the quadratic dependence of Transformers. In contrast, our work only uses a small 512 token sliding window attention to produce a model with linear dependency on the sequence length with a constant KV cache size. Other works like Griffin and Samba (which has been published after the submission of this manuscript) also use sliding windows, but manage to slightly outperform Transformer leveraging much longer sliding windows sizes.\n\n**Impact of the window length and profiling of the KV Cache.** The window length measures short-term eidetic memory and is similar to the impact of changing the window size on Transformer models like Mistral, the larger the better the results but at a higher FLOPs count, all our experiments leverage a fixed window size of 512. \nB'MOJO's forward time is constant with respect to the sequence length, and its KV cache is comprised of the sliding window cache, the fading memory and eidetic memory tokens (see Figure 7 for more details), whose number is fixed by the user. \n\n\n### Table A2: Profiling time (in ms) of a forward call for various model sizes, batch size = 1, measured on A100 40Gb\n\n|fwd profiling time in ms||Context Length|| B'MOJO's rel. improvement||\n|---|---|---|---|---|---|\n|Models|1024|2048|4096|||\n|---|---|---|---|---|---|\n|**1.4B**||||||\n|Transformer w/o GQA|72|117|oom|||\n|Transformer|67|110|oom|||\n|BMOJO-F|56|106|207|29%|10%|\n|BMOJO-F with GQA|53|99|190|26%|11%|\n|BMOJO|62|107|229|16%|9%|\n|BMOJO with GQA|58|103|224|16%|7%|\n|---|---|---|---|---|---|\n|**3B**||||||\n|Transformer w/o GQA|83|173|oom|||\n|Transformer |76|161|oom|||\n|BMOJO-F|78|157||6%|10%|\n|BMOJO-F with GQA|73|142||4%|13%|\n|BMOJO|81|163|311|2%|6%|\n|BMOJO with GQA|75|148|286|1%|8%|", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bnsX8nkxMj", "reviewer": "Reviewer_Wgie", "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.", "questions": "See weaknesses section.", "limitations": "The authors discuss the limitation of their work in Section 5.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-written and generally pleasant to read.\n* The idea of combining eidetic and fading memory, while not new, is very interesting and B'MOJO represents an original and elegant way to do so.\n* B'MOJO is shown to integrate the benefits of transformers and modern state-space models. The model compares favourably with Mamba on most tasks. While inferior compared to Mistral7B on zero-shot tasks, B'MOJO inherits its length generalization performance from SSMs, outperforming Mistral7B in this latter case. The model is shown to be marginally better than baselines in terms of training time.", "weaknesses": "* Generally Mistral 7B seems to perform comparably or even better than the proposed model, with the exception of the length generalization task. In particular, Mistral is often significantly better on the zero-shot evaluation tasks, also on those involving long contexts. Do the authors have an explanation for this? \n* 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?\n* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns. \n\n**Mistral 7B seems to perform comparably or even better.** As we clarify in general comments, in a true apples-to-apples comparison in short contexts a Transformer represents a paragon, not a baseline. The primary strength of the B’MOJO model relative to Transformers is in the long context regime. However, performance in long context tasks is sensitive to scale, especially for hybrid state space models like B’MOJO, and in Table 2 we provided experiments only at a scale that is reproducible in academic environments. However, in order to address the reservations raised in this review, we also ran larger-scale experiments up to 1.4B scale (Table A1) featuring NVIDIA’s RULER benchmark using our larger models: B’MOJO-F 1.4B (with only fading memory) and B’MOJO 1.4B (with fading + eidetic memory), where we show that both B’MOJO and B’MOJO-F outperforms Transformer models on long context tasks as expected.\n\n**Eidetic memory can in principle be unbounded, is there a criterion to select a specific subset?** In the proposed implementation we bound the number of tokens to a maximum size (to keep the cost of attention manageable). However, the token span can be arbitrarily large. We are yet to explore a selection mechanism (e.g. Landmark attention).\n\n**Clearly state the dimension of each matrix/vector.** We apologize for the confusion and have revised our draft to include dimensions of each matrix and vector.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kJFq2jQdw6", "reviewer": "Reviewer_PLgm", "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.  \n\n---\n\nThis 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.  \n\nnotes from reading the paper: \n  -New recurrent architecture called BMojo.  \n  -Adds eidetic and fading memory, aiming to outperform Mamba.  \n  -Substantially improved OOD length generalization.  \n  -Transductive inference for sample-specific inference.  \n  -\"Unpredictable\" tokens are added to a sliding window that is attended over.  \n  -Mamba has only fading memory.", "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.  \n  \n  -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": "The limitations seem to be fairly discussed.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "-The introduction and exposition are both well written.  \n  -The improved OOD generalization results are impressive.  \n  -The idea of adding the most surprising tokens to be attended seems like an interesting idea.", "weaknesses": "-While reading this paper, I got confused about the structure of the block, and the overall architecture.  This could be my fault, but it's possible that other readers will also get confused.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and your suggestions. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns. \n\n**Confused about the structure of the block, and the overall architecture.** We have made revisions to improve the clarity of our work. Please see Figure 6 and 7 in the appendix for more details.\n\n**B'Mojo algorithm block (algorithm 1) refers to a single BMojo layer.** Yes.\n\n**Odd that every layer in your architecture (even the first layer) is predicting the final output tokens.** Each block is only predicting representations that are passed to the subsequent block, and not predicting the final tokens. Only the final block predicts the final tokens.\n\n**How are the spaces, x, y defined in a more formal sense.** The spaces x and y are indeed \\in \\mathbb{R}^d, where d is the embedding dimension of the Mamba block.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9qLEYCe4L5", "reviewer": "Reviewer_MJ98", "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.", "questions": "Many of my major questions are listed in the weaknesses. Here are a few others.\n\nMajor: \n- 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?\n- 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.\n- 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?\n\nOther:\n- 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": "Adequately discussed.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Considers the important and relevant problem of balancing memory and efficiency in modern sequence models\n- Provides a nice, high level overview of stochastic realization theory and its connections with modern sequence models\n- The idea of using an error function inspired by the idea of an innovation process to determine a storage/eviction policy in sliding window attention appears to be novel", "weaknesses": "- There are several related lines of work that are not discussed well enough and/or should be cited:\n    - 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\n    - 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\n    - 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 \n  - 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. \n\n- 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. \n  - 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.\n  - I have other questions below which may help me better appreciate the value of this framework.\n- The experimental results seem to generally be weak\n  - MQAR: \n     - Too many details are missing to judge the usefulness of this experiment. E.g. What task sequence length vs sliding window lengths were used?  \n    -  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)?\n  - Language task: \n     - What dataset was used? I do not see this information anywhere.    \n     - 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. \n    - 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.\n    - 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and your insightful comments. We've addressed the main concerns shared with other reviewers in the global comment, here we will address specific concerns.\n\n**Related work References.** We are grateful for the references and associated insights, we have revised our manuscript to incorporate and discuss these works. We acknowledge that there are many relevant related works, our goal is to not propose yet another one, but rather a general framework that encompasses them using Stochastic Realization Theory.  Indeed, the referenced hybrid models are special cases of B’MOJO, for example Block State Transformers can be obtained by removing B’MOJO’s the eidetic memory. Similarly, GSS can be obtained by removing the eidetic memory and some of the sliding window attention layers. We have incorporated more discussions in our related work section. \n\n**Unique Insights Gained From Our Framework.** 1) Long term Eidetic Memory is unique to our work and provides crucial performance gains over models without it (see Figure 3 for perplexity gains at scale with eidetic memory, and Table A1 with results on NVIDIA’s RULER benchmark). 2) Our use of interleaved memory chunking, efficient innovation selection and input dependent SSMs (like Mamba) as opposed to slow long convolutions (like the ones used in Block-State Transformers) are specific innovations to ensure our method runs efficiently on modern hardware (see Section 3.3 and Appendix B2). These are important differentiators which we hope could inspire further developments in hybrid SSM architecture design. \n\n**Difference between B'MOJO/B’MOJO-F and Hybrid baseline.** Succinctly, a vanilla hybrid layer can be written as “output = Attention(SSM(input_tokens))”, while our B’MOJO-F layers can be expressed as “output = Attention(Cat(Inputs, SSM(past inputs)))”. In particular, B’MOJO-F is not a simple stacking of SSM (Mamba) and Attention as in our baseline hybrid model and some recent hybrid models in the literature (like Jamba/Griffin). While vanilla hybrid models are limited to interleaved stacking of SSM and Attention layers, B’MOJO-F allows the attention module to attend to both the input tokens (to the SSM) and the output representations of the SSM (the fading memory). A mechanism that allow the Attention layer to merge information from \"memory tokens\" to the layers' input tokens. \nDifferently from B’MOJO, B’MOJO-F does not use the innovation selection mechanism (no long term eidetic memory). Please see Figure 6 in the appendix for more details. \n\n**Motivation of the Innovation Selection Mechanism.** The \"state\" of a stochastic realization is a sufficient statistic for (next token) prediction, meaning that the state is as good as the entire past for the purpose of minimizing the prediction error (i.e. making the prediction error unpredictable). Thus, the predictability of the prediction error, which is what the innovation test measures, is the natural test to know if the state has captured the past well enough. In practice, whenever the innovation test measures high residuals we know that the state must be refreshed with new information which in turn will help the predictor in reducing future prediction residuals.\n\n**Innovation Selection's efficient implementation.** To have an efficient implementation we picked the squared loss and coupled it with a linear predictor which takes as input the most recent past states. However, having a learnable linear predictor would increase the number of parameters of the model and it would require to modify the training loss. To avoid this, we restrict the predictor to a moving average which we implement with short 1D grouped convolutions. To preserve efficiency, we leverage the same 1D convolution implementation used in the Mamba layers which uses a kernel size of length 4. \n\n**Experimental Results seem to be generally weak**.\nWhen performing apples-to-apples comparisons, our largest B’MOJO model (1.4B) outperforms Mamba 1.4B in pre-training perplexity by 10% and by 1% on zero shot downstream evaluations  (see Table 1, Figures 3, 5). Furthermore, as suggested by the reviewers we assessed the long context recall capabilities of B'MOJO on a more difficult needle-in-a-haystack benchmark (see the global comment). In these recall intensive tasks we found that B'MOJO significantly outperforms Mamba by 4% accuracy @2k context sizes and by 11% @4k. As well as outperforming our Transformer model on longer contexts sizes. \n\n**Synthetic MQAR experiments.**\nWe use a sequence length of 256 and an attention window length of 32 so that key-value pairs stay outside the sliding window context. We tried to strike a balance between detail and readability. We plan to make our work reproducible not just by listing numbers, but by releasing our source code upon completion of the review process. As we show in Figure 2, Panel 4, adding eidetic memory tokens (the difference between B’MOJO and B’MOJO-F) steadily increases performance. \n\n\n**Why is B'MOJO faster than Mamba and the Transformer baseline?** \n B’MOJO is faster than the Transformer baseline because it replaces MLP layers with Mamba layers. In Table A3, we report the profiling results for different sequence lengths and show that MLP layers are slower than Mamba layers as the sequence length increases. \nFurthermore, B’MOJO is faster than the Mamba baseline since it replaces half of the Mamba layers with sliding window attention layers of length <1k. In Table A3 we report the profiling results for different sequence lengths (a similar observation can be found in Figure 8 of the original Mamba paper). \n\n### Table A3: Profiling time (in ms) of a forward call of basic blocks, measured on A100 40Gb\n\n| Context Length | 1024 | 2048 | 4096 | 8192 | 16384 |\n|----------------|------|------|------|------|-------|\n| Blocks | | | | | |\n| Mamba | 1.1 | 2.0 | 3.2 | 6.0 | 11.9 |\n| Full Attention | 1.0 | 1.9 | 5.8 | 18.8 | 94.6 |\n| MLP | 1.2 | 2.2 | 3.9 | 7.4 | 34.6 |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their constructive feedback and the positive comments on the novelty of our method (MJ98, Plgm, Wgie) as well as its theoretical motivation and connection with Stochastic Realization Theory (MJ98, Wgie). Furthermore, we are glad that reviewers appreciated our efficient implementation which allows us to process longer contexts “more efficiently than standard attention” (majr) and our “impressive OOD generalization results” (Plgm) outperforming the transformer baseline in length generalization (wgie).

**Main evaluation benchmarks and paragons.** Reviewers MJ98, majr and Wgie are concerned that our Transformer baseline seems to perform comparably or better than our proposed model, with the exception of the length generalization task. Our experimental results are of two main types, (i) a generic comparison on short context benchmarks typically used to assess LLMs’ zero-shot performance on short contexts, and (ii) specific long context/recall-based tasks where finite-context models are ill-fit. For a fair apples-to-apples comparison all our models are trained using the same pre-training data, tokenizer, and context length. 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, long context evaluation and length generation to assess our models. 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. While the latter tasks may be a minority by frequency, they carry outsize weight in terms of business value (e.g. to support more factual queries and reduce hallucinations, RAG, …). Unlike the paragon, our baselines Mamba and a Hybrid model are surpassed by B’MOJO in both (i) and (ii). 

**Better assessment of long context recall capabilities.** Reviewers MJ98 and majr suggested we further assess the longer context recall capabilities of B’MOJO on a larger and more difficult needle-in-a-haystack benchmark. Following their suggestions, we report our results of our larger models (1.4B) in Table A1 below using NVIDIA’s RULER benchmark. Our goal is to test a) how much adding eidetic memory helps complementing fading memory, b) how B’MOJO compares with a Transformer and c) how it compares with a Mamba model on longer contexts.
We test our pre-trained models on 2k tokens at varying context lengths (2k and 4k) and compare our models with Transformer and Mamba baselines pre-trained from scratch on the same data. We find:

a) B'MOJO (with fading + long term eidetic memory) is strictly stronger than BMOJO-F (with only fading memory), and the gap is higher with longer context lengths (51% vs 48% accuracy @2k and 19% vs 10% accuracy @4k respectively). This showcases that our innovation selection mechanism helps the model to recall information from the past (in line with our synthetic experiments in Figure 1). 

b) The Transformer baseline only has high recall when it is tested on the same context length used during pre-training (2k) and struggles with longer sequences (as we show in Figure 5 on the length generalization experiments). On the other hand, B’MOJO can recall information beyond the pre-training context length, outperforming the Transformer baseline on 4k context length. 

c) Both B’MOJO and B’MOJO-F outperform the Mamba baseline and the gap increases as the context length increases, showcasing that Mamba does not preserve enough information from the older past and this is especially evident when the context length increases (accuracy @4k is 8% for Mamba and 19% for B'MOJO). 



### Table A1: Long context evaluation with RULER (needle in a haystack)
| Context Length | Model        | S-NIHA | MK-NIAH | MV-NIAH | MQ-NIAH | Average |
|----------------|--------------|--------|---------|---------|---------|---------|
| 2048           | Transformer  | 100    | 95      | 62      | 61      | 79      |
|                | Mamba        | 100    | 32      | 29      | 28      | 47      |
|                | B'MOJO-F      | 90     | 36      | 35      | 31      | 48      |
|                | B'MOJO        | 90     | 45      | 37      | 33      | 51      |
|----------------|--------------|--------|---------|---------|---------|---------|
| 4096           | Transformer  | 0      | 0       | 0       | 0       | 0       |
|                | Mamba        | 9      | 12      | 5       | 7       | 8       |
|                | B'MOJO-F      | 10     | 16      | 5       | 8       | 10       |
|                | B'MOJO        | 22     | 21      | 17      | 17      | 19      |

Meta-review: 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.
Decision: Accept (poster)