Mamba: Linear-Time Sequence Modeling with Selective State Spaces论文总结

向岸看

已于 2024-09-12 16:09:39 修改

阅读量281

点赞数 8

文章标签： Mamba S6 SSM

于 2024-09-07 11:33:06 首次发布

本文链接：https://blog.csdn.net/qq_45981086/article/details/141675075

版权

题目：Mamba: Linear-Time Sequence（线性时间序列） Modeling with Selective State Spaces（选择性状态空间）

论文：Mamba: Linear-Time Sequence Modeling with Selective State Spaces

源码：state-spaces/mamba: Mamba SSM architecture (github.com)

推荐课程：序列模型、状态空间结构模型(SSM)、Mamba模型原理与代码精讲_哔哩哔哩_bilibili

0. 预备知识：RNN、LSTM、SSM、S4

一、摘要

二、引言

三、State Space Models（状态空间模型）

四、Selective State Space Models （选择性状态空间模型 S6）

4.1 动机：Selection as a Means of Compression（选择作为压缩的一种手段）

4.2 Improving SSMs with Selection（使用选择机制改进SSMs）

4.3 Overview of Selective Scan: Hardware-Aware State Expansion（选择性扫描概述：硬件感知状态扩展）

五、结论

0. 预备知识：RNN、LSTM、SSM、S4

RNN（循环神经网络）

通常情况下，深度神经网络都是在水平方向上延伸的，如CNN。虽然在延伸的过程中增加了隐藏层数量，但是没有考虑单个隐藏层在时间维度上的变换，导致像CNN一样的网络没有办法学习到序列数据的上下文关系，因此不适合直接处理时间序列数据（如文本）。RNN通过建立隐藏层不同时刻之间的迭代关系（加性），让神经网络拥有短期记忆的能力。

线性神经网络的输出：   $\mathrm{S=f(W_{in}X + b)}$

RNN的输出：


代码实现：
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(RNN, self).__init__()

        self.hidden_size = hidden_size

        self.u = nn.Linear(input_size, hidden_size)
        self.w = nn.Linear(hidden_size, hidden_size)
        self.v = nn.Linear(hidden_size, output_size)

        self.tanh = nn.Tanh()
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, inputs, hidden):

        u_x = self.u(inputs)

        hidden = self.w(hidden)
        hidden = self.tanh(hidden + u_x)

        output = self.softmax(self.v(hidden))

        return output, hidden
迭代：
    for i in range(line_tensor.size()[0]):
        output, hidden = rnn(line_tensor[i], hidden)
LSTM（长短期记忆网络）

LSTM在RNN的基础上，添加一条长期记忆链（顶部）和三个控制门（遗忘门、记忆门/输入门、输出门）。长期记忆链（顶部）和短期记忆链（底部）通过三个控制门实现相互更新。

遗忘门（Sigmoid + 逐元素相乘）：用于根据 短期记忆链（rnn）选择性删除 长期记忆链 中不重要的特征。它通过一个 sigmoid 函数来生成一个 0 到 1 之间的权重，表示每个信息在单元状态中的保留程度。

记忆门（Sigmoid + tanh + 逐元素相乘 + 逐元素相加）：用于根据 短期记忆链（rnn）选择性地增强或减弱 长期记忆链 中重要的特征。

输出门（Sigmoid + tanh + 逐元素相乘）：用于根据 长期记忆链 选择性地增强或减弱 短期记忆链（rnn） 中重要的特征。

