Mamba: 효율적인 시퀀스를 위한 선택적 상태 공간 모델

선택적 상태 공간을 사용하는 모델인 맘바는 구조화된 상태 공간으로 긴 시퀀스를 효율적으로 모델링하여 다른 시퀀스 모델 기반 기준선보다 뛰어난 성능을 발휘합니다.

### Mamba: 선택적 상태 공간 및 하드웨어 인식 알고리즘을 사용한 효율적인 시퀀스 모델링
* 🐍  맘바는 효율적인 장거리 시퀀스 모델링을 위해 **선택적 상태 공간과 하드웨어 인식 알고리즘**을 도입한 새로운 시퀀스 모델링 아키텍처입니다.
* 📊  상태 공간 모델에 대한 이전 작업을 기반으로 하며 **선택 메커니즘, 하드웨어 인식 알고리즘, 단순화된 아키텍처 설계**에 기여합니다.
* 🔍  이 연구의 동기는 시퀀스 모델링에서 효과와 효율성 사이의 균형을 탐구할 필요성에서 비롯되었으며, 합성 작업, 언어 모델링,  평가 대상: **DNA 모델링 및 오디오 모델링**
* ⚙️   Mamba는 **난해성, 확장성, 추론 처리량** 측면에서 뛰어난 성능을 보여주며 Transformers 및 이전 상태 공간 모델과 같은 기준선을 뛰어넘습니다.
* 🦠  그러나 텍스트 및  DNA와 같은 개별 양식에는 제한이 있을 수 있으며 **더 큰 모델 크기에서의 성능은 아직 평가되지 않았습니다.**

A few days ago, Mamba, linear time sequence modeling with Selective State spaces, appeared on archive. It is causing some excitement, partly because it shows strong empirical results, partly because it has an interesting design, and partly because Mamba is a great name. In this video, I'll give a high-level summary of the paper. To understand the motivation for this work, it's helpful to go back in time to the Distant Memories of 2020, and in particular, the long-range Arena, a benchmark for efficient Transformers. The starting premise was that Transformers do not scale very well to long sequence lengths, largely because of quadratic self-attention complexity. Given the proliferation of efficient alternatives to Transformers that were emerging at the time, this work proposed a benchmark specifically focused on evaluating model quality under long context scenarios. There were tasks like long list Ops, where the model is given an input consisting of a long list of nested operators and needs to calculate an answer, and fun tasks like pass Pathfinder. Where the model gets a flattened image that looks like this, and

needs to determine whether the two circles can be joined. This is a positive example because they can. This is a negative example because they cannot. Broadly, they found that if you wanted to maintain performance encapsulated here in this L score, it's quite hard to do much better than the vanilla Transformer independently. And a little earlier, a creative approach to tackling the long-range dependency problem was introduced in this work on leand memory units. In the context of recurrent neural networks, as a refresher, Leon pooms are these beautiful creatures. You can construct them by making a sequence of polynomials that are orthogonal to each other. By projecting values onto a basis of these pooms, you can efficiently store a history of the inputs to your system. That's the key idea. In this funky looking equation, you can reconstruct the input U by reassembling it from the components of the memory vector denoted here by m. This idea can then be baked into a layer where the memory block is implemented with simple linear operations. Building on that idea and

others hypothesized that recurrent memory with optimal polynomial projections had the insight that we can phrase memory as a technical problem of online function approximation. Where a function is summarized by storing its optimal coefficients in terms of some basis functions. Using this insight, the authors constructed a simple model based on AJ pooms that worked well on tasks like permuted mnist, outperforming other sequence model-based baselines like LSTM and dilated RNN. This was followed by the work combining recurrent convolutional and continuous-time models with linear state-space layers, which proposed a unifying framework for sequence modeling based on a standard state-space representation. This work showed how starting from an implicit continuous-time state-space model, you could discretize to produce a system that can be interpreted as a recurrent network with the benefit of efficient inference, or as a convolutional model which benefits from parallelizable training. Unfortunately, a naive implementation of these models requires a massive amount of memory. In fact, for reasonably sized models, such as when the state dimension is 256, the linear state-space layer uses

orders of magnitude more memory than comparably sized RNNs or CNN. That brings us to the S4 model introduced in efficiently modeling long sequences with structured state spaces. What we'd really like to be able to do is use the convolutional interpretation we've just seen to efficiently train our state space model. Now, if we want to train state space models using the convolutional representation, we can do a little bit of unrolling and a little bit of vectorization to get this L-dimensional kernel, where L is the sequence length. The scary part here is this a raised to the power of I term, marking back to your linear algebra classes. As soon as you see this, you may feel a strong and appropriate urge to try to diagonalize A. The authors discuss diagonalization; unfortunately, it cannot be used directly for the hippo matrix due to numerical issues, specifically, it has entries exponentially large in these state size. However, we can address the hippo matrix with an adjustment. The authors note that although the hippo matrix is not normal, it can be

