PaperlayerView exact revision on arXiv

Source-linked AI summary

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

arXiv:2312.00752v2cs.LGcs.AI

TL;DR

Transformer가 sequence modeling을 지배하고 있지만, language 같은 주요 modality에서는 효율적인 대안들이 그 성능을 따라잡지 못했다. Mamba는 attention-free architecture에 selective state space를 도입해 Transformer 수준의 결과를 달성하면서 sequence length에 선형적으로 확장된다.

  • 문제

    기존 sequence model은 Transformer 수준의 modeling quality와 효율적인 linear scaling 사이의 격차에 직면하며, 특히 language와 그 밖의 discrete modality에서 두드러진다.

  • 방법

    Mamba는 SSM parameter를 input-dependent하게 만들어 정보를 선택적으로 retain하거나 forget할 수 있게 하며, simplified attention-free architecture에 selective state space를 결합한다.

  • 결과

    Mamba는 language, audio, genomics 전반에서 state-of-the-art 성능을 달성하고, 강력한 Transformer model과 대등하거나 이를 능가하며 context length가 1M까지 늘어날수록 성능이 향상된다.

  • 시사점 및 한계

    Selective state space는 긴 context가 필요한 여러 domain에서 Mamba를 범용 sequence-model backbone의 강력한 후보로 만든다.

  • 시사점 및 한계

    실증 평가는 small model을 사용하므로, 더 큰 parameter scale에서 Mamba가 유리하게 비교될지는 remains unclear하다.

Abstract

from arXiv · show

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5$\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

1 서론

Mamba는 입력에 의존하는 파라미터를 갖는 selective state space를 도입해 선형 시퀀스 확장성을 유지하면서 content-based filtering을 가능하게 한다. synthetic, audio, genomics, language 태스크 전반에서 효율적인 long-context 연산으로 높은 성능을 달성한다.

  • 선택 메커니즘: 입력에 의존하는 SSM 파라미터를 통해 Mamba는 무관한 정보는 선택적으로 필터링하고 관련 정보는 무기한 유지하며, 이산적이고 정보가 밀집된 데이터에서 기존 모델이 보인 약점을 해결한다.이 메커니즘은 selective-copy 및 induction-head 태스크에서 동기를 얻었으며 content-based selection을 가능하게 한다.
  • Hardware-aware 알고리즘: hardware-aware recurrent scan은 확장된 상태를 구체화하지 않고 selective SSM을 계산해 비용이 큰 GPU 메모리 계층 간 전송을 피한다.이 변경으로 time- 및 input-invariant한 기존 SSM에서 사용하던 convolution 연산을 대체한다.
  • 아키텍처: Mamba는 기존 SSM 설계와 Transformer MLP block을 selective state space를 포함하는 단일 homogeneous architecture로 결합한다.결과 모델은 완전한 recurrent 구조이며, 학습 중 연산량과 메모리가 시퀀스 길이에 대해 선형으로 확장된다.
  • 실증적 검증: synthetic copying 및 induction-head 태스크에서 Mamba는 문제를 쉽게 해결하고 1M 토큰을 넘어 해법을 extrapolate한다.이 태스크들은 large language model에 중요한 능력을 포착하기 위해 제안된 핵심 테스트로 제시된다.
  • 실증적 검증: Mamba는 audio waveform과 DNA sequence에서 기존 state-of-the-art 모델을 능가하며, million-length sequence까지 더 긴 context에서 성능이 향상된다.비교 대상에는 SaShiMi, Hyena, Transformer가 포함되며 downstream 및 pretraining 성능 향상이 나타난다. speech-generation FID는 절반 이상 감소한다.
  • 실증적 검증: Mamba는 pretraining perplexity와 downstream evaluation에서 Transformer 수준의 language 성능을 달성한 것으로 보고된 최초의 linear-time sequence model이며, 5× generation throughput을 제공한다.최대 1B 파라미터까지의 scaling 실험에서 modern LLaMA 기반 Transformer training recipe를 포함한 광범위한 baseline을 능가한다.

2 상태 공간 모델

