Paper deep dive
MixFormer: Linear Transformer with Mixture of Memory Experts
Yu Guo, Lei Duan
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:State Space Models (SSMs), as a mainstream research direction of linear Transformers, aim to achieve higher efficiency than standard Transformers in long-context modeling. However, existing SSMs suffer from limited input adaptivity and constrained memory capacity, leading to information loss when modeling ultra-long sequences. To address these limitations, we propose MixFormer, a novel linear Transformer that integrates a Mixture-of-Memory-Experts (MoE) mechanism. Specifically, the model maintains differentiated memory states through multiple collaborating memory experts and employs a novel Time-Aware Linear Attention (TALA) mechanism, which leverages learnable exponential decay functions and positional biases to dynamically update memory. This design enables the model to selectively reinforce important historical information while effectively mitigating memory dilution, substantially improving long-range dependency modeling. Experiments on long-sequence text and image generation tasks demonstrate that MixFormer not only achieves significant performance gains but also provides a more sustainable computational backbone for the next generation of web infrastructure.
Tags
Links
- Source: https://arxiv.org/abs/2608.09468v1
- Canonical: https://arxiv.org/abs/2608.09468v1
Trouble viewing inline? Open PDF directly →
Full Text
56,374 characters extracted from source content.
Expand or collapse full text
MixFormer: Linear Transformer with Mixture of Memory Experts Yu Guo School of Computer Science, Sichuan University Chengdu, China gyguoyugy@gmail.com &Lei Duan School of Artificial Intelligence, Sichuan University Chengdu, China leiduan@scu.edu.cn Corresponding author. Abstract State Space Models (SSMs), as a mainstream research direction of linear Transformers, aim to achieve higher efficiency than standard Transformers in long-context modeling. However, existing SSMs suffer from limited input adaptivity and constrained memory capacity, leading to information loss when modeling ultra-long sequences. To address these limitations, we propose MixFormer, a novel linear Transformer that integrates a Mixture-of-Memory-Experts (MoE) mechanism. Specifically, the model maintains differentiated memory states through multiple collaborating memory experts and employs a novel Time-Aware Linear Attention (TALA) mechanism, which leverages learnable exponential decay functions and positional biases to dynamically update memory. This design enables the model to selectively reinforce important historical information while effectively mitigating memory dilution, substantially improving long-range dependency modeling. Experiments on long-sequence text and image generation tasks demonstrate that MixFormer not only achieves significant performance gains but also provides a more sustainable computational backbone for the next generation of web infrastructure. Keywords Transformer ⋅· Large Language Model ⋅· Linear Attention ⋅· State Space Model ⋅· Mixture of Experts 1 Introduction Transformer [37] has become the foundational architecture for large language models (LLMs) [35, 36] and multimodal models [7, 6, 4], having been successfully applied in fields such as Natural Language Processing (NLP) and Computer Vision (CV). The powerful representation learning capability of Transformers stems from the self-attention mechanism, which leverages the softmax function to capture full pairwise token interactions. This results in quadratic time and space complexity (O(L2)O(L^2)) with respect to the context length, making it unsuitable for very long-range sequence modeling. In the World Wide Web (W), the Transformer architecture has deeply penetrated a wide range of core applications — from semantic understanding in search engines and personalized generation in recommendation systems to real-time interactive processing in web browsers. Its self-attention–based parallel computation offers highly efficient sequence modeling capabilities for massive web-scale data. Recently, efficient Transformer models [34, 42, 24, 13] have attracted increasing research interest, with a primary focus on reducing the computational cost of the self-attention mechanism. Among these efforts, linear Transformers have emerged as one of the mainstream solutions for building efficient Transformers. The core module of linear Transformers lies in constructing linear attention mechanisms that approximate the behavior of the softmax function through carefully designed kernel functions, thereby reducing the computational cost from O(N2)O(N^2) to O(N)O(N), where N is the context length. In addition, other linear Transformers further extend linear attention into a recursive formulation, allowing each token at a given time step to inherit the computation results from the previous step. Such methods are typically categorized as State Space Models (SSMs) [9, 11, 10, 14, 31]. Unlike the self-attention mechanism, which requires computing an explicit attention score matrix, SSMs do not maintain the correspondence between each query and historical tokens when handling variable-length contexts. Instead, an SSM preserves a fixed-size intermediate feature map as a global historical memory state matrix, which encodes and compresses information from all previous time steps while discarding their explicit temporal indices. Specifically, during the recursive process at each time step, the SSM’s linear attention updates the memory state matrix through additive operations. This update strategy discards the positional indexing of the original queries, making it impossible to retrieve information from specific time steps. Existing approaches primarily enhance the memory capability of SSM to capture historical information by introducing RNN-style gating mechanisms. These methods employ either fixed decay coefficients (e.g., the constant parameter proposed by [32]) or learnable weight matrices (e.g., the work of [40]), enabling the model to retain past states to some extent. However, although such methods can marginally increase the capacity of the SSM’s memory state, they still suffer from the following inherent limitations: (1) Existing approaches typically update the memory state using fixed coefficients or a limited number of learnable parameters, resulting in insufficient input adaptivity. The model should dynamically decide, at each time step, which information to preserve over the long term and which to gradually decay based on its contextual importance. (2) Most current SSMs employ a single memory state matrix to store historical information. Under ultra-long context inputs, this single-state architecture inherently limits storage capacity, making it difficult to prevent early critical information from being overwritten or diluted by subsequent inputs, thereby leading to irreversible information loss in long-range dependency modeling. To bridge these gaps, we propose a novel linear Transformer with Mixture of Memory Experts (MixFormer). The proposed architecture is built upon the linear Transformer, supporting both parallel training and recurrent inference modes. Specifically, MixFormer introduces the Mixture-of-Experts (MoE) into state memory modeling for the first time, where multiple memory experts collaboratively construct the final memory state through weighted fusion. Within each expert, we design a novel Time-Aware Linear Attention (TALA) to replace the conventional multi-head attention layer. By integrating exponential decay with learnable parameters, TALA dynamically adjusts the memory update strategy according to the current input, enabling the model to reinforce important historical information while progressively attenuating less relevant content. Furthermore, the multi-expert parallel architecture breaks the capacity limitation of a single memory state by forming a distributed memory storage. Each expert employs a differentiated memory strategy via the TALA mechanism, focusing on distinct temporal patterns such as long-range dependencies or short-term dynamics. Through a content-adaptive expert activation mechanism, MixFormer selectively strengthens the preservation of salient information and suppresses redundant signals. After multi-layer stacking, the model achieves a robust capability for ultra-long context modeling. We conduct extensive experiments on both image generation (including conditional and unconditional generation) and long-sequence text modeling tasks. The results demonstrate that the proposed MixFormer consistently achieves significant performance gains across all tasks. Furthermore, the effectiveness of MixFormer is validated through an in-depth analysis of its causal masking mechanism and the collaborative division of multiple memory experts. Finally, with the growing demand for web applications that support long-text understanding and multimodal content processing, MixFormer is expected to provide a more sustainable computational backbone for the next generation of web infrastructure. The main contributions are as follows: • We propose a novel linear Transformer with mixture-of-memory-experts (MixFormer). To the best of our knowledge, this is the first work to incorporate multi-memory-experts into State Space Models (SSMs). Maintaining multiple differentiated memory states in parallel enhances global perception over ultra-long sequences. • We propose a Time-Aware Linear Attention (TALA) mechanism based on linear attention, which integrates an exponential decay function with learnable pair-wise position biases to significantly enhance the model’s capability for long-range dependency modeling. • Experimental results on both text and image tasks demonstrate that MixFormer achieves strong performance while supporting efficient parallel training and fast autoregressive inference, providing a promising solution for building efficient Transformer models. Figure 1: Linear Attention vs SSMs. (a) is the workflow of the Linear Attention, and (b) is the State Space Models. Figure 2: (a) is the overall architecture of our proposed MixFormer, (b) is the Time-Aware Linear Attention (TALA) layer, and (c) is the Mixture of Memory Experts (MoME) layer. 2 Preliminaries Figure 3: Dot-product Attention vs Linear Attention. (a) is the illustration of the Dot-product Attention, and (b) is the Linear Attention. To better understand the working details of our proposed MixFormer, we provide a detailed review of various forms of dot-product and linear attention. 2.1 Dot-product Attention Figure 3 (a) illustrates the dot-product attention mechanism in the Vanilla Transformer [37]. Given the input X, it is mapped into three different semantic space matrices: Q, K, and V, where Q,K,V∈RL×dQ,K,V∈ R^L× d, L denotes the context length, and d is the hidden state dimension. Next, the attention weight matrix between Q and K is computed using the softmax function. Finally, this matrix is used to compute the weighted sum with V, the final result, as follows: Q Q =XWq, =XW_q, (1) K K =XWk, =XW_k, V V =XWv =XW_v Attention(Q,K,V) Attention(Q,K,V) =Softmax(QKTdk)V, =Softmax( QK^T d_k)V, where WqW_q, WkW_k, and WvW_v are the trainable matrix, dkd_k is hidden layer dimensional number. Generally, we refer to the attention mechanism as Dot-Product Attention or Softmax Attention. Its time and space complexity are both quadratic with respect to the sequence length (L). 2.2 Generalization of Attention According to Equation (1), the attention vector for the i-th token is calculated as follows: Attention(Q,K,V)i=∑j=1neqiTkj∑j=1neqiTkjvj,Attention(Q,K,V)_i= _j=1^n e^q_i^Tk_j _j=1^ne^q_i^Tk_jv_j, (2) where q, k, and v are column vectors. The eqi⊤kje^q_i k_j is used as a similarity function between two vectors. Let the similarity function be denoted as sim(qi,kj)sim(q_i,k_j). This generalizes Equation (2) to Equation (3), following: Attention(Q,K,V)i=∑j=1nsim(qi,kj)∑j=1nsim(qi,kj)vj,Attention(Q,K,V)_i= _j=1^n sim(q_i,k_j) _j=1^nsim(q_i,k_j)v_j, (3) where sim(qi,kj)>=0sim(q_i,k_j)>=0, following previous work [15, 30], the similarity function must ensure the non-negativity of the attention map. 2.3 Linear Attention The dot-product attention uses the softmax function to compute the similarity between Q and K, resulting in a time and space complexity of O(N2)O(N^2). Therefore, reducing the similarity function’s time and space complexity is key to achieving an efficient transformer. Figure 3 (b) illustrates linear attention, which uses a carefully designed similarity function, reducing the time and space complexity from O(N2)O(N^2) to O(N)O(N). Many previous works further extend sim(⋅,⋅)sim(·,·) to kernel-based methods [28, 12, 21], which approximate the softmax operation. The formula is as follows: sim(qi,kj)=ϕ(qi)Tϕ(kj),sim(q_i,k_j)=φ(q_i)^Tφ(k_j), (4) where ϕ(⋅)φ(·) can be seen as kernel smoother. By substituting the kernel similarity function into Equation (3), we obtain: Attention(Q,K,V)i=∑j=1nϕ(qi)Tϕ(kj)∑j=1nϕ(qi)Tϕ(kj)vjAttention(Q,K,V)_i= _j=1^n φ(q_i)^Tφ(k_j) _j=1^nφ(q_i)^Tφ(k_j)v_j (5) Then, by exploiting the associativity of matrix multiplication, the computation order of q, k, and v is altered. It can be simplified to: Attention(Q,K,V)i=ϕ(qi)T∑j=1n(ϕ(kj)vj)ϕ(qi)T∑j=1nϕ(kj)Attention(Q,K,V)_i= φ(q_i)^T _j=1^n(φ(k_j)v_j)φ(q_i)^T _j=1^nφ(k_j) (6) The intermediate results are reduced from an L×L× L matrix to a d×d× d matrix, and the overall computational complexity becomes (d2L)O(d^2L), which scales linearly with the sequence length L. 2.4 Autoregressive Generation In the inference process of Linear Transformers, the prediction of the current token depends solely on the previously generated token sequence and is not influenced by any subsequent tokens. Following the autoregressive generation formulation of language models in [15], we reformulate Equation (3) into an autoregressive computation paradigm: Attention(Q,K,V)i=ϕ(qi)TSiϕ(qi)TZi,Attention(Q,K,V)_i= φ(q_i)^TS_iφ(q_i)^TZ_i, (7) where SiS_i and ZiZ_i are two newly defined variables, as follows: Si=∑j=1i(ϕ(kj)vj),Zi=∑j=1iϕ(kj) S_i= _j=1^i(φ(k_j)v_j),Z_i= _j=1^iφ(k_j) (8) Si=Si−1+ϕ(kj)vjT S_i=S_i-1+φ(k_j)v_j^T Zi=Zi−1+ϕ(kj), Z_i=Z_i-1+φ(k_j), where S0=0S_0=0 and Z0=0Z_0=0. Therefore, the linear attention mechanism enables efficient autoregressive inference with both linear time complexity and constant memory usage. 3 Related Work 3.1 Linear Transformers Linear Transformers represent one of the main directions in the development of efficient Transformers. As shown in Figure 1 (a), many studies aim to approximate dot-product attention through carefully designed kernel function variants, thereby changing the computation order of QKVQKV and reducing the cost complexity from O(L2)O(L^2) to O(L)O(L). Such designs enable linear Transformers to achieve performance comparable to that of vanilla Transformers [37]. Li et al.[17] proposed a Taylor approximation-based linear attention mechanism. This method approximates the softmax function using the first-order Taylor expansion of the exe^x term in the softmax function. Katharopoulos et al. [15] proposed replacing kernel methods with non-negative activation functions, achieving up to 4000× faster autoregressive prediction on very long sequences. Shen et al. [30] proposed a method that ingeniously applies the softmax function separately to the rows of Q and the columns of K. Peng et al. [26] proposed random feature methods that offer a straightforward way of learning with recency bias through an optional gating mechanism to approximate the softmax function. Choromanski et al. [2] introduced Performers, which employ a novel fast attention via positive orthogonal random characteristics (FAVOR) approach to go beyond the limitations of softmax-based attention. Han et al. [12] proposed a novel focusing function that makes similar vectors more concentrated and dissimilar vectors more dispersed. Meng et al. [21] proposed POLAFormer, which improves the focusing function by making its focusing capability learnable, thereby further enhancing model performance. These efforts have advanced the development of linear Transformers. However, kernel-based approximation methods often yield attention maps that fail to adequately focus on important features. On the other hand, these methods fail to account for the recursive formulation of linear attention, thereby overlooking the update mechanism of the state memory matrix during inference. In contrast, our proposed MixFormer integrates both the parallel and recurrent forms of linear attention, while further enhancing the dynamic temporal awareness of the state memory matrix. 3.2 State Space Models Another line of research reformulates the computation paradigm of linear Transformers into an iterative process, enabling parallel modeling while maintaining autoregressive inference. As shown in Figure 1 (b), these methods draw inspiration from the autoregressive nature of recurrent neural nets (RNNs) and can be regarded as RNN-style Transformers. Such architectures are typically categorized as State Space Models (SSMs) [9, 8, 5]. Existing studies have proposed various algorithmic improvements to SSMs, aiming to optimize the update mechanism of their historical information feature matrices. Han et al. [32] proposed the Retentive Network (RetNet), which uses a fixed parameter γ∈[0,1]γ∈[0,1] as a computational factor to update historical information in the SSM. Han et al. [32] proposed the Retentive Network (RetNet), which introduces a fixed parameter γ∈[0,1]γ∈[0,1] as a computational factor to update historical information in the SSM. RetNet demonstrates the advantage of SSMs over kernel-based linear Transformers, enabling both Parallel training and sequential inference. Meanwhile, it employs chunkwise Parallelism to model ultra-long sequences, where inputs are processed in Parallel within chunks and sequentially across chunks. Yang et al. [40] further extended this idea with Gated Linear Attention (GLA), where a learnable weight matrix (W) is employed to control the update of historical information at each time step. GLA uses an input-dependent forgetting gate and dynamically adjusts the update weights of the sequence. Others of the SSM focus on dynamically updating the memory state at each time step through an exponential decay mechanism. Zhai et al. [41] proposed a learned pair-wise positional bias based on exponential decay, enabling SSMs to achieve dynamic temporal awareness. Peng et al. [25] introduced RWKV, which leverages token-shifted weighted exponential decay to further optimize the memory modeling mechanism. As the model depth increases, RWKV empowers SSMs with long-range temporal dependency modeling capabilities. Gu and Dao [8] proposed the Mamba model, which designs the parameters of the SSM as functions of the input, thereby enabling dynamic control of information flow. Subsequently, Dao and Gu [5] introduced Mamba-2, an improved version of Mamba based on the State Space Duality (SSD) framework, which enhances inference efficiency. However, merely enhancing the representational capacity of the memory state in SSMs is insufficient to address the problem of knowledge forgetting in long sequences. Our proposed MixFormer introduces a Mixture-of-Memory-Experts mechanism, deploying multiple functionally specialized memory experts to capture different types and scales of temporal knowledge. This multi-expert collaborative architecture significantly improves the model’s ability to retain information over long-range dependencies, effectively mitigating progressive knowledge decay in long-sequence tasks. 4 Methodology 4.1 Problem Definition As shown in Figure 2 (a), based on the Transformer architecture [37], our proposed MixFormer adopts a similar multi-block stacking structure. Each MixFormer block consists of a Time-Aware Linear Attention (TALA) and a Feed-Forward Network (FFN) layer, each followed by a Group Norm (GN) layer. Given an input sequence X=(x1,…,xn)X=(x_1,…,x_n), where n denotes the context length, MixFormer employs parallel encoding during training and autoregressive generation during inference. The input vectors X are first encoded through an embedding layer into the initial representations X0=(x10,…,xn0)∈Rn×dX^0=(x_1^0,…,x_n^0)∈ R^n× d, where d denotes the hidden dimension. These representations are then iteratively refined through L stacked MixFormer blocks as follows: Xl=MixFormer_Blockl(Xl−1),l∈[1,L]X^l=MixFormer\_Block_l(X^l-1), l∈[1,L] (9) where l denotes the index of the current block layer, and L is the total number of layers in the model. Each block progressively enhances the contextual representations. MixFormer_Block(⋅)MixFormer\_Block(·) denotes the mapping process of the current input within this module. MixFormer adopts a pre-norm strategy to enhance the training stability during deep stacking, ensuring effective gradient flow across multiple layers. Meanwhile, Group Norm (GN) [38] is introduced to address potential statistical biases arising from chunkwise computation, further improving the model’s training stability and convergence performance. The TALA module incorporates time-aware linear attention to effectively capture long-range temporal dependencies among tokens, while supporting both parallel training and recursive inference modes. The notation used throughout the methodology is summarized in Table 1. 4.2 Symbol Definition Table 1: Illustration of mathematical symbols. Symbol Description X Model inputs XlX^l Input to the l-th MixFormer block Q Query projection matrix K Key projection matrix V Value projection matrix WQ,WK,WVW_Q,W_K,W_V Learnable parameter matrices ω Exponentially decayed weighting G Learnable pair-wise positional bias SnS_n Intermediate memory state matrix O Model output B Chunk length 4.3 Time-Aware Linear Attention Figure 4: The recurrence of the Time-Aware Linear Attention. To enhance the memory state perception of long sequences in the state space model, we propose a novel Time-Aware Linear Attention (TALA) layer to replace the standard multi-head self-attention layer. Built upon the principles of linear attention and SSM design, this layer supports both parallel training and sequential inference. The parallel computation structure is illustrated in Figure 2 (b), while the recursive computation process is shown in Figure 4. In the TALA layer, given an input sequence X, it is first projected into three vector spaces of the same dimension, as follows: Q=XWQ,K=XWK,V=XWV, Q=XW_Q,K=XW_K,V=XW_V, (10) where WQ∈Rd×dW_Q∈R^d× d, WK∈Rd×dW_K∈R^d× d, and WV∈Rd×dW_V∈R^d× d are learnable parameter matrices used to generate the query, key, and value vectors, respectively. d is the hidden dimension. As shown in Section 2, linear attention changes the computation order of QKVQKV and requires the non-negativity of Q and K. Following previous studies [1, 19, 39], TALA applies a kernel function transformation to ensure the non-negativity of Q and K, which is defined as follows: ϕ(Q)=eQ‖Q‖,ϕ(K)=eK‖K‖, φ(Q)=e Q||Q||,φ(K)=e K||K||, (11) where ||⋅||||·|| denotes the L2L_2 normalization applied to the row vectors of Q and the column vectors of K. Meanwhile, we design a hybrid mechanism that integrates positional awareness with dimension-wise adaptability to enhance the temporal awareness of key–value (KV) pairs. This mechanism applies exponentially decayed weighting based on relative positions to emphasize recent information, and employs a learnable parameter matrix to achieve adaptive calibration across feature dimensions. While ensuring the non-negativity of K, it significantly improves the model’s ability to capture temporal dependencies in long sequences. The detailed computation process is as follows: TALAtt(Qt,Kt,Vt)=ϕ(Qt)⊙∑i=1tϕ(Ki)(ωGi)t−iVi∑i=1tϕ(Ki)(ωGi)t−i, TALAtt(Q_t,K_t,V_t)=φ(Q_t) Σ^t_i=1φ(K_i)(ω G_i)^t-iV_iΣ^t_i=1φ(K_i)(ω G_i)^t-i, (12) where ⊙ denotes element-wise multiplication, ω represents exponentially decayed weighting, and G∈RT∗TG∈ R^T*T is a learnable pair-wise positional bias. TALAtt(⋅)TALAtt(·) denotes the time-aware linear attention. In addition, TALA can be extended to a multi-head attention mechanism to enhance its feature extraction capability. By splitting the hidden dimension d into multiple head dimensions dhead_head and assigning independent decay parameters ω to each attention head, the model can capture temporal dynamics across multiple scales. Based on this, Equation LABEL:TALA-eq can be extended to the following multi-head time-aware linear attention formulation: ω=1−e−h,h∈[1,⋯,d/dhead] ω=1-e^-h,h∈[1,·s,d/d_head] (13) headi=TALAtt(Qt,Kt,Vt,ωi), head_i=TALAtt(Q_t,K_t,V_t, _i), Y=GroupNormh(Concat(head1,⋯,headh)), Y=GroupNorm_h(Concat(head_1,·s,head_h)), where d is hidden dimension, dhead_head is head dimension. The GroupNorm is applied to the output of each head. The Parallel Representation of TALA. During training, MixFormer can leverage parallel matrix operations. The parallel computation architecture of TALA is illustrated in Figure 2 (b), and its computations are defined as follows: Q=ϕ(XWQ),K=ϕ(XWK),V=XWV, Q=φ(XW_Q),K=φ(XW_K),V=XW_V, (14) TALAtt(Q,K,V)=Q⊙(KT⊙D)V, TALAtt(Q,K,V)=Q (K^T D)V, D=ωn−mGnm,n≥m0,n<m, D= \ array[]lω^n-mG_nm,n≥ m\\ 0,\ n<m,\\ array . where WQ∈Rd×dW_Q∈ R^d× d, WK∈Rd×dW_K∈ R^d× d, and WV∈Rd×dW_V∈ R^d× d are learnable parameter matrices, D is a causal masking matrix, and GnmG_nm is a learnable pair-wise positional bias. It can be seen that the proposed MoFormer reduces computational complexity by leveraging a linear attention mechanism, while also enabling efficient Parallel training on GPUs. The Recurrent Representation of TALA. During inference, MixFormer can reformulate TALA as a recurrent computation, enabling autoregressive inference. This allows GPU memory consumption to remain constant, rather than growing with sequence length. The recurrent computation structure of TALA is illustrated in Figure 4, at the i-th time step, the computation proceeds through the following process: S0=0, S_0=0, (15) Sn=ωG⊙Sn−1+KnTVn, S_n=ω G S_n-1+K_n^TV_n, On=QnSn, O_n=Q_nS_n, where G is the same as in Equation (14). The TALA maintains linear-time complexity during inference. The Chunkwise Recurrent Representation of TALA. Chunkwise computation, an efficient training paradigm for long sequences, achieves significant acceleration by combining the advantages of parallel and sequential computation [32, 40]. In this approach, the input sequence is divided into consecutive chunks. For each chunk, fully parallelized computation is performed to leverage hardware acceleration, while a recurrent state propagation mechanism is established across chunks to preserve long-range dependencies. Under the linear attention framework, this hybrid computation structure not only enables faster training compared to conventional methods but also exhibits notable I/O-aware benefits through optimized memory access. Specifically, let B denote the chunk length. The computation of the i-th chunk in TALA is defined as follows: Ri=(ωG)B⊙Rn−1+K[i]TV[i], R_i=(ω G)^B R_n-1+K_[i]^TV_[i], (16) TALAtt(X[i])=Q[i]⊙(K[i]T⊙D)V[i]⏟inner−chunk+Q[i]Ri−1⏟cross−chunk, TALAtt(X_[i])= Q_[i] (K_[i]^T D)V_[i]_inner-chunk+ Q_[i]R_i-1_cross-chunk, where [i][i] is the i-th chunk. 4.4 Mixture of Memory Experts In State Space Models (SSMs), a single memory state architecture suffers from inherent limitations in long-sequence reasoning: as the sequence length increases, information from early inputs inevitably decays. While increasing the memory state dimension can partially mitigate this issue, it fundamentally fails to address the challenge of information forgetting under long-range dependencies. We propose a Mixture-of-Memory-Experts (MoME) mechanism based on the Mixture-of-Experts (MoE) framework [29, 23] to overcome this problem. This mechanism deploys multiple functionally heterogeneous memory expert networks, endowing the model with dynamic and differentiated memory capabilities: each expert specializes in memory patterns of varying temporal spans and information importance, and the system adaptively activates the most relevant subset of experts according to the current contextual state. This design not only significantly enhances the model’s ability to recall historical information but also achieves long-term retention of critical information through a multi-granularity memory preservation strategy. The MoME mechanism in TALA extends the memory state by maintaining multiple memory experts. As shown in Figure 2 (c), at each time step, the system dynamically activates a subset of experts based on the current input and constructs the current memory state by weighted aggregation of their outputs. This mechanism significantly enhances the model’s ability to capture temporal features and salient information. The computation is formally defined as follows: Yms=∑i=1nEiMSi, Y_ms= _i=1^nE_iMS_i, (17) where YmsY_ms denotes the weighted output of multiple memory experts, MSiMS_i represents the i-th memory expert, EiE_i is a router weight vector that determines the contribution of each memory state expert, as follows: Ei=Softmax(KnVnWg), E_i=Softmax(K_nV_nW_g), (18) ∑j=1keEij=1, _j=1^k_eE_i^j=1, where wg∈ℝ1×kewg ^1× k_e is a learnable parameter matrix, and kek_e is the number of experts. The current time step’s KnK_n is processed through the kernel transformation defined in Equation (11). Meanwhile, TALA enables the memory update at each time step to dynamically emphasize salient information while suppressing less relevant content. With the stacking of multiple layers and the collaborative effect of the MoME mechanism, the model can effectively retain and utilize information from early tokens even after processing ultra-long sequences. 4.5 Overall Architecture of MixFormer The proposed MixFormer consists of N stacked MixFormer blocks (Equation (9)). The input to the first block is the embedding of X, comprising both token and positional representations, while each subsequent block receives the output of the previous block as its input. The computation is defined as follows: Yl=TALA(GN(Xl))+Xl, Y^l=TALA(GN(X^l))+X^l, (19) Xl+1=FFN(GN(Yl))+Yl, X^l+1=FFN(GN(Y^l))+Y^l, where GN(⋅)GN(·) denotes the Group Normalization operation (GN). MixFormer adopts a structure similar to that of mainstream Transformers and employs a Feed-Forward Network (FFN) to enable the modeling of complex nonlinear transformations. The computation is defined as follows: FFN(X)=max(0,XW1)W2,FFN(X)=max(0,XW_1)W_2, (20) where W1W_1 and W2W_2 are learnable weight matrices. 4.6 Complexity Analysis To verify the linear-time complexity of MixFormer, we conduct a detailed complexity analysis as follows. Let N denote the sequence length (context length), d the hidden state dimension, d′d the expansion dimension of the feed-forward network (FFN), where d′>d >d, and k the number of memory experts. The total computational cost of MixFormer comprises four main components: projections (Proj), the time-aware linear attention (TALA), mixture-of-memory experts (MoME), and the feed-forward network (FFN). As shown in Equation 21, the Projection cost (Query, Key, Value, and Output projections) is 4Nd24Nd^2, the time-aware linear attention cost is Nd2Nd^2, the MoME cost is NkNk, and the FFN cost is 2dd′2d . Under the long-context setting (N≫dN d), the total computational complexity of MixFormer converges to O(N)O(N). 4Nd2⏟Proj+Nd2⏟TALA+Nk⏟MoME+2dd′⏟FFN 4Nd^2_Proj+ Nd^2_TALA+ Nk_MoME+ 2d _FFN (21) 4.7 Differences from Previous Methods To compare MixFormer with other efficient Linear Transformers, we conduct a comparative analysis across four key dimensions, as shown in Table 2. For training parallelism, MixFormer, like the standard Transformer, supports fully parallelized training. In contrast, Recurrent Neural Networks (RNNs) and RWKV [25] suffer from sequential training bottlenecks. MixFormer achieves performance comparable to that of standard Transformers, while significantly outperforming Linear Transformers—whose kernel approximations often lead to degraded attention focus—as well as traditional RNN-based models. For inference efficiency, MixFormer and linear Transformers both achieve O(1)O(1) inference time per token, significantly outperforming the O(N2)O(N^2) complexity of standard Transformers. For memory consumption, MixFormer maintains linear O(N)O(N) memory usage with respect to sequence length, in contrast to the quadratic memory growth of standard Transformers. MixFormer demonstrates significant advantages in its memory state mechanism. Unlike conventional Transformers that rely on explicit attention score matrices and lack an inherent memory design, existing linear Transformers and SSMs typically maintain a single static memory state, which limits their ability to model complex dependencies over long sequences. In contrast, MixFormer uses a mixture-of-expert memory mechanism, where a dynamic routing strategy assigns input features to multiple specialized memory modules, enabling the model to adaptively activate historical information across different dimensions. This multi-granularity collaborative memory architecture allows MixFormer to capture fine-grained temporal patterns and retain information more effectively in long-sequence tasks. Finally, compared to S4 [9] and RetNet [32], MixFormer achieves stronger overall performance while preserving both training Parallelism and linear time and memory complexity. In summary, the proposed MixFormer effectively integrates the advantages of existing efficient Transformers and achieves significant performance improvements. Table 2: Comparison between MixFormer and other efficient linear Transformer models from multiple perspectives. Including training parallelism, inference cost, memory consumption on long sequences, and the number of memory states. Architectures Training Parallelization Inference Cost Long-Sequence Memory Complexity Memory State Transformer Y O(N)O(N) O(N2)O(N^2) N/a Linear Transformer Y O(1)O(1) O(N)O(N) Singe Recurrent N N O(1)O(1) O(N)O(N) Singe RWKV N O(1)O(1) O(N)O(N) Singe S4 Y O(1)O(1) O(NlogN)O(NlogN) Singe RetNet Y O(1)O(1) O(N)O(N) Singe MixFormer (Ours) Y O(1)O(1) O(N)O(N) Multiple 5 Experiments Table 3: Comparison of MixFormer with other efficient Transformers on the LRA benchmark (%\%). The best results are highlighted in bold, and the second-best results are underlined. The results of other models follow [33]. Model ListOps Text Retrieval Image Pathfinder Avg Transformer 36.37 64.27 57.46 42.44 71.40 54.39 Local Attention 15.82 52.98 53.39 41.46 66.63 46.06 Sparse Trans. 17.07 63.58 59.59 44.24 71.71 51.24 Longformer 35.63 62.85 56.89 42.22 69.71 53.46 Linformer 35.70 53.94 52.27 38.56 76.34 51.36 Reformer 37.27 56.10 53.40 38.07 68.50 50.67 Sinkhorn Trans. 33.67 61.20 53.83 41.23 67.45 51.39 Synthesizer 36.99 61.68 54.67 41.61 69.45 52.88 BigBird 36.05 64.02 59.29 40.83 74.87 55.01 Linear Trans. 16.13 65.90 53.09 42.34 75.30 50.55 Performer 18.01 65.40 53.82 42.77 77.05 51.41 MixFormer (Ours) 37.35 66.56 60.50 42.65 76.55 56.72 5.1 Experimental Setups The main hyperparameter settings of MixFormer are summarized in Table 4. In this experiment, MixFormer was configured with two parameter variants: A0.3B-2B and A1B-7B. The A0.3B-2B model uses a hidden dimension of 1024 for the Q/K/V/O projections, a feed-forward network dimension of 896, 8 attention heads, 64 memory state experts, and 12 stacked MixFormer blocks. For training, we adopt a learning rate of 1×10−41× 10^-4 with the Adam optimizer, and train the model on a corpus containing 15 billion tokens. The total number of parameters is 2B, of which 0.3B are active parameters. The A1B-7B model uses a hidden dimension of 2048 for the Q/K/V/O projections, a feed-forward network dimension of 1024, 16 attention heads, 64 memory state experts, and 16 stacked MixFormer blocks. For training, we adopt a learning rate of 1×10−51× 10^-5 with the Adam optimizer, and train the model on a corpus containing 15 billion tokens. The total number of parameters is 7B, of which 1B is an active parameter. All experiments are conducted on four NVIDIA RTX 4090 GPUs. Table 4: Experimental Setups. HyperParameters A0.3B-2B A1B-7B Hidden Dimension 1024 2048 FFN Dimension 896 1024 Num of Heads 8 16 Num of Layers 12 16 Num of Memory Experts 64 64 LR 1e-4 1e-5 Training Tokens 15B 30B 5.2 Long Context Ability To evaluate the effectiveness of MixFormer in long-context scenarios, we conduct experiments on the Long-Range Arena (LRA) benchmark [33], a standard suite for long-sequence modeling. The Long Range Arena (LRA) benchmark systematically integrates five core tasks characterized by pronounced long-sequence dependencies. In the mathematical expression parsing task ListOps [22], models are required to parse deeply nested operator structures, posing a significant challenge to their hierarchical reasoning ability. The text classification task based on IMDb movie reviews [20] demands that models extract key sentiment features from lengthy reviews, testing their capacity for semantic filtering under noisy conditions. The AAN document retrieval task [27] evaluates how effectively models capture semantic associations across long academic documents through citation-based matching. The Pathfinder task [18], a synthetic visual challenge, requires the establishment of long-range spatial dependencies across pixel sequences. Finally, the CIFAR-10 image classification task [16], which flattens images into sequences, innovatively examines a model’s ability to preserve visual features in the absence of two-dimensional structural priors. The experimental results, as shown in Table 3, demonstrate that MixFormer outperforms all compared efficient Transformer models in terms of overall average accuracy. Specifically, it achieves the best performance on ListOps, text classification, and document retrieval tasks, and ranks second on CIFAR-10 and Pathfinder. This can be attributed to MixFormer’s use of linear attention as the backbone, which enables longer sequence outputs under limited memory consumption, thereby enhancing contextual representation learning. In addition, the MoME mechanism significantly broadens the model’s perceptual scope, enabling it to continuously focus on critical information in long-context scenarios. By coordinating multiple specialized experts, this mechanism ensures that important content is effectively preserved and reinforced throughout long-sequence propagation, thereby enhancing the model’s ability to capture long-range dependencies. Overall, MixFormer exhibits strong capability in modeling long-range dependencies across diverse domains. 5.3 Image Generation To evaluate the effectiveness of MixFormer in image generation tasks, we conduct both conditional and unconditional generation experiments on the MNIST dataset [3]. Figure 5 (top row) is the result of unconditional generation. While standard Transformers can be directly applied to image generation, they suffer from substantial memory overhead when modeling long sequences in pixel-by-pixel generation. In contrast, MixFormer not only significantly reduces memory consumption and generation time but also produces high-quality samples with sharp boundaries and minimal noise. Figure 5 (bottom row) presents an image completion task, where the input consists of occluded images and MixFormer is responsible for generating the complete outputs. We observe that MixFormer demonstrates strong long-range dependency modeling capabilities, accurately replicating the stroke style and width characteristics of the original images. These advantages highlight the potential of MixFormer as an efficient and effective solution for long-sequence image generation tasks. In summary, the proposed MixFormer performs remarkably well in both unconditional and conditional image generation tasks, demonstrating its effectiveness in modeling long-range pixel dependencies. Figure 5: Experimental results on image generation tasks. The top row is the result of unconditional generation, while the bottom row is conditional generation. (a) is the occluded input image, (b) is the output generated by MixFormer, and (c) is the corresponding ground-truth image. 5.4 CIFAR 10 To evaluate the performance of MixFormer on multi-channel image tasks, we conduct an image completion experiment on the CIFAR-10 [16]. Compared with single-channel inputs, the serialized representation of multi-channel images results in substantially longer sequences, imposing greater demands on the model’s ability to capture long-range dependencies. As shown in Figure 6, when 50% of image pixels are randomly masked, MixFormer successfully reconstructs visually coherent images with well-preserved structural integrity, clear boundaries, and minimal noise—significantly outperforming baseline models in texture continuity and detail preservation. Within the MoME architecture, different memory experts specialize in capturing local texture features and global structural information. When parts of an image are occluded, the model adaptively activates the corresponding structural experts to recover the overall layout while engaging texture experts to refine fine-grained details. In summary, the synergistic interaction between the MoME and the TALA mechanisms endows MixFormer with distinct advantages in multi-channel image completion tasks. Figure 6: Experimental results on CIFAR 10. (a) is the occluded input image, and (b) is the output generated by MixFormer. 5.5 Interpretability Analysis We conduct an interpretability analysis by visualizing intermediate outputs between the standard dot-product attention and reflection attention. Figure 7 (a) shows the intermediate result of the standard softmax attention: the matrix QKTQK^T after softmax normalization. Softmax attention not only strengthens the connections between important tokens but also suppresses irrelevant ones. Moreover, it preserves full pairwise interactions among all tokens, ensuring comprehensive global information exchange. Any given token can retrieve interaction information with all other tokens through the attention map. In contrast, Figure 7 (b) presents the KTVK^TV output in the reflection attention, which forms a fixed-size d×d× d matrix known as the State Space Model (SSM). This matrix serves as a compressed global memory of historical information. Due to its fixed dimensionality, the SSM does not retain explicit step-wise positional indices. At each time step, the token updates the SSM by accumulating information from all previous steps. At each time step, token representations are selectively written into the SSM based on the current input, and the output is generated using the updated global memory. Figure 7: Visualized intermediate outputs of (a) standard dot-product attention and (b) reflection attention. In other words, softmax attention produces explicit token-pair attention scores, while linear attention constructs an implicit global attention matrix. During autoregressive inference, MixFormer enforces causal masking through the matrix D (Equation 14). As illustrated in Figure 8, the visualization of the D matrix clearly reveals the temporal evolution of memory attention values—an effect that arises from our exponential decay mechanism and learnable pairwise interaction biases. In summary, the causal mask analysis demonstrates that MixFormer effectively balances strict causal constraints with global information flow. Meanwhile, the multi-expert case study reveals a natural functional differentiation among memory experts—where some focus on short-term feature extraction while others specialize in long-term dependency preservation—offering a mechanistic explanation for MixFormer’s superior performance on long-sequence tasks. Figure 8: Visualization of the causal mask used by MixFormer. 5.6 Mixture-of-Memory-Experts Analysis To examine the activation behavior of the MoME mechanism, we conduct a visualization analysis. As shown in Figure 9, when the current token is “Paris” and the context is “The Eiffel Tower is located in the city of”, the MoME mechanism exhibits significantly stronger expert activations for the tokens “Eiffel”, “Tower”, and “city”. Through the collaboration of multiple memory experts, the model effectively identifies the strong semantic association between the landmark entity “Eiffel Tower” and “Paris”. Meanwhile, the expert system captures the spatial hierarchical relationship between “city” and the target location “Paris”. In particular, the long-term memory experts successfully preserve information about distant core entities, ensuring accurate contextual recall during generation. This observation confirms that the MoME mechanism achieves dynamic selection and persistent retention of crucial historical information through functional specialization among experts. Figure 9: Visualization of memory-expert activations. 6 Conclusion In this work, we proposed MixFormer, a novel linear Transformer architecture incorporating the MoE to address the limitations of existing SSMs in ultra-long sequence modeling. By maintaining multiple differentiated memory states in parallel and introducing the Time-Aware Linear Attention (TALA) mechanism, MixFormer dynamically updates memory with input-adaptive exponential decay and learnable positional biases. This enables the model to selectively retain important historical information while mitigating memory dilution, effectively enhancing its capacity for long-range dependency modeling. Extensive experiments on both long-sequence text and image generation tasks demonstrate that MixFormer consistently outperforms prior approaches, achieving robust performance under both parallel training and recurrent inference modes. Finally, MixFormer provides a promising solution for efficient Transformer design with strong modeling capabilities for ultra-long contexts. In future work, we plan to improve MixFormer in two directions. First, we aim to scale up the model parameters to further enhance its performance. Second, we will extend the application of MixFormer to a broader range of tasks and domains. References [1] H. Cai, C. Gan, and S. Han (2022) Efficientvit: enhanced linear attention for high-resolution low-computation visual recognition. arXiv preprint arXiv:2205.14756 3 (1). Cited by: §4.3. [2] K. M. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Q. Davis, A. Mohiuddin, L. Kaiser, et al. (2021) Rethinking attention with performers. In Proceedings of the 2021 International Conference on Learning Representations, ICLR, Cited by: §3.1. [3] G. Cohen, S. Afshar, J. Tapson, and A. Van Schaik (2017) EMNIST: extending mnist to handwritten letters. In 2017 international joint conference on neural networks (IJCNN), p. 2921–2926. Cited by: §5.3. [4] W. Dai, N. Lee, B. Wang, Z. Yang, Z. Liu, J. Barker, T. Rintamaki, M. Shoeybi, B. Catanzaro, and W. Ping (2024) Nvlm: open frontier-class multimodal llms. arXiv preprint arXiv:2409.11402. Cited by: §1. [5] T. Dao and A. Gu (2024) Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 2024 International Conference on Machine Learning, ICML, p. 10041–10071. Cited by: §3.2. [6] M. Deitke, C. Clark, S. Lee, R. Tripathi, Y. Yang, J. S. Park, M. Salehi, N. Muennighoff, K. Lo, L. Soldaini, et al. (2024) Molmo and pixmo: open weights and open data for state-of-the-art multimodal models. arXiv e-prints, p. arXiv–2409. Cited by: §1. [7] A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang, A. Fan, et al. (2024) The llama 3 herd of models. arXiv e-prints, p. arXiv–2407. Cited by: §1. [8] A. Gu and T. Dao (2024) Mamba: linear-time sequence modeling with selective state spaces. In Proceedings of the 2024 International First Conference on Language Modeling, CoLM, Cited by: §3.2. [9] A. Gu, K. Goel, and C. Re (2022) Efficiently modeling long sequences with structured state spaces. In Proceedings of the 2022 International Conference on Learning Representations, ICLR, Cited by: §1, §3.2, §4.7. [10] A. Gu, I. Johnson, K. Goel, K. Saab, T. Dao, A. Rudra, and C. Ré (2021) Combining recurrent, convolutional, and continuous-time models with linear state space layers. Proceedings of the 2021 Advances in Neural Information Processing Systems, NeuraIPS 34, p. 572–585. Cited by: §1. [11] A. Gupta, A. Gu, and J. Berant (2022) Diagonal state spaces are as effective as structured state spaces. Proceedings of the 2022 Advances in Neural Information Processing Systems, NeuraIPS 35, p. 22982–22994. Cited by: §1. [12] D. Han, X. Pan, Y. Han, S. Song, and G. Huang (2023) Flatten transformer: vision transformer using focused linear attention. In Proceedings of the 2023 International Conference on Computer Vision, ICCV, p. 5961–5971. Cited by: §2.3, §3.1. [13] K. Han, Y. Wang, H. Chen, X. Chen, J. Guo, Z. Liu, Y. Tang, A. Xiao, C. Xu, Y. Xu, et al. (2022) A survey on vision transformer. IEEE Transactions on Pattern Analysis and Machine Intelligence 45 (1), p. 87–110. Cited by: §1. [14] R. Hasani, M. Lechner, T. Wang, M. Chahine, A. Amini, and D. Rus (2023) Liquid structural state-space models. In Proceedings of the 2023 International Conference on Learning Representations, ICLR, Cited by: §1. [15] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret (2020) Transformers are rnns: fast autoregressive transformers with linear attention. In Proceedings of the 2020 International Conference on Machine Learning, ICML, p. 5156–5165. Cited by: §2.2, §2.4, §3.1. [16] A. Krizhevsky, G. Hinton, et al. (2009) Learning multiple layers of features from tiny images. Cited by: §5.2, §5.4. [17] R. Li, J. Su, C. Duan, and S. Zheng (2020) Linear attention mechanism: an efficient attention for semantic segmentation.. arXiv preprint arXiv:2007.14902. Cited by: §3.1. [18] D. Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre (2018) Learning long-range spatial dependencies with horizontal gated recurrent units. Proceedings of the 2018 Advances in Neural Information Processing Systems, NeuraIPS 31. Cited by: §5.2. [19] J. Lu, J. Yao, J. Zhang, X. Zhu, H. Xu, W. Gao, C. Xu, T. Xiang, and L. Zhang (2021) Soft: softmax-free transformer with linear complexity. Proceedings of the 2021 Advances in Neural Information Processing Systems, NeuraIPS 34, p. 21297–21309. Cited by: §4.3. [20] A. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts (2011) Learning word vectors for sentiment analysis. In Proceedings of the 2011 Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, ACL, p. 142–150. Cited by: §5.2. [21] W. Meng, Y. Luo, X. Li, D. Jiang, and Z. Zhang (2025) PolaFormer: polarity-aware linear attention for vision transformers. In Proceedings of the 2025 International Conference on Learning Representations, ICLR, Cited by: §2.3, §3.1. [22] N. Nangia and S. Bowman (2018) ListOps: a diagnostic dataset for latent tree learning. In Proceedings of the 2018 Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, ACL, p. 92–99. Cited by: §5.2. [23] S. Narayan, S. B. Cohen, and M. Lapata (2018) Don’t give me the details, just the summary! topic-aware convolutional neural networks for extreme summarization. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, EMNLP, p. 1797–1807. Cited by: §4.4. [24] L. Papa, P. Russo, I. Amerini, and L. Zhou (2024) A survey on efficient vision transformers: algorithms, techniques, and performance benchmarking. IEEE Transactions on Pattern Analysis and Machine Intelligence 46 (12), p. 7682–7700. Cited by: §1. [25] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, et al. (2023) RWKV: reinventing rnns for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, p. 14048–14077. Cited by: §3.2, §4.7. [26] H. Peng, N. Pappas, D. Yogatama, R. Schwartz, N. Smith, and L. Kong (2021) Random feature attention. In Proceedings of the 2021 International Conference on Learning Representations, ICLR, Cited by: §3.1. [27] D. R. Radev, P. Muthukrishnan, V. Qazvinian, and A. Abu-Jbara (2013) The acl anthology network corpus. Language Resources and Evaluation 47 (4), p. 919–944. Cited by: §5.2. [28] I. Schlag, K. Irie, and J. Schmidhuber (2021) Linear transformers are secretly fast weight programmers. In Proceedings of the 2021 International Conference on Machine Learning, ICML, p. 9355–9366. Cited by: §2.3. [29] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean (2017) Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. In Proceedings of the 2017 International Conference on Learning Representations, ICLR, Cited by: §4.4. [30] Z. Shen, M. Zhang, H. Zhao, S. Yi, and H. Li (2021) Efficient attention: attention with linear complexities. In Proceedings of the 2021 International Conference on Computer Vision, ICCV, p. 3531–3539. Cited by: §2.2, §3.1. [31] J. T. Smith, A. Warrington, and S. W. Linderman (2023) Simplified state space layers for sequence modeling. In Proceedings of the 2023 International Conference on Learning Representations, ICLR, Cited by: §1. [32] Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei (2023) Retentive network: a successor to transformer for large language models. arXiv preprint arXiv:2307.08621. Cited by: §1, §3.2, §4.3, §4.7. [33] Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler (2020) Long range arena: a benchmark for efficient transformers. In Proceedings of the 2020 International Conference on Learning Representations, ICLR, Cited by: §5.2, Table 3. [34] Y. Tay, M. Dehghani, D. Bahri, and D. Metzler (2022) Efficient transformers: a survey. ACM Computing Survey 55 (6). Cited by: §1. [35] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. (2023) Llama: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: §1. [36] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. (2023) Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: §1. [37] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Proceedings of the 2017 Advances in Neural Information Processing Systems, NeuraIPS, p. 5998–6008. Cited by: §1, §2.1, §3.1, §4.1. [38] Y. Wu and K. He (2018) Group normalization. In Proceedings of the European conference on computer vision (ECCV), p. 3–19. Cited by: §4.1. [39] Y. Xiong, Z. Zeng, R. Chakraborty, M. Tan, G. Fung, Y. Li, and V. Singh (2021) Nyströmformer: a nyström-based algorithm for approximating self-attention. In Proceedings of the 2021 Conference on Artificial Intelligence, AAAI, p. 14138–14148. Cited by: §4.3. [40] S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim (2024) Gated linear attention transformers with hardware-efficient training. In Proceedings of the 2024 International Conference on Machine Learning, ICML, p. 56501–56523. Cited by: §1, §3.2, §4.3. [41] S. Zhai, W. Talbott, N. Srivastava, C. Huang, H. Goh, R. Zhang, and J. Susskind (2021) An attention free transformer. arXiv preprint arXiv:2105.14103. Cited by: §3.2. [42] B. Zhuang, J. Liu, Z. Pan, H. He, Y. Weng, and C. Shen (2023) A survey on efficient training of transformers. In Proceedings of the 2023 International Joint Conference on Artificial Intelligence, IJCAI, p. 6823–6831. Cited by: §1.