arXiv reader

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu and Tri Dao

arXiv
2312.00752v2 · PDF · 31 May 2024
Categories
cs.LG
Licence
CC BY 4.0 The site shows the text.

Abstract

Foundation models.

Items 4

Theorem

Theorem 1

QK9W

When N=1, \bm{A}=-1, \bm{B}=1, s_\Delta=\mathsf{Linear}(x), and \tau_\Delta=\mathsf{softplus}, then the selective SSM recurrence (?) takes the form

\begin{aligned} g_t &= \sigma(\mathsf{Linear}(x_t)) \\ h_{t} &= (1-g_t) h_{t-1} + g_t x_t . \end{aligned}
Proof (Proof of Theorem 1) of Theorem 1

Consider a selective SSM (?) with N=1, \bm{A}=-1, \bm{B}=1, s_\Delta=\mathsf{Linear}(x), \tau_\Delta=\mathsf{softplus}. The corresponding continuous-time SSM (?) is

\begin{aligned} h(t) = -h(t) + x(t) \end{aligned}

which is also called a leaky integrator.

The discretization step size is

\begin{aligned} \Delta_t &= \tau_\Delta(\mathsf{Parameter} + s_\Delta(x_t)) \\ &= \mathsf{softplus}(\mathsf{Parameter} + \mathsf{Linear}(x_t)) \\ &= \mathsf{softplus}(\mathsf{Linear}(x_t)) \end{aligned}

where we observe that the parameter can be viewed as a learnable bias and folded into the linear projection.

Now applying the zero-order hold (ZOH) discretization formulas:

\begin{aligned} \overline{\bm{A}}_t &= \exp(\Delta \bm{A}) = \frac{1}{1 + \exp(\mathsf{Linear}(x_t))} = \sigma(-\mathsf{Linear}(x_t)) \\&= 1 - \sigma(\mathsf{Linear}(x_t)) \\ \overline{\bm{B}}_t &= (\Delta \bm{A})^{-1} (\exp(\Delta \bm{A}) - \bm{I}) \cdot \Delta \bm{B} = -(\exp(\Delta \bm{A}) - \bm{I}) = 1 - \overline{\bm{A}} \\&= \sigma(\mathsf{Linear}(x_t)) . \end{aligned}

Thus the final discrete recurrence (?) is

\begin{aligned} g_t &= \sigma(\mathsf{Linear}(x_t)) \\ h_{t} &= (1-g_t) h_{t-1} + g_t x_t \end{aligned}

as desired.

Remark

Remark 3.1

H1CF

For brevity in our experimental results, we sometimes abbreviate selective SSMs as S6 models, because they are S4 models with a selection mechanism and computed with a scan.

Remark

Remark E.1

GBUJ

We also note that the schedule was not tuned, and we never experimented with turning off sequence length warmup for these pretraining experiments. We later found that SLW did not help noticeably for audio pretraining at similar lengths (?), and it is possible that it is not necessary for DNA pretraining either.

Source and licence

The text on this page is from arXiv:2312.00752v2 by Albert Gu and Tri Dao, under the licence CC BY 4.0. The corpus made these changes: it read the items out of the LaTeX source, it expanded the macros of the paper, it put each proof after its statement, and it converted the text to HTML. The formulas are the TeX of the paper, and KaTeX renders them. A formula or a reference can look different from the PDF. The page is published under CC BY 4.0.

The record of each item is in analysis/objects/2312/2312.00752.jsonl.