Structured SSM은 latent state를 이용한 sequence transformation으로, 시간 불변 dynamics와 structured state matrix를 통해 효율성을 확보하며 recurrence 또는 convolution으로 계산할 수 있다. 이 절에서는 content-based modeling의 한계로 linear time invariance를 지적하고, 효율적인 계산을 유지하면서 이를 제거하는 방향을 제시한다.

  • Structured SSM: SSM은 sequence를 higher-dimensional latent state를 거쳐 변환하며, (Δ, A, B, C) 네 parameter로 transformation을 정의한다.Structured SSM은 각 input channel에 대해 model을 독립적으로 적용한다.
  • 계산: Discretization 이후 SSM은 linear-recurrence와 global-convolution 계산을 모두 지원하며, 일반적으로 parallel training에는 convolution을, autoregressive inference에는 recurrence를 사용한다.두 mode는 서로 다른 access pattern을 활용한다. training에서는 전체 sequence를, inference에서는 한 번에 하나의 timestep을 처리한다.
  • Linear Time Invariance: LTI SSM은 시간에 따라 (Δ, A, B, C)를 고정하므로 recurrence와 convolution이 동등해지지만, 특정 data type을 model링하는 능력은 제한된다.따라서 이 연구는 그에 따른 efficiency bottleneck을 다시 도입하지 않으면서 LTI constraint를 제거하는 것을 목표로 한다.
  • Structure와 Dimensions: Diagonal structure는 state matrix를 compact하게 만들지만, batch size B, length L, D channel에 대해 SSM을 적용하려면 여전히 O(BLDN) time and memory가 필요하다.Input 하나당 전체 hidden state의 dimension은 DN이므로, 이는 이후에 다루는 근본적인 efficiency bottleneck을 만든다.
  • SSM Architectures: SSM은 standalone sequence transformation으로서 end-to-end architecture에 삽입할 수 있으며, linear attention, H3, Hyena, RetNet, RWKV와 같은 관련 design에도 포함된다.이 architecture들은 SSM computation의 recurrent, convolutional, gated 또는 attention-like variant를 사용한다.

3 선택적 상태 공간 모델

Mamba는 입력에 의존하는 매개변수를 통해 콘텐츠 인식형 필터링과 시퀀스 정보 전파를 가능하게 하는 선택적 상태 공간 모델을 도입한다. 이 접근법은 convolution을 hardware-aware recurrent computation으로 대체하고, attention이나 MLP block이 없는 단순화된 architecture에 선택적 SSM을 결합한다.

  • 3.1 동기: LTI 모델은 선택적 복사와 induction-head task에서 실패한다. 일정한 dynamics로는 입력 콘텐츠에 따라 hidden state를 선택하거나 업데이트할 수 없기 때문이다.선택적 복사에는 무관한 토큰 사이에서 관련 토큰을 기억하는 능력이 필요하고, induction head에는 context-aware retrieval이 필요하다.
  • 3.1 동기: 선택적 SSM은 시퀀스 상호작용을 입력 의존적으로 만들어 모델이 관련 입력에 집중하고 이를 sequential state로 필터링할 수 있게 한다.이는 효율성과 유효성 사이의 tradeoff를 다룬다. 효율적인 모델은 context를 압축하는 반면, 유효한 모델은 필요한 정보를 유지한다.
  • 3.2 선택적 상태 공간 모델: S6 mechanism은 Δ, B, and C를 입력의 함수로 만들어 SSM을 time-invariant에서 time-varying으로 바꾸고, computation을 recurrent scan으로 제한한다.이로써 convolution과의 동치성이 사라지는 동시에 입력 의존적인 state transition이 가능해진다.
  • 3.3 Hardware-Aware Algorithms: Kernel fusion, parallel scan, and recomputation은 convolutional computation을 잃더라도 modern hardware에서 선택적 recurrent SSM을 효율적으로 만든다.순진한 recurrent form은 O(BLDN) FLOPs를 사용하고 convolution은 O(BLDlog(L)) FLOPs를 사용한다. recurrent computation은 긴 시퀀스와 중간 정도의 N에서 더 적은 FLOPs를 사용할 수 있다.
  • 3.4 Mamba: Mamba는 linear-attention-inspired component와 MLP component를 attention이나 별도의 MLP block이 없는 하나의 homogeneous stacked block으로 결합해 기존 SSM architecture를 단순화한다.선택적 모델은 더 긴 context에서 무관한 history를 제거하기 위해 state를 reset할 수도 있다.

4 실증 평가