decomposed as the sum of a norm normal and low-rank matrix. Even this is still not useful by itself, and additional three new techniques are needed. First, instead of computing the kernel Kar directly, they instead compute its spectrum, then determine Kar by applying an inverse FFT. Next, we use everyone's favorite low-rank trick, the Woodbury identity, that lets us perform an efficient matrix update. Third, it is shown that the diagonal matrix case is equivalent to the computation of a Koshi kernel, a well-studied problem with stable near-linear algorithms. All of this comes together in Theorem 3. It says that given any step size Delta, computing the state-space model convolution filter Kar can be reduced to four Koshi multiplies, requiring only Big O of n plus L operations, ignoring logarithmic factors, and Big O of n plus L space, where again L is the sequence length and N is the state size. Empirically, the proposed S4 architecture does very well on the long-range arena, even doing well on the path X task where the model must

determine if two points are connected on a fairly large image. This task is not too hard when the input is arranged as an image, but it's pretty difficult when it's a flattened sequence, as evidenced by the fact that many other models fail on this task. If you'd like to gain some intuition for how S4 works, I highly recommend the annotated S4 by Sasha Rush and Sid Caram Chetti. Just as with the legendary annotated Transformer blog post back in 2018, the annotated S4 takes you step by step through the paper and provides snippets of code that implement the mathematics. It also includes nice visualizations to help you build intuition for how State space models work. All of this brings us back to remember this work. Builds on the state space model approach with a few key contributions. First, a selection mechanism. Since previous state space models lack the ability to efficiently select data in an input-dependent manner, they design a simple selection mechanism by parameterizing the state space model parameters based on the input.

unfortunately, this change makes it hard for us to use the efficient convolutional trick we saw in S4. So, the second contribution is a hardware-aware algorithm that computes the model recurrently with a scan instead of convolution. This leads to an implementation that is faster than previous methods, both in theory scaling linearly in sequence length compared to pseudo-linear for all convolution-based SSMS, and on Modern Hardware up to three times faster on A100 GPUs. Spicy! The third contribution is to simplify prior deep sequence model architectures by combining the design of Prior State space model architectures with the MLP block of Transformers into a single block, leading to a simple and homogeneous architecture design called Mamba that incorporates selective State Spaces by building on a state space model Foundation. Mamba is related to many previous state space model architectures. In addition to the models we met earlier, there is linear retention, which can be viewed as a degenerate linear SSM, the gloriously named Hungry Hungry Hippos or H3, Hyena RNet, and RW KV.

the core motivation behind Mamba is to use selection as a means of compression. The authors point out that you can view various sequence models as making different trade-offs when tackling a fundamental problem of sequence modeling, i.e., compressing context into a smaller state. At one extreme, we have attention, which is both effective and inefficient because it explicitly does not compress context at all. At the other extreme, Recurrent models are efficient because they have a finite state, but their effectiveness is limited by how well this state has compressed the context. To explore this trade-off, the authors focus on two synthetic tasks. One previously well-studied task is copying. The job of the model is to copy color sequences in the input into the output, delayed by some constant offset. Standard convolutional models with fixed kernels can solve this without any problem. The first of the harder tasks they consider is selective copying. This time, we have random spaces symbolized by the white blocks in between the input sequence elements. To successfully copy the color outputs to the output while ignoring the white blocks, we need a mechanism that

behaves differently at different inputs. The second task they look at is induction heads. Here, in order to predict the appropriate color at the question mark block, the model needs to be able to perform retrieval back over the input sequence based on the context provided by the black square. In order to predict that blue is the next color in the sequence, to build a model capable of solving these tasks, the model starts from the S4 State space model, which uses parameters a, b, c, and Delta, which don't depend on the input sequence. Now, for the proposed algorithm to allow the model to behave differently on different inputs, b, c, and Delta are altered to be time-varying. They now depend on the input sequence. This gives the model more power but means we can no longer use the efficient S4 convolution trick. To get around this, the authors develop a selective scan based on Hardware-aware State expansion. In effect, the two big challenges to be tackled once we give up on convolution are the sequential nature of recurrence and the large memory usage. There are three techniques that come