代码实现：
class CustomLSTM(nn.Module):
    def __init__(self, input_sz, hidden_sz):
        super().__init__()
        self.input_sz = input_sz
        self.hidden_size = hidden_sz
        self.W = nn.Parameter(torch.Tensor(input_sz, hidden_sz * 4))
        self.U = nn.Parameter(torch.Tensor(hidden_sz, hidden_sz * 4))
        self.bias = nn.Parameter(torch.Tensor(hidden_sz * 4))
        self.init_weights()
 
    def init_weights(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)
 
    def forward(self, x, init_states=None):
        """Assumes x is of shape (batch, sequence, feature)"""

        bs, seq_sz, _ = x.size()
        hidden_seq = []

        if init_states is None:
            h_t, c_t = (torch.zeros(bs, self.hidden_size).to(x.device), 
                        torch.zeros(bs, self.hidden_size).to(x.device))
        else:
            h_t, c_t = init_states
 
        HS = self.hidden_size

        for t in range(seq_sz):
            x_t = x[:, t, :]

            # batch the computations into a single matrix multiplication
            gates = x_t @ self.W + h_t @ self.U + self.bias

            i_t, f_t, g_t, o_t = (
                torch.sigmoid(gates[:, :HS]),             # input
                torch.sigmoid(gates[:, HS:HS*2]),         # forget
                torch.tanh(gates[:, HS*2:HS*3]),
                torch.sigmoid(gates[:, HS*3:]),           # output
            )

            # forget gate + input gate
            c_t = f_t * c_t + i_t * g_t
            
            # output gate
            h_t = o_t * torch.tanh(c_t)
            
            hidden_seq.append(h_t.unsqueeze(0))

        hidden_seq = torch.cat(hidden_seq, dim=0)
        hidden_seq = hidden_seq.transpose(0, 1).contiguous()
        return hidden_seq, (h_t, c_t)
SSM（State Space Model 状态空间模型）

参考：一文通透想颠覆Transformer的Mamba：从SSM、HiPPO、S4到Mamba_mamba模型-CSDN博客

状态空间模型（SSMs）是用于描述当前状态表示并根据某些输入进行下一个状态预测的模型。SSMs的输入和输出一般是连续数据。