합성 과제와 language, DNA, audio 평가 전반에서 Mamba의 selective SSM은 content-based sequence reasoning, 우수한 scaling, long-context 활용, 경쟁력 있는 generation quality를 가능하게 한다. Ablation 결과는 selective SSM이 성능 향상의 주된 원천임을 추가로 보여준다.

  • 4.1 합성 과제: Mamba의 selective SSM은 induction-heads를 완벽하게 해결하고 training보다 4000× 긴 million-length sequences로 일반화되지만, 다른 어떤 방법도 2×를 넘지 못한다.중간의 관련 없는 정보는 무시하면서 관련 token을 선택적으로 기억한다.
  • 4.2 Language Modeling: Mamba는 강력한 Transformer++ recipe와 성능을 맞춘 최초의 attention-free model이며, 특히 sequence length가 증가할수록 scaling 이점이 커진다.비교에는 Chinchilla protocol에 따라 약 125M부터 1.3B parameters까지의 model이 포함된다.
  • 4.3 DNA Modeling: 약 40M parameters에서 Mamba는 Transformer++ 및 HyenaDNA와 성능을 맞추면서 약 3× to 4× fewer parameters를 사용하고, model size가 커질수록 매끄럽게 향상된다.Mamba는 최대 1M의 context lengths에서도 이점을 얻는 반면, HyenaDNA는 sequence length가 증가할수록 성능이 저하된다.
  • 4.4 Audio Modeling: 작은 Mamba-UNet은 더 큰 GAN- 및 diffusion-based speech model을 능가하며, parameter-matched larger model은 fidelity metrics를 크게 추가 향상시킨다.SC09에서 Mamba는 outer blocks의 S4+MLP도 일관되게 능가하며, center blocks에서는 S4+MLP와 MHA+MLP보다 높은 순위를 기록한다.
  • 4.6 Ablations: Selective SSM은 static alternatives보다 성능을 크게 향상시키며, Mamba는 H3와 유사한 성능을 보이고 selective layer를 사용하면 약간 더 우수하다.Δ와 (B, C) projections를 static에서 selective로 변경하는 것이 가장 큰 이점을 제공한다.
  • 4.6 Ablations: Selective SSM state size를 늘리면 parameters가 1%만 추가되어도 perplexity가 1.0을 넘게 향상된다.더 큰 projection dimensions는 일반적으로 parameters를 소폭 늘리면서 더 완만한 이득을 제공한다.

5 논의

논의에서는 selective SSM 설계 선택이 표현력과 상태 용량을 향상시키는 한편, 모달리티 간 절충, 미해결된 downstream 능력, 더 큰 규모에서의 불확실성을 강조한다.

  • Ablations: Δ의 표현력: Selective Δ projection은 dimension 1에서도 성능을 향상시키며, 더 큰 projection을 사용하면 추가적인 성능 향상이 있지만 parameter 비용은 modest하다.SSM state size는 N = 16으로 고정된다.
  • Ablations: SSM state dimension: SSM state dimension N을 늘리면 B와 C가 selective일 때에만 무시할 수 있는 parameter 및 FLOP 비용으로 성능이 유의하게 향상된다.비교에는 constant B와 C 및 selective B와 C가 모두 포함되며, Δ projection은 64로 고정된다.
  • No Free Lunch: Continuous-Discrete Spectrum: Selection mechanism은 SSM이 text와 DNA 같은 discrete modality에서의 약점을 극복하도록 돕지만, LTI SSM이 강점을 보이는 데이터에서는 성능을 저해할 수 있다.이는 audio waveform ablation과 함께 논의된 continuous–discrete spectrum 절충을 반영한다.
  • Downstream Affordances: SSM이 fine-tuning, prompting, in-context learning, RLHF와 같은 downstream affordance에서 Transformer-based foundation model과 대등한지는 여전히 미해결된 질문이다.논의에서는 pretrained-model interaction과 관련된 adaptation, instruction tuning, quantization도 언급한다.
  • Scaling: Mamba의 평가는 small model만 사용했으므로, 더 큰 규모에서의 경쟁력, scaling behavior, engineering requirement는 여전히 미해결되었다.논의에서는 이러한 한계를 7B parameter 이상에서 Llama, RWKV, RetNet을 평가한 선행 연구와 대조한다.

6 결론

이 논문은 문맥 의존적 추론을 가능하게 하면서 시퀀스 길이에 대한 선형 스케일링을 유지하는 selective state space model을 제안한다. Attention-free architecture에서 Mamba는 다양한 도메인에 걸쳐 state-of-the-art 성능을 달성하고, 강력한 Transformer와 대등하거나 이를 능가하며 foundation model의 적용 범위를 넓힐 가능성을 제시한다.

  • 결론: Selective state space model은 시퀀스 길이에 대해 선형적으로 스케일링하면서 문맥 의존적 추론을 가능하게 한다.선택 메커니즘은 structured state space model에 대한 이 논문의 핵심 기여다.
  • 결론: Mamba의 단순한 attention-free architecture는 다양한 도메인에서 state-of-the-art 결과를 달성하며, 강력한 Transformer model과 대등하거나 이를 능가한다.이는 이 접근법이 범용 sequence modeling architecture로서 경쟁력을 갖춘다는 점을 보여준다.
  • 결론: Selective state space model은 다양한 도메인에서 foundation model을 지원할 수 있다.결론에서는 실증된 설정을 넘어 광범위하게 적용될 가능성을 강조한다.