into play kernel Fusion parallel scan and recomputation ation a key idea for getting around the sequential processing problem comes from simplified State space layers for sequence modeling or S5. which highlighted the benefits of using parallel scans to maintain efficiency while avoiding the use of convolutional tricks used in S4. Mamba uses this parallel scan idea in combination with efficient use of memory. in particular, instead of preparing the scan input in GPU high bandwidth memory. they load the state space model parameters Direct from the slow high bandwidth memory to fast SRAM where they perform the discretization and recurrence. then they write the final results back to high bandwidth memory. last recomputation is used to reduce the memory requirements. this allows them to avoid saving the intermediate states which are necessary for back propagation. putting this all together here is the selective State space model with Hardware aware State expansion. the first thing we see is that BT delt T and CT all depend on the input XT via the selection mechanism. the second thing we see is that the parameters are loaded into fast GPU SRAM.

indicated in this diagram by the color orange, where the update is performed. Then, the output is ultimately stored back into TPU high-bandwidth memory, shown in green. Next, we come to the simplified State space model architecture. Here, the authors combine the Hungry Hungry Hippos block with a gated MLP block to produce their MBER block. The shapes here indicate that the dimensionality is expanded inside the block. This block is repeated and interleaved with standard normalization and residual connections to form the Mamba architecture. I'll mention a few other model details. The authors note that most prior State space models use complex numbers in their state, but it has been empirically observed that completely real-valued State space models seem to work fine, and possibly even better in some settings. So, they use real values as the default, which work well for all but one of their tasks. Next, we come to the empirical evaluation. First, we have the synthetic tasks described earlier. First, as a bit of notation, the authors abbreviate selective State space models as S6 models because they are S4 models with a

selection mechanism and computed with a scan. That's a lot of s's on the selective copying task. S6 works well with every architecture, but S4 and hyena work comparatively poorly on the induction heads task. M, shown in brown, appears up here at the top, where we see it succeeds on test sequence lengths of up to a million tokens, which is 4,000 times longer than it saw during training, while none of the other methods compared to generalize to Beyond twice their training length. Next, we have experiments on language modeling, which follow the training recipe described in the GBD3 work and train on the pile data set. Here, the metric on the y-axis is perplexity, so lower is better. We can see that across a sequence length of 248 and a sequence length of 8,192, that Mamba, shown in purple, outperforms the other baselines and matches Transformer Plus+, which is highlighted in Orange and represents the strongest Transformer recipe that the authors know of based on the palm and llama architectures. That means rotary.

embeddings swiglo MLPs RMS Norm instead of layer Norm no linear bias and Higher Learning rates on zero shot Downstream. If valuations there are lots of comparisons generally we see Mamba achieving an average accuracy of baselines at twice the model size. Next, there are studies on DNA modeling on human genome data. Mamba shown in Orange scales better than hyena DNA shown in blue and Transformer Plus+ shown in red. When you scale up the number of parameters, it also scales better than hyena DNA. When you scale up the sequence length, member models also do well relative to a hyena DNA Baseline when fine-tuning for a species DNA classification task. Here the metric on the Y AIS is accuracy so higher is better with longer sequences. As we move to the right along the x-axis, we see the orange and green mber curves Rising. I won't go through them in detail, but there are other experiments showing that mber mostly works well on audio modeling and generation. However, there is one ablation in the appendix identifying a

case when a naive application of Mamba with the selection mechanism seems to HT performance here lower is better. And the Orange Line denotes the default M configuration. The authors suggest that this may be a case where audio waveforms actually benefit from linear time-invariant models, which have a matching inductive bias. Next, we have speed and memory benchmarks. The authors provide an implementation of the scan operation that is pretty fast. They compare the scan versus convolution and attention on an a100 GPU. Here the Y AIS measures time, so lower is better. We see the red line representing their implementation is considerably below the other mechanisms at longer sequence lengths. In terms of inference throughput, where higher is better, we find that Mamba shown in blue and green can achieve five times higher throughput than Transformers, benefiting from its recurrent nature. In terms of limitations, the authors highlight that there is no free lunch on the continuous-discrete Spectrum. The selection mechanism overcomes the weaknesses of Prior State space models on discrete modalities such as text and DNA, but this can impede their performance on data.

that linear time-invariant state-space models excel on which we saw with the audio oblation they also note that the empirical evaluation is limited to small model sizes below the threshold of most strong open-source llms therefore it remains to assess whether Mamba still compares favorably at larger sizes that's it we've reached the end if you'd like to try it out the authors have released code on GitHub

Mamba - a replacement for Transformers?

Video descriptions/tutorials for:
- AI/machine learning research
- Algorithms/data structures

(Optional) ways to support the channel: 
- https://patreon.com/samuel_albanie
- https://www.buymeacoffee.com/samuelalbanie