state equation（状态公式）： $\mathrm{h'(t)=Ah(t)+Bx(t)}$ (RNN)（底部分支）

output equation（输出公式）： $\mathrm{y(t)=Ch(t)+Dx(t)}$ （根据RNN更新后的隐藏状态h(t)，在经过一个矩阵C加权，与经过矩阵D加权的x(t)进行加和）（本质上就是权重矩阵 + 跳连接操作）（顶部分支）

更简洁的SSM模型图示：

S4（Structured State Space for Sequences 结构化的状态空间模型）

S4论文：Efficiently Modeling Long Sequences with Structured State Spaces

状态空间模型SSM的离散化。由于除了连续的输入数据之外，还会通常碰到离散的输入(如文本序列、图像)，因此我们还需要将模型离散化。这里就要使用 零阶保持技术 (Zero-order hold 技术）。

1）S4的离散化表示（Zero-order hold 技术）

首先，每次收到离散信号时，我们都会保留其值，直到收到新的离散信号，如此操作导致的结果就是创建了 SSM 可以使用的连续信号
保持该值的时间由一个新的可学习参数表示，称为步长(size)—— $\Delta$  ，它代表输入的阶段性保持(resolution)

模型公式更改（可以针对A、B按如下方式做Zero-order hold，实现离散化）：

最终使我们能够从连续 SSM 转变为离散SSM，使得SSM的输入和输出结果不再是一个函数到函数， $x(t)\rightarrow y(t)$ ，而是一个序列到序列， $x_k \rightarrow y_k$ 。

y2的展开如下：

推而广之，可得

$y_{k}=C \bar{A}^{k} \bar{B} x_{0}+C \bar{A}^{k-1} \bar{B} x_{1}+\cdots+C \bar{A} \bar{B} x_{k-1}+C \bar{B} x_{k}$

在离散化的SSM中，根据时间步长进行迭代，在每个时间步长中，我们计算当前输入( $\mathrm{Bx_k}$ )如何影响前一个状态( $\mathrm{Ah_{k-1}}$ )，然后计算预测输出( $\mathrm{Ch_k}$ )。

代码实现：Structured State Space for Sequences
def random_SSM(rng, N):
    a_r, b_r, c_r = jax.random.split(rng, 3)
    A = jax.random.uniform(a_r, (N, N))
    B = jax.random.uniform(b_r, (N, 1))
    C = jax.random.uniform(c_r, (1, N))
    return A, B, C

# 离散化
def discretize(A, B, C, step):
    I = np.eye(A.shape[0])
    BL = inv(I - (step / 2.0) * A)
    Ab = BL @ (I + (step / 2.0) * A)
    Bb = (BL * step) @ B
    return Ab, Bb, C

def scan_SSM(Ab, Bb, Cb, u, x0):
    def step(x_k_1, u_k):
        x_k = Ab @ x_k_1 + Bb @ u_k
        y_k = Cb @ x_k
        return x_k, y_k

    return jax.lax.scan(step, x0, u)

def run_SSM(A, B, C, u):
    L = u.shape[0]
    N = A.shape[0]
    Ab, Bb, Cb = discretize(A, B, C, step=1.0 / L)

    # Run recurrence
    return scan_SSM(Ab, Bb, Cb, u[:, np.newaxis], np.zeros((N,)))[1]
2）S4的卷积化表示（S4实现并行计算）

一文通透想颠覆Transformer的Mamba：从SSM、HiPPO、S4到Mamba_mamba模型-CSDN博客

根据推导式 $y_{k}=C \bar{A}^{k} \bar{B} x_{0}+C \bar{A}^{k-1} \bar{B} x_{1}+\cdots+C \bar{A} \bar{B} x_{k-1}+C \bar{B} x_{k}$ ，可得卷积核的计算公式：

我们只要通过矩阵乘积就可以得到我们想要的卷积核，并且可以实现并行计算。（这是对RNN的一个极大的提示，之前RNN由于串行计算，导致训练速度缓慢，而S4并行计算实现了加速训练）

3）HiPPO（High-order Polynomial Projection Operator 高阶多项式投影）为S4增加长距离依赖

问题：S4本质上还是一个RNN模型，RNN有一个致命的缺点，它只能保持短期记忆没有办法进行长距离依赖的提取。

如我们之前在循环表示中看到的那样，矩阵 $A$ 捕获先前状态的信息来构建新状态( $h_k = \overline{A} h_{k-1} + \overline{B} x_k$ ，当 $k = 5$ 时，则有 $h_5 = \overline{A} h_{4} + \overline{B} x_5$ )。这就要求SSM中设置矩阵 $A$ 时不能像RNN一样在迭代过程中会遗忘掉非常靠前的信息。怎么做呢？

解决方法：一种方法是可以通过使用HiPPO矩阵。

对于 $A_{nk}$ 而言，假设 $n$ (代表行)和 $k$ (代表列) 的范围是 0 到 2，由上式可得：

原理：HiPPO 矩阵可以产生一个隐藏状态来记住其历史(从数学上讲，它是通过跟踪 正交多项式的系数来实现的，这使得它能够逼近所有以前的历史)，使得在被应用于循环表示和卷积表示中时，可以处理远程依赖性。

举例：假设 $h_k = \overline{A} h_{k-1} + \overline{B} x_k$ 中，前一个隐藏状态 $h_{k-1}=(a_1,a_2,a_3)$ ，代入HiPPO矩阵所得的矩阵A，可得 $Ah_{k-1}$ 的乘积结果：

由上述乘积结果可知，在新的隐藏状态 $h_k$ 中，每一个元素值 都是由前一个隐藏状态各个分量元素通过不同的加权权重组成的，通过迭代这种方法“可以很好地捕获最近的token并衰减旧的token”，以此来处理远程依赖性。

注意，HiPPO矩阵 也是可学习的，例如，HiPPO-LegT矩阵 在 HiPPO矩阵 公式前添加了一个可学习参数 1/θ 使其变得可学习：

一、摘要

研究背景 / 研究问题：许多次二次时间架构，如线性注意力、门控卷积和循环模型，以及结构化状态空间模型 (SSM) 用于解决transformer在长序列上的计算效率低下的问题，但它们在语言等重要模态上的表现不如注意力。这些模型的一个关键弱点是无法进行基于内容的推理。

主要工作：

首先，简单地将SSM参数作为输入的函数，解决了离散模态的弱点，允许模型根据当前标记有选择地沿着序列长度维度传播或忘记信息。( 选择性扫描算法 )
第二，这种改变未使用高效卷积，但设计了一种 递归模式下的硬件感知并行算法 。将这些 selective SSMs 集成到一个简化的端到端神经网络架构中，无需注意力，甚至无需MLP模块 (Mamba)。

研究成果：Mamba 具有快速推理 (比Transformer高5倍的吞吐量) 和序列长度的线性缩放，其性能在高达百万长度的真实数据上有所提高。作为通用序列模型骨干，Mamba在语言、音频和基因组学等多种模态上取得了最先进的性能。在语言建模方面，Mamba-3B模型的表现优于相同大小的transformer，并在预训练和下游评估中比transformer的小两倍。吞吐量比Transformers高5倍。

二、引言

研究背景：基础模型（FM）的主干通常是单一类型的序列模型：Transformer 和 Self-Attention —> 根本的缺点：无法对有限窗口之外的任何东西进行建模，并且相对于窗口长度进行二次缩放

相关工作（SSMs）：结构化状态空间序列模型（SSM）可以被解释为递归神经网络（RNN）和卷积神经网络（CNN）的组合 —> 对离散的和信息密集的数据（如文本）建模时效率较低

主要工作：

Selection Mechanism. 设计了一种简单的选择机制，通过基于输入对SSM参数进行参数化。这使得模型可以过滤掉不相关的信息，并无限期地记住相关信息。
Hardware-aware Algorithm. 所有先前的SSM模型必须是时间和输入不变的，以便在计算上高效。本文通过硬件感知算法克服了这一问题，该算法通过扫描而不是卷积来循环计算模型，但不具体化扩展状态，以避免GPU内存层次结构的不同级别之间的IO访问。（在A100 GPU上比先前方法速度快了3倍）
Architecture（总体结构）：先前SSM架构的设计 + 变压器的MLP模块

三、State Space Models（状态空间模型）

以Structured State Space Models（S4）为例，S4 模型由四个参数（Δ、A、B、C）定义：

离散化：

S4的两个主要问题：

线性时间不变性(LTI) ：换句话说（Δ、A、B、C），以及（A、B）对于所有时间步长都是固定的。（对于 SSM 生成的每个token，矩阵A 、B、C都是相同的，使得SSM无法针对输入做针对性的推理，去除LTI约束的方法是参数化SSM，但这又会造成无法使用卷积进行高效的并行训练，带来了效率瓶颈）

结构和维度：为了在D个通道，批大小B和长度L的输入序列x上进行操作，SSM独立应用于每个通道。请注意，在这种情况下，每个输入的总隐藏状态具有DN维，并且在序列长度上计算它需要O(BLDN)时间和内存; 这就是 S6 讨论的基本效率瓶颈的根源。（这种方法并不高效）

四、Selective State Space Models （选择性状态空间模型 S6）

4.1 动机：Selection as a Means of Compression（选择作为压缩的一种手段）

作者认为序列建模的一个基本问题是将上下文压缩成一个更小的状态。

以Transformer和RNN/S4为例：

Transformer 的注意力机制虽然很有效（状态空间大），但效率低，因为它需要存储整个上下文，导致推理和训练时间较长。
SSM 递归模型因为它们的状态是有限的(单纯时不变导致状态空间很小)，效率高但有效性受限于状态的压缩能力。

如下图所示：

然而，Transformer/S4的有效性受到这种状态对上下文的压缩程度的限制。为了理解这一原理，作者提出关注两种能力：

选择性复制任务（抓重点的能力）：从内容感知推理，以便能够记住相关的标记并过滤掉不相关的标记。（从大量信息中选择和记住关键的信息，忽略不相关的部分。类似于在人群中找到你的朋友或者在一篇文章中找到关键词。比如从下面句子中找出名词。）

诱导头任务（上下文联想/推理能力）：由上下文感知推理来知道何时在适当的上下文中产生正确的输出。（在处理连续的信息时，能够保持逻辑一致性和上下文的连贯性。比如下图的回答能力，只用了单样本学习。）

那么具体怎么改进呢？我们之前说过SSM 有LTI线性时不变模型的强假设（ABC是固定的），这样就导致难以有效的选择上下文信息。这样是不行的，为了让SSM转变为时变模型，作者提出 “使用选择机制改进SSMs”。

4.2 Improving SSMs with Selection（使用选择机制改进SSMs）

将选择机制纳入模型的一种方法是让影响序列交互的参数(例如，RNN的递归动态或CNN的卷积核)依赖于输入。算法 1 和 2 说明了本文使用的主要选择机制：

主要改进：B，C的矩阵大小从原来的(D,N)变为(B,L,N)，步长 $\Delta$ 的大小由原来的D变成了(B,L,D)。这三个参数分别对应batch size、sequence length、hidden state size。对应的选择函数分别为

$s_{B}(x)=\operatorname{Linear}_{N}(x)$

$s_{C}(x)=\operatorname{Linear}_{N}(x)$

$s_{\Delta}(x)=\operatorname{Linear}_{D}(x)$

$\tau_{\Delta}=\text { softplus }$

特别是，作者强调这些参数现在具有长度维度L，这意味着模型已经从时间不变更改为时变。

SSM —> Mamba：

Q：为什么增加序列长度维度L之后，SSM从时间不变更改为时变模型呢？
A：我们来逐个分析 $\Delta$ 和 B，C 矩阵。

B，C 矩阵，引入了线性变换

首先我们来回顾一个SSM中B，C 矩阵，它们是直接由参数得到的，因此是固定不变的。而 Mamba 是通过 $S_B(x)$ ， $S_C(x)$ ，由输入 $x$ 经过线性变化得到，这意味着对于每个输入token，我们现在有不同的B和C矩阵，从而解决了内容感知的问题。

步长 $\Delta$ ，类似遗忘门

首先来看Softplus函数，与LSTM中Sigmoid函数的作用相同，Softplus函数会根据当前输入的重要与否，选择性地增强或减弱对应的步长 $\Delta$ （矩阵）中参数。如下图所示，较小的步长 $\Delta$ 导致忽略特定单词，而是更多地使用先前的上下文，而较大的步长则更专注于输入单词而不是上下文：

注意：矩阵 A 保持不变，因为我们希望状态本身保持静态，但通过 $\Delta$ ，B 和 C 影响的方式是动态的，简言之， $A$ 离散化之后 $\overline{\boldsymbol{A}}=\exp (\Delta \boldsymbol{A})$ ， $\Delta$ 的“输入数据依赖性”能够让整体的 $\bar{A}$ 与输入相关。

4.3 Overview of Selective Scan: Hardware-Aware State Expansion（选择性扫描概述：硬件感知状态扩展）

1）并行扫描

由于这些矩阵现在是动态的，所以不能使用卷积表示来计算，因为它假定了一个固定的核。我们只能使用循环表示，并丢失卷积提供的并行化。（效率下降到与RNN类似）

为了实现并行化，先探讨如何使用循环计算输出，如下图所示：

每个状态是前一个状态(乘以 $\bar{A}$ )和当前输入(乘以 $\bar{B}$ )的总和。这被称为扫描操作，可以很容易地用for 循环计算。相反，Mamba 通过并行扫描进行实现并行化。

假设我们执行操作的顺序与关联属性无关，因此，我们可以分段计算序列并迭代地组合它们：

（这里的B，其实是 $\bar{B}$ ）

我们将相关推导展开来看：

所有的 $\bar{A}$ 和 $\bar{B}$ 都是可以由输入 $X$ 得到的。H2可以直接由H1得到，也可以由H0。

H3可以直接由H2得到，也可以由H1或H0得到。

我们可以不遵守顺序加乘的规则，直接越级计算。这样可以使扫描操作的时间复杂度由 O(n) 降低到 O(n/t)。时间复杂度 O(n/t) 中的 t ，通常代表用于执行任务的处理器或计算单元的数量。

2）内核融合

问题：近期 GPU 的一个缺点是频繁地在 SRAM（高速缓存，直接用于GPU计算）和 DRAM（显存，负责存储）之间复制信息成为瓶颈。（IO缓慢）

解决方法：内核融合。Mamba，就像 flash Atention 一样，限制需要从 DRAM 到 SRAM 读取的次数。简而言之，隐藏状态 $h$ 更新的计算过程在 SRAM 中进行，其他参数A，B，C的更新过程都在 DRAM 中进行，这样就减少了从 DRAM 到 SRAM 读取的次数。

3) 重新计算

反向传播需要中间结果的值以计算梯度。如果所有中间结果都存储，会占用大量内存。重新计算策略可以减少所需存储的中间结果数量，只有在反向传播需要时才计算这些中间结果，从而节省内存。

4.4 一种简化的SSM结构

这受到门控注意力单元 (GAU) (Hua et al. 2022) 的启发，Mamba采用类似的架构，并参考H3架构，在主分支上添加了一个卷积，并用激活函数取代了第一个乘法门。该体系结构包括通过一个可控的扩展因子 E 来扩展模型维度 D。

六、结论

主要工作：在结构化状态空间模型中引入了一种选择机制，允许它们在序列长度线性伸缩的情况下执行上下文相关推理。

实验结果：当被整合到一个简单的免关注架构中时，Mamba在不同的域集合上实现了最先进的结果，在这些域中，它匹配或超过了强大的Transformer模型的性能。

展望：对选择性状态空间模型在不同领域构建基础模型的广泛应用感到兴奋，特别是在需要长上下文的新兴模式中，如基因组学，音频和视频。我们的结果表明，Mamba是一个强有力的候选者，作为一个通用的序列模式骨架。

七、源码

7.1 Mamba architecture

先来看Mamba的整体架构，

位置：mamba_ssm/modules/block.py

# Mamba architecture
class Block(nn.Module):
    def __init__(
        self, dim, mixer_cls, mlp_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False
    ):
        """
        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"

        This Block has a slightly different structure compared to a regular
        prenorm Transformer block.
        The standard block is: LN -> MHA/MLP -> Add.
        [Ref: https://arxiv.org/abs/2002.04745]
        Here we have: Add -> LN -> Mixer, returning both
        the hidden_states (output of the mixer) and the residual.
        This is purely for performance reasons, as we can fuse add and LayerNorm.
        The residual needs to be provided (except for the very first block).
        
        Args:
            @param dim: 维度
            @param mixer_cls: 混合器, Mamba, Mamba2
            @param mlp_cls: 中间维度
            @param norm_cls: 正则化方法
            @param fused_add_norm: 加法和归一化（normalization）
            @param residual_in_fp32: 在计算过程中使用 32 位浮点数来存储和处理残差
        """
        super().__init__()
        self.residual_in_fp32 = residual_in_fp32
        self.fused_add_norm = fused_add_norm
        self.norm = norm_cls(dim)
        self.mixer = mixer_cls(dim)
        if mlp_cls is not nn.Identity:
            self.norm2 = norm_cls(dim)
            self.mlp = mlp_cls(dim)
        else:
            self.mlp = None
        if self.fused_add_norm:
            assert RMSNorm is not None, "RMSNorm import fails"
            assert isinstance(
                self.norm, (nn.LayerNorm, RMSNorm)
            ), "Only LayerNorm and RMSNorm are supported for fused_add_norm"

    def forward(
            self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None, **mixer_kwargs
    ):
        r"""Pass the input through the encoder layer.

        Args:
            hidden_states: the sequence to the encoder layer (required).
            residual: hidden_states = Mixer(LN(residual))
        """
        if not self.fused_add_norm:
            residual = (hidden_states + residual) if residual is not None else hidden_states
            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
            if self.residual_in_fp32:
                residual = residual.to(torch.float32)
        else:
            hidden_states, residual = layer_norm_fn(
                hidden_states,
                self.norm.weight,
                self.norm.bias,
                residual=residual,
                prenorm=True,
                residual_in_fp32=self.residual_in_fp32,
                eps=self.norm.eps,
                is_rms_norm=isinstance(self.norm, RMSNorm)
            )

        # 隐藏状态
        hidden_states = self.mixer(hidden_states, inference_params=inference_params, **mixer_kwargs)

        if self.mlp is not None:
            if not self.fused_add_norm:
                residual = hidden_states + residual
                hidden_states = self.norm2(residual.to(dtype=self.norm2.weight.dtype))
                if self.residual_in_fp32:
                    residual = residual.to(torch.float32)
            else:
                hidden_states, residual = layer_norm_fn(
                    hidden_states,
                    self.norm2.weight,
                    self.norm2.bias,
                    residual=residual,
                    prenorm=True,
                    residual_in_fp32=self.residual_in_fp32,
                    eps=self.norm2.eps,
                    is_rms_norm=isinstance(self.norm2, RMSNorm)
                )
            hidden_states = self.mlp(hidden_states)

        return hidden_states, residual

    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)