A 논의: 선택 메커니즘 · B 관련 연구 · B.1 S4 변형 및 파생 모델

이 논문은 selection을 gating, hypernetwork, data-dependence 같은 포괄적 명칭과 구분하고, 대신 이를 RNN gating 및 input-dependent SSM discretization에 기반해 설명한다. 또한 S6를 S4에서 파생된 모델군에 배치하며, 대체로 non-selective하고 흔히 LTI인 기존 structured SSMs에 대비해 selective한 SISO recurrent formulation임을 강조한다.

  • A 논의: 선택 메커니즘: Selection은 포괄적인 gating, hypernetwork, data-dependent 명칭과 구별된다. 이러한 범주에는 단순한 multiplicative construction과 standard attention mechanism이 포함되기 때문이다.GLU와 유사한 transformation은 activation으로만 작동하면서 세 가지 명칭을 모두 만족할 수 있으므로, 이러한 명칭만으로는 selection을 특징짓기 어렵다.
  • A 논의: 선택 메커니즘: Selection은 전통적인 RNN gating과 가장 밀접하게 관련되며, Δ의 input-dependent discretization을 통해 SSMs와 연결된다 (Funahashi and Nakamura 1993; Gu, Dao, et al. 2020; Tallec and Ollivier 2018).이 논문은 gating이라는 용어가 현대적으로 여러 의미로 과도하게 사용되는 것을 피하기 위해 “selection”을 사용한다.
  • B 관련 연구: 관련 연구 개요에서는 Mamba를 recurrent layer인 S4, S5, quasi-RNNs 및 end-to-end architecture인 H3, RetNet, RWKV와 나란히 다룬다.이 모델들은 논문의 접근법과 가장 밀접하게 관련된 방법들 가운데 하나로 제시된다.
  • B.1 S4 변형 및 파생 모델: S4는 diagonal 및 diagonal-plus-low-rank structure를 갖는 structured SSMs를 도입하고, HIPPO가 continuous-time online memorization과 맺는 연관성에 동기를 얻은 효율적인 convolutional algorithm을 사용했다 (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021).이후 DSS는 diagonal structured SSMs의 경험적 효과를 확립했고, S4D는 이 결과를 이론적으로 확장했다 (Gupta, Gu, and Berant 2022; Gu, Gupta, et al. 2022).
  • B.1 S4 변형 및 파생 모델: S5는 parallel scan으로 recurrent하게 계산된 최초의 S4 model이지만, SISO에서 MIMO로 전환해 유효 state dimension을 줄였다. 반면 S6는 SISO dimension을 유지하면서 scan을 공유한다 (Smith, Warrington, and Linderman 2023).제공된 문단은 S5와 구별되는 S6의 특징으로 SISO formulation과 shared scan을 제시한다.
  • B.1 S4 변형 및 파생 모델: Lu et al. (2023)은 meta-RL에서 S5에 hard-coded reset mechanism을 사용하지만, 이 논문은 selection을 learnable하고 input-dependent한 대안으로 규정한다.이들의 mechanism은 episode boundary에서 A를 0으로 수동 설정한다. 이에 대해 이 논문은 selective SSMs가 이러한 reset을 자동으로 학습할 수 있는지 검증할 것을 제안한다.
  • B.1 S4 변형 및 파생 모델: 대부분의 다른 structured SSM 변형은 여전히 non-selective하고 대체로 strictly LTI이며, 빠른 autoregressive inference를 직접 수행할 수 없는 convolutional method도 여기에 포함된다.Liquid S4는 제한적인 input-dependent transition을 추가하면서도 LTI에 가깝게 유지되는 반면, SGConv, Hyena, LongConv, MultiresConv, Toeplitz Neural Network는 long convolution kernel에 초점을 둔다.

B.2 SSM 아키텍처 … C Selective SSM의 메커니즘