Other related content:
- Twitter: https://twitter.com/SamuelAlbanie
- Research lab: https://caml-lab.com/

FAQ:
Software used to make videos
- keynote on mac (to make slides)
- Adobe Premiere Pro for editing


Samuel Albanie

시퀀스 모델링에서 맘바의 흥미로운 등장

맘바는 강력한 경험적 결과와 흥미로운 설계를 갖춘 새로운 시퀀스 모델링 접근법으로, 연구 커뮤니티에서 큰 관심을 불러일으키고 있습니다.

이 논문에서는 긴 컨텍스트 시나리오에서 모델 성능을 평가하기 위한 벤치마크를 소개하며, 긴 목록 작업 및 패스 패스 파인더와 같은 작업을 예로 들어 설명합니다.

또한 표준 상태 공간 표현에 기반한 시퀀스 모델링을 위한 통합 프레임워크를 제안하여 효율적인 추론과 병렬화 가능한 훈련의 이점을 보여줍니다.

S4 모델: 구조화된 상태 공간으로 긴 시퀀스를 효율적으로 모델링하기

S4 모델은 긴 시퀀스를 모델링할 때 RNN과 CNN보다 메모리 효율이 높은 대안을 제공합니다.

이 모델은 컨볼루션 해석, 언롤링, 벡터화를 통해 계산을 몇 개의 코시 배수로 줄입니다.

이 모델의 성공은 경험적으로 입증되었으며, 장거리 작업에서 우수한 성능을 발휘하고 까다로운 평평한 수열 문제에서도 탁월한 성능을 발휘합니다.

S4의 접근 방식은 주석이 달린 S4에서 자세히 설명되고 시각화되어 있으며, 모델의 수학과 구현에 대한 단계별 가이드를 제공합니다.

하드웨어 인식 알고리즘으로 시퀀스 모델링 개선하기

새로운 하드웨어 인식 알고리즘은 컨볼루션 대신 스캔을 사용하여 모델을 더 빠르게 계산하는 방법을 제공하여 시퀀스 길이를 선형적으로 확장하고 최신 GPU에서 최대 3배 빠른 성능을 제공합니다.

Mamba 아키텍처는 선행 상태 공간 모델 아키텍처의 설계와 Transformers의 MLP 블록을 결합하여 컨텍스트를 더 작은 상태로 압축하고 시퀀스 모델링에서 절충점을 찾는 것을 목표로 심층 시퀀스 모델을 간소화합니다.

이 모델은 효율적인 S4 컨볼루션 트릭을 포기하고 하드웨어 인식 상태 확장에 기반한 선택적 스캔을 사용하여 시간에 따라 변하는 매개변수를 변경함으로써 선택적 복사 및 유도 헤드와 같은 까다로운 작업을 해결합니다.

맘바 모델: 상태 공간 레이어에 대한 하드웨어 인식 접근 방식

맘바 모델은 병렬 스캔과 효율적인 메모리 사용을 통해 처리 효율성을 개선하고 메모리 요구 사항을 줄입니다.

이 모델은 하드웨어 인식 상태 확장 방식을 사용하여 업데이트를 위해 매개변수를 빠른 GPU SRAM에 로드하고 출력을 TPU 고대역폭 메모리에 저장합니다.

이 모델은 또한 대부분의 작업에서 효율성을 입증한 완전한 실수값 상태 공간을 활용하며, Hungry Hungry Hippos 블록과 게이트 MLP 블록을 결합하여 아키텍처에 MBER 블록을 생성합니다.

다양한 작업에서 맘바의 성능 비교

맘바는 선택적 복사 및 유도 헤드 작업에서 다른 방법보다 성능이 뛰어납니다.

언어 모델링에서 Mamba는 다양한 서열 길이에 걸쳐 더 낮은 난해성을 보이며, 이는 Transformer Plus+의 성능과 일치합니다.

또한, DNA 모델링에서는 매개변수와 서열 길이를 늘릴 때 Mamba가 하이에나 DNA 및 Transformer Plus+보다 더 잘 확장되어 종 DNA 분류에서 더 높은 정확도를 보여줍니다.

맘바의 선택 메커니즘이 성능에 미치는 영향

Mamba의 선택 메커니즘은 선형 시간 불변 모델 덕분에 오디오 파형에서 더 나은 성능을 보여줍니다.

스캔 작업의 구현은 a100 GPU의 컨볼루션 및 주의에 비해 현저히 빠릅니다.

맘바는 트랜스포머보다 최대 5배 높은 추론 처리량을 보여주지만, 데이터에 대한 성능은 이전 상태 공간 모델의 약점으로 인해 저해될 수 있습니다.