이 논문은 Selective SSM을 gated SSM, recurrent, linear-attention, long-context 아키텍처와 비교하며, 입력 의존적 선택과 원칙에 따른 SSM parameterization을 강조한다. 또한 discretization을 통해 Selective SSM의 메커니즘을 formalize하고, long-context 검증을 기존 주장과 대조한다.

  • B.2 SSM 아키텍처: Selective SSM 아키텍처는 GSS, Mega, H3, Selective S4, RetNet, RWKV, GAU와 비교되며, recurrence, attention 유사 메커니즘, gating, state expansion을 결합하는 방식에서 차이를 보인다.GSS는 model dimension을 축소하는 반면, 이 논문의 아키텍처는 SSM state size를 키우기 위해 이를 확장한다. GAU 역시 이 아키텍처에 영감을 준다.
  • B.2 SSM 아키텍처: Selective S4 (J. Wang et al. 2023)는 입력을 외부에서 mask하며, 논문은 masking이 Selective Copying에서 관련 token 사이의 간격을 does not alter the spacing 않기 때문에 이를 selection과 구분한다.따라서 논문은 Selective S4를 selection mechanism이 아닌 architectural gating으로 규정한다.
  • B.3 RNN과의 관계: RNN과 SSM은 모두 latent-state recurrence를 사용한다. 초기 RNN은 sequential-efficiency와 vanishing-gradient problems에 직면했지만, 이후 SSM theory는 parallelization과 정교한 parameterization을 통해 이를 다뤘다.Orthogonal 및 unitary RNN은 대신 transition matrix를 제약해 eigenvalue를 조절하고 vanishing gradient를 방지하지만, 다른 한계는 유지한다.
  • B.3 RNN과의 관계: strongly typed RNN, QRNN, SRU와 같은 Gated RNN은 selective SSM으로 볼 수 있지만, 성능에 중요한 state expansion과 selective B,C parameter가 부족하다.논문은 또한 이들의 heuristic gate를 SSM theory에 기반한 selection-plus-discretization formulation과 대조한다.
  • B.4 Linear Attention: Linear Attention (Katharopoulos et al. 2020)은 kernel attention과 recurrent autoregressive models를 연결하며, 다양한 kernel, normalization, positional, sampling 기반 변형을 이끈다.이 절은 efficient-attention 변형에 대한 폭넓은 survey로 Tay, Dehghani, Bahri, et al. (2022)도 제시한다.
  • B.5 Long Context Models: 논문은 synthetic task, 더 짧은 평가, 또는 data와 compute의 교란으로 제한된 기존 long-context 주장과 달리, 더 긴 context에서 성능이 향상됨을 meaningfully demonstrate improving performance 한 최초의 접근법 중 하나라고 주장한다.Recurrent Memory Transformer, LongNet, Hyena 및 HyenaDNA, Sparse Transformer 등이 예로 제시되며, 이들이 보고한 context 능력은 comparable하게 검증되지 않았다.
  • C Selective SSM의 메커니즘: Selective SSM mechanics proof는 Algorithm 2를 N=1, A=-1, B=1로 특수화하고, 입력 의존적 sΔ=Linear(x)와 τΔ=softplus를 사용한 뒤 zero-order-hold discretization을 적용해 final discrete recurrence를 얻는다.이 discretization은 parameter를 linear projection에 fold된 learnable bias로 해석하는 것도 가능하게 한다.

D Selective SSM을 위한 하드웨어 인지 알고리즘 … E.2 언어 모델링

Mamba의 selective SSM은 하드웨어 인지 fused scan을 사용해 선형 스케일링을 유지하면서 memory traffic과 실행 시간을 크게 줄이며, recomputation으로 최적화된 Transformer와 견줄 만한 activation memory를 유지한다. 제시된 실험 세부사항은 selective copying과 induction-head 평가를 위한 synthetic-task 설정과 학습 프로토콜을 설명한다.

  • D Selective SSM을 위한 하드웨어 인지 알고리즘: Selective SSM scan은 sequence length에 대해 여전히 선형이며 이론적 비용은 O(BLDN)이지만, foundation-model 학습을 위해서는 하드웨어 인지 구현이 필요하다.Input-dependent selectivity는 nonselective SSM에서 사용할 수 있는 convolution 및 FFT 구현을 제거하므로, 이 방법은 parallel associative scan을 사용한다.
  • D Selective SSM을 위한 하드웨어 인지 알고리즘: Discretization, associative scan, output multiplication을 융합하면 20-40×의 speedup이 발생하며, memory I/O를 O(N)만큼 줄인다.Fused kernel은 중간 계산을 SRAM에 유지하고 최종 output만 HBM에 기록한다.
  • D Selective SSM을 위한 하드웨어 인지 알고리즘: SRAM에 들어가지 않을 만큼 긴 sequence에서는 intermediate scan state를 사용해 chunked fused scan을 chunk 간에 계속 수행한다.이는 scan recurrence를 유지하면서 fused 구현을 SRAM 용량 이상으로 확장한다.
  • D Selective SSM을 위한 하드웨어 인지 알고리즘: Selective SSM은 token당 16 bytes의 activation memory를 사용하므로, selective-SSM layer 두 개는 token당 32 bytes를 사용하는 FlashAttention layer 하나와 MLP layer 하나에 필적한다.이 비교는 mixed-precision FP16 또는 BF16 학습을 가정한다.
  • E.1 Synthetic Tasks: Synthetic-task 실험은 length-4096 sequence, 16-token vocabulary, 16개의 memorized data token, 두 개의 layer, model dimension D=64를 사용한다.Vocabulary에는 Figure 2의 white noise token이 포함된다.
  • Selective Copying.: Selective-copying 평가는 무작위로 생성한 training data, 고정된 무작위 생성 validation set, 그리고 10 to 50 epochs 범위의 model별 reporting checkpoint를 사용한다.Mamba와 MHA-Abs 결과는 25 epochs 후, MHA-RoPE와 MHA-xPos는 50 epochs 후, LTI H3와 Hyena는 10 epochs 후에 보고한다.
  • Selective Copying.: Selective-copying 학습은 weight decay 없는 Adam과 constant learning rate 2e−4 및 1e−3을 사용하며, 각 model의 더 나은 결과를 보고한다.Attention과 Hyena는 1e−3에서 학습되지 않았고, H3는 2e−4에서 더 잘 generalize했으며 Mamba는 1e−3에서 더 잘 extrapolate했다.

E.2.1 Scaling Law 세부 사항

Scaling-law 실험에서는 GPT2 tokenizer를 사용해 Pile에서 모델을 학습했으며, batch size와 token budget을 조정한 것을 제외하면 대체로 GPT3 recipe를 따랐다. Transformer, Hyena, H3++, RWKV, RetNet, Mamba의 standard 및 improved-recipe 버전을 비교했다.

  • Scaling-law 실험에서는 GPT2 tokenizer를 사용해 Pile에서 모든 모델을 학습했으며, 대체로 GPT3 recipe를 따랐다.
  • Model Size: Scaling-law 모델 크기와 hyperparameter는 1.3B batch size와 training budget을 소폭 수정한 GPT3 specification을 따랐다.1.3B batch size는 1M token에서 0.5M token으로 변경했으며, training step과 total token은 Chinchilla scaling law에 맞춰 조정했다.
  • Training Recipe: 기본 training recipe는 AdamW, 1.0에서의 gradient clipping, cosine decay를 적용한 linear warmup, GPT3 peak learning rate를 사용했다.
  • Training Recipe: 개선된 recipe는 peak learning rate를 GPT3 값의 5×로 높이고, 1e−5까지 decay했으며, linear bias를 제거하고 RMSNorm을 사용하고 AdamW β를 (.9, .95)로 변경했다.이러한 변경은 PaLM (Chowdhery et al. 2023)과 LLaMa (Touvron et al. 2023)에서 사용한 recipe에서 영감을 얻었다.
  • Architecture and Training Detail: 비교에는 GPT3-style Transformer, improved Transformer++, Hyena, H3++, RWKV, RetNet, standard Mamba architecture가 포함됐다.Transformer++는 rotary positional encoding과 SwiGLU를 사용했으며, Hyena와 H3++는 명시된 architectural modification을 사용했고, Mamba는 improved recipe를 사용했다.

E.2.2 추가 Scaling Law Ablation

추가 ablation 결과, Mamba를 MLP 또는 attention block과 interleaving해도 성능 변화는 작지만, training recipe와 확장된 SSM state dimension은 더 큰 영향을 미친다. Inner LTI SSM의 선택은 영향이 거의 없다.

  • Mamba Architecture: Block Interleaving: Mamba-MLP는 homogeneous Mamba보다 조금 낮은 성능을 보이지만 Transformer++를 제외한 모든 모델보다 여전히 우수하며, Mamba-MHA는 조금 더 우수하다.이 ablation에서는 Mamba block을 homogeneous하게 쌓는 대신 standard MLP 또는 multi-head attention block과 Mamba block을 interleaving한다.
  • H3 Architecture: Training Recipe: 개선된 training recipe에서 큰 향상이 나타나는 반면, inner LTI SSM을 교체해도 영향이 거의 없으며 head dimension을 확장하면 성능이 향상된다.이 결과는 SSM 성능에서 training choice와 확장된 state dimension의 중요성을 뒷받침하며, Hyena와 S4의 선택은 상대적으로 중요하지 않음을 보여준다.

E.2.3 다운스트림 평가 세부사항 … E.4 오디오 세부사항

부록에서는 언어, 유전체학, 장기 컨텍스트 분류를 위한 평가 프로토콜을 자세히 설명하며, 데이터셋 구성, 모델 스케일링, 최적화, 시퀀스 길이 스케줄을 다룬다. 또한 각 컨텍스트 길이에 맞춘 pretrained model을 사용해 Great Apes DNA 분류 결과를 보고한다.

  • E.2.3 다운스트림 평가 세부사항: 다운스트림 언어 평가는 LAMBADA, HellaSwag, PIQA, ARC-challenge, ARC-easy, WinoGrande에서 EleutherAI’s LM evaluation harness를 사용한다.LAMBADA, WinoGrande, PIQA, ARC-easy에는 accuracy를 보고하고, HellaSwag과 ARC-challenge에는 sequence length로 정규화한 accuracy를 사용한다.
  • E.3.1 Pretraining 세부사항: HG38 training set은 Enformer의 genome split을 따르며, 2^17 = 131072 길이의 34021개 segment와 approximately 4.5 billion DNA tokens를 포함한다.Segment는 chromosome number와 start 및 end index로 표현하며, 더 긴 sequence를 위해 확장할 수 있다.
  • E.3.1 Pretraining 세부사항: Context length L <= 2^17에서는 각 genome segment를 겹치지 않는 length-L subsegment로 나누는 반면, 더 긴 context에서는 지정된 경계에 맞춰 segment당 두 개의 sample을 생성한다.길이 2^18에서는 default보다 4× 많은 token을 사용하고, 길이 2^20에서는 16× 많은 token을 사용한다.
  • E.3.2 스케일링: 모델 크기 세부사항: 스케일링에서는 Transformer++, HyenaDNA, standard Mamba를 비교하며, attention과 MLP block을 모두 포함하는 Transformer layer와 parameter count를 맞추기 위해 Mamba의 block 수를 두 배로 늘린다.Transformer++는 RoPE positional encoding을 사용하고, HyenaDNA는 multi-head attention을 MLP로 parameterize한 H3 global-convolution block으로 대체한다.
  • E.3.2 스케일링: 모델 크기 세부사항: Learning-rate sweep에서는 Transformer와 HyenaDNA에 2e-3, Mamba에 8e-3를 선택했으며, Mamba는 더 안정적이고 높은 rate에서 추가로 향상되었다.Sweep에는 {1e-3, 2e-3, 4e-3, 8e-3}를 사용했으며, Mamba는 일치하는 2e-3 rate에서도 baseline을 능가했다.
  • E.3.3 스케일링: 컨텍스트 길이 세부사항: Context-length scaling에서는 2^10 = 1024부터 각 2의 거듭제곱 길이에서 두 epoch 동안 sequence-length warmup을 사용하고, sequence length가 두 배가 될 때마다 batch size를 절반으로 줄인다.이 schedule은 튜닝하지 않았으며, 이를 비활성화해 DNA pretraining에서의 이점을 검증하지도 않았다.
  • E.3.4 Species (Great Apes) 분류: Species classification에서는 prediction에 final causal output만 사용하고, pretraining 중 batch_size × sequence_length를 일정하게 유지하며, 2^14 = 16384를 초과하는 context에는 추가 warmup을 적용한다.2^20 = 1048576-context model은 6회의 warmup epoch 후 maximum length에서 4 epoch를 사용한다.
  • E.3.4 Species (Great Apes) 분류: Great Apes DNA classification은 일치하는 context length를 가진 pretrained model을 사용해 2^10 = 1024부터 2^20 = 1048576까지의 context length에서 평가한다.Table 13에는 accuracy를 보고하며, random guessing의 accuracy는 20%다.

E.4.1 YouTubeMix 오디오 사전학습 · E.4.2 SC09 음성 생성

YouTubeMix 실험은 raw signal에 가까운 위치에서는 audio waveform이 LTI processing을 선호하는 반면, tokenization과 compression 이후에는 selection의 해로움이 줄어듦을 보여준다. SC09 training은 autoregressive settings를 사용했지만, 작은 dataset으로 인해 가장 큰 model에서 likelihood metrics의 overfitting이 발생했다.

  • E.4.1 YouTubeMix 오디오 사전학습: YouTubeMix models는 총 15개의 Mamba blocks, pooling factor p=16, outer dimension D=64, 약 3.5M개의 parameters를 사용했다.Dataset은 vocabulary가 256개의 discrete tokens인 8-bit mu-law encoding을 사용했다.
  • E.4.1 YouTubeMix 오디오 사전학습: YouTubeMix clips는 최대 1분 또는 length 960000까지였으며, 가장 긴 hardware-efficient sequence는 468 × 2048 = 958464로 설정했다.더 긴 clips는 두 번의 pooling stages 이후 subsampling하고 원하는 sequence lengths로 나누었다.
  • E.4.1 YouTubeMix 오디오 사전학습: Audio waveforms는 smooth하고 uniformly sampled된 continuous structure와 일관되게 selection-based Mamba-S6보다 Mamba-S4를 유의하게 선호한다.Figure 10은 LTI models가 이에 부합하는 inductive bias를 제공하기 때문에 selection을 제거하면 long-form audio waveform pretraining이 향상된다고 보고한다.
  • E.4.1 YouTubeMix 오디오 사전학습: Inner U-Net layers에서만 selection을 ablate하면 performance differences가 크게 줄어들며, 이는 tokenization 이후의 layers가 LTI로 유지될 필요가 없음을 시사한다.Outer layers는 Mamba-S4로 유지하고, inner layers는 signal이 tokenized되고 compressed된 이후에 작동한다.
  • E.4.1 YouTubeMix 오디오 사전학습: Inner layers에서 selection을 제거해도 real-valued SSM은 complex-valued SSM보다 여전히 성능이 낮다.이 비교는 outer U-Net layers가 Mamba-S4로 유지되는 setting에서 보고되었다.
  • E.4.2 SC09 음성 생성: SC09 autoregressive training은 0.002 learning rate, 200000 steps, batch size 16, 10% warmup, gradient clipping at 0.1을 사용했다.이 settings는 대체로 autoregressive language-modeling protocol을 따랐다.
  • E.4.2 SC09 음성 생성: Large SC09 Mamba model은 stage당 15 layers, D=96, pooling factor 4를 사용했지만, 작은 dataset으로 인해 substantial BPB or NLL overfitting이 발생했다.Generated samples의 automated metrics는 likelihood overfitting에도 불구하고 training 전반에 걸쳐 계속 향상되었다.
  • E.4.2 SC09 음성 생성: SC09 architecture ablations는 S4+MLP, Transformer, Mamba blocks 전반에 걸쳐 stage당 8 layers, D=64, pooling factor 4, 대략 6M개의 parameters를 사용했다.S4+MLP, Transformer, Mamba blocks는 comparable total model sizes가 되도록 parameterized했다.

E.5 효율성 벤치마크

효율성 벤치마크는 convolution 및 attention과 비교해 selective SSM scan을 평가하고, end-to-end inference와 memory 측정도 함께 수행한다. Mamba의 memory footprint는 최적화 수준이 가장 높은 Transformer와 비슷하며, scan 구현은 discretization과 scanning을 융합해 HBM에 대규모 parameter를 materialize하지 않는다.

  • Scan 연산: 핵심 연산 벤치마크는 각 핵심 연산 외부의 비용을 제외하고, A100 80GB PCIe GPU에서 selective SSM scan을 convolution 및 attention과 비교해 측정한다.Convolution은 PyTorch FFT 구현을 사용하고, attention은 causal FlashAttention-2 (Dao 2024)를 사용한다.
  • Scan 연산: Fused selective SSM scan은 discretization과 parallel scanning을 결합해 HBM에 대규모 A, B, C parameter를 materialize하지 않는다.Unfused PyTorch baseline은 이 parameter들을 HBM에 materialize하는 반면, 제안 구현은 두 연산을 융합한다.
  • Memory 벤치마크: Mamba의 memory footprint는 길이 2048인 sequence로 학습한 125M model에서 최적화 수준이 가장 높은 Transformer와 비슷하다.이 벤치마크는 A100 80GB GPU 한 대에서의 training memory를 torch.compile kernel fusion과 FlashAttention-2를 사용하는 Transformer와 비교한다.
  • End-to-end 추론: End-to-end inference에서는 Mamba 1.4B 및 6.9B model을 batch size 1부터 128까지의 1.3B 및 6.7B GPT-3-style Transformer와 비교한다.측정은 A100 GPU에서 2048-token prompt로부터 128 token을 생성하며, 세 번 실행한 결과를 평균한다.
Loading 2312.00752v2…