State-space models (Mamba, S4) — Hinglish notes

6.1.11 · AI-ML › Scaling & Efficient Architectures

Intuition Core Idea

State-space models (SSMs) sequence modeling ke liye Transformers ka ek alternative hain jo quadratic ki jagah linear-time complexity achieve karte hain. Ye sequences ko continuous dynamical systems ki tarah treat karte hain jo time ke saath discretize hote hain, aur ek hidden state maintain karte hain jo learned dynamics ke according evolve karta hai. Ise ek signal processing mein filter ki tarah socho jo important patterns ko selectively yaad rakhta hai aur noise ko bhool jaata hai, lekin end-to-end seekha jaata hai.

YE KYUN IMPORTANT HAI: Transformers $O (n^{2})$ scale karte hain sequence length mein kyunki full attention matrices banti hain. SSMs $O (n)$ scale karte hain, jo unhe bahut lambi sequences (100K+ tokens) ke liye viable banata hai jo attention ke saath GPU memory exhaust kar deti hain.

The State-Space Foundation

CONTINUOUS KYUN? Bhalein hum discrete tokens process karte hain, continuous formulation hamein control theory aur signal processing ke powerful tools deta hai. Hum ise carefully discretize karenge.

From Continuous to Discrete: The Key Step

Neural networks ke liye jo discrete sequences process karti hain, hume continuous dynamics ko discretize karna hoga. Zero-order hold assumption ka use karke (input timesteps ke beech constant hota hai), exact solution hai:

YE SPECIFIC DISCRETIZATION KYUN? Matrix exponential $e^{A Δ}$ constant input ke liye ODE ka exact solution hai. Simple methods use karne se (jaise Euler) lambi sequences par errors accumulate ho jaati hain.

$e^{A Δ}$ KAISE COMPUTE KARO? General matrices ke liye ye costly hai, lekin S4 structured matrices use karta hai (agla section) jahan ye efficient ho jaata hai.

S4: Structured State Spaces

Original SSM approach ek wall se takra gayi: state dimension $N$ bada (thousands) hona chahiye taaki sufficient memory capacity ho, lekin isse $A$ matrix huge aur expensive ho jaati hai.

The Convolution View

Yahan ek crucial dual perspective hai: SSM ko ya to dekha ja sakta hai:

Recurrent: $h_{k} = \overset{ˉ}{A} h_{k - 1} + \overset{ˉ}{B} x_{k}$ (generation ke liye)
Convolutional: $y = \overset{ˉ}{K} * x$ (training ke liye)

SSM convolution kernel hai:

CONVOLUTIONAL VIEW KYUN? Convolutions FFT use karke $O (n lo g n)$ mein compute ho sakti hain, aur ye sequence ke across perfectly parallelize hoti hain (recurrent form ke unlike). Training convolution use karti hai, inference recurrence use karta hai.

DPLR STRUCTURE KAISE HELP KARTA HAI? Diagonal-plus-low-rank $A$ ke saath, powers $\overset{ˉ}{A}^{k}$ aur kernel $\overset{ˉ}{K}$ Cauchy kernel trick aur FFT ka use karke efficiently compute ho sakte hain. Details mein complex analysis hai, lekin result hai $O (N lo g N)$ computation per layer.

Worked example Worked Example: 2D SSM Step

Intuition ke liye ek tiny 2D state ke saath ek timestep trace karte hain.

Diya gaya:

State dimension $N = 2$
Discrete matrices: $\overset{ˉ}{A} = [0.9 0 0 0.5]$ , $\overset{ˉ}{B} = [0.1 0.2]$ , $C = [1 - 0.5]$
Previous state: $h_{0} = [0.3 0.1]$
Current input: $x_{1} = 2.0$

Step 1: State update karo $h_{1} = \overset{ˉ}{A} h_{0} + \overset{ˉ}{B} x_{1} = [0.9 0 0 0.5] [0.3 0.1] + [0.1 0.2] \cdot 2.0$ $\overset{ˉ}{A}$ SE MULTIPLY KYUN? Hum old state ko decay kar rahe hain—0.9 ka matlab hai pehli dimension mein lambi memory hai (dheere decay hoti hai), jabki 0.5 ka matlab hai doosri dimension jaldi bhool jaati hai.

$= [0.27 0.05] + [0.2 0.4] = [0.47 0.45]$

Step 2: Output compute karo $y_{1} = C h_{1} = [1 - 0.5] [0.47 0.45] = 0.47 - 0.225 = 0.245$

NEGATIVE WEIGHT KYUN? Dono state dimensions ke opposing effects ho sakte hain, jisse model complex patterns seekh sakta hai (jaise, "output high ho jab pehli dimension high ho LEKIN doosri low ho").

Mamba: Selective State Spaces

S4 ek breakthrough tha lekin iska ek key limitation tha: matrices $\overset{ˉ}{A}, \overset{ˉ}{B}, C$ input-independent hain. Har token par same filter apply hota hai.

SELECTIVE KYUN? Socho ek document process kar rahe ho: model shayad pehle sentence mein mention kiya gaya naam poori sequence ke liye yaad rakhna chahta hai, lekin filler words ko turant bhoolna chahta hai. Fixed SSMs ye choice nahi kar sakte—Mamba kar sakta hai.

YE EFFICIENTLY KAISE COMPUTE HOTA HAI? Ye convolution view tod deta hai (kernel ab input-dependent hai), isliye FFT use nahi kar sakte. Mamba specialized hardware-aware kernels use karta hai jo recurrent computation ko careful memory management ke saath fuse karte hain, aur FFT-based S4 se comparable speed achieve karte hain bawajood recurrent hone ke.

Worked example Mamba Selection Example

Scenario: Sentence process kar rahe hain "John went to the store and bought milk."

Traditional S4: Har word ko same decay rate milta hai. Agar long memory set karein, to "went" aur "and" ko utna hi yaad rakhenge jitna "John" ko (wasteful). Agar short memory ho, to "John" ko bahut jaldi bhool jaate hain.

Mamba:

Token "John": model large $Δ$ seekhta hai → slow decay → $\overset{ˉ}{A} \approx 0.98$
Token "and": model small $Δ$ seekhta hai → fast decay → $\overset{ˉ}{A} \approx 0.3$
Token "milk": model large $Δ$ seekhta hai → next sentence ke liye yaad rakho

SELECTION KAISE HOTA HAI: $Linear_{Δ}$ network token embedding dekhta hai aur ek scalar output karta hai. Agar token mein high "semantic importance" features hain, to bada $Δ$ output karta hai. Ye task loss se end-to-end seekha jaata hai.

The Mamba Architecture Block

Ek full Mamba block selective SSM ko gating aur projections ke saath combine karta hai:

Formula Mamba Block

Input x (B, L, D)
 ↓
LayerNorm
 ↓
Project to 2*E dimensions
 ↓ (split)
 ↓           ↓
SSM path   Gate path
(selective  (Conv1D 
 SSM)       + SiLU)
 ↓           ↓
 └── ⊙ ────┘  (element-wise multiply)
      ↓
  Project back to D
      ↓
  Residual + Input

GATING KYUN? GLU (Gated Linear Units) ki tarah, ye model ko information flow control karne deta hai. Gate decide karta hai ki SSM-processed information kitni "pass through" ho.

Training vs Inference Modes

Common mistake Common Misconception: "SSMs long-range dependencies model nahi kar sakte"

Galat idea: "RNNs jaise recurrent models fail hue kyunki wo door tak yaad nahi rakh sakte. SSMs recurrent hain, to ye bhi fail honge."

Ye sahi kyun lagta hai: Vanilla RNNs mein lambi sequences par vanishing gradients ki problem hoti hai.

Sach: SSMs isse bachte hain in tareekon se:

Careful initialization: $A$ matrix ko specific eigenvalues (HiPPO initialization) ke saath initialize kiya jaata hai jo stable, long-term memory channels banate hain
Large state dimension: RNNs ke liye $N = 16 - 64$ , SSMs ke liye $N = 1000 - 8000$
Structured parameterization: DPLR structure gradients ko stable rakhta hai
Selection mechanism (Mamba): Dynamically choose kar sakta hai kya yaad rakhna hai Practice mein, Mamba 100K tokens se separated strings copy karne jaise long-range tasks par Transformers se match ya exceed karta hai.

Common mistake Common Misconception: "Linear complexity matlab hamesha faster"

Galat idea: " $O (n)$ , $O (n^{2})$ ko beat karta hai, to SSMs hamesha Transformers se faster hain."

Ye sahi kyun lagta hai: Big-O notation long-term scaling batata hai.

Sach:

Choti sequences ke liye (<2K), Transformers aksar faster hote hain kyunki highly optimized CUDA kernels (FlashAttention) aur hardware acceleration hoti hai
SSMs ke higher constant factors hain (complex number arithmetic, recurrence mein sequential dependencies)
Crossover point jahan SSMs win karte hain wo hardware ke hisaab se around 4K-8K tokens hai
Bahut lambi sequences ke liye (>32K), SSMs dominate karte hain

Fix: Architecture choose karo apni typical sequence length ke basis par. Chatbots ke liye (kuch K tokens), Transformers theek hain. Genomics ya lambe documents ke liye (100K+ tokens), SSMs use karo.

Connections to Other Models

Transformers: SSMs attention ka ek alternative hain, flexibility ko efficiency ke saath trade-off karte hain
RNNs and LSTMs: SSMs "spiritual successor" hain better long-term memory aur parallelization ke saath
Convolutional Neural Networks: Convolution view SSMs ko time-series convolutions se link karta hai
Linear Attention: Kuch linear attention variants mathematically kuch SSM configurations ke equivalent hain
Mixture of Experts: Mamba ko MoE ke saath combine karke aur greater efficiency hasil ki ja sakti hai

Recall Ek 12-saal ke bacche ko explain karo

Socho tum ek lamba gaana sun rahe ho aur use yaad karne ki koshish kar rahe ho. Tumhara dimaag har ek awaaz ko equally yaad nahi karta—kuch parts (chorus, melody) tum zyada yaad rakhte ho, jabki kuch parts (ek dheemi drum beat) tum jaldi bhool jaate ho.

Ek Transformer mein tumhare dimaag ki tarah hota hai jaise ek bade kaagaz par har ek awaaz likh do aur baar baar poora kaagaz dekho. Agar gaana bahut lamba hai (jaise 100,000 notes!), to tumhara kaagaz dev ho jaata hai aur sab check karne mein bahut time lagta hai.

Ek State-Space Model (SSM) ek choti notebook ki tarah hai jahan tum summary notes rakhte ho. Jab tum koi nayi awaaz suno, to tum apni notebook ko purane notes aur nayi awaaz ke basis par update karte ho. Teri notebook choti rehti hai chahe gaana kitna bhi lamba ho!

S4 ne notebook ko super organized (structured) banaya taaki tum use bahut tezi se dekh sako. Mamba ne ise aur bhi smarter banaya: jab koi important awaaz aaye (jaise singer kahe "I love you"), notebook kehti hai "Ise pen mein likho!" taaki ye wahan lamba rahe. Jab boring awaazein aayein, ye kehti hai "Pencil mein likho" taaki ye jaldi fade ho jaayein.

Isliye SSMs bahut bahut lambi sequences handle kar sakte hain (jaise poori Harry Potter book ek saath padho!) jabki Transformers beech mein hi thak jaate.

Flashcards

#flashcards/ai-ml

Ek continuous state-space model define karne wale do key equations kya hain? :: $h^{'} (t) = A h (t) + B x (t)$ (state evolution) aur $y (t) = C h (t) + D x (t)$ (observation)

Sequence length

n

ke liye SSMs ki time complexity Transformers se kaise compare hoti hai?

SSMs

O (n)

O (n lo g n)

hain, jabki Transformers

O (n^{2})

hain attention matrix ki wajah se

Neural networks ke liye ek continuous SSM ko hum kaise discretize karte hain?

\overset{ˉ}{A} = e^{A Δ}

aur

\overset{ˉ}{B} = (A^{- 1} (e^{A Δ} - I)) B

define karo, phir recurrence

h_{k} = \overset{ˉ}{A} h_{k - 1} + \overset{ˉ}{B} x_{k}

use karo

A

matrix par kya structural constraint lagata hai?

S4 diagonal-plus-low-rank (DPLR) structure use karta hai:

A = Λ - p q^{*}

, jo FFT ke zariye efficient computation enable karta hai

Ek SSM ke do computational views kya hain?

Recurrent view (inference ke liye):

h_{k} = \overset{ˉ}{A} h_{k - 1} + \overset{ˉ}{B} x_{k}

, aur Convolutional view (training ke liye):

y = \overset{ˉ}{K} * x

SSM convolution kernel

\overset{ˉ}{K}

kya hai?

\overset{ˉ}{K} = (C \overset{ˉ}{B}, C \overset{ˉ}{A} \overset{ˉ}{B}, C \overset{ˉ}{A}^{2} \overset{ˉ}{B}, \dots, C \overset{ˉ}{A}^{L - 1} \overset{ˉ}{B})

—system ka impulse response

S4 par Mamba ka key innovation kya hai?

Mamba SSM parameters ko input-dependent (selective) banata hai:

\overset{ˉ}{B} (x), C (x), Δ (x)

input token par depend karte hain, jisse model choose kar sakta hai kya yaad rakhna hai

Mamba ka selectivity mechanism kaise kaam karta hai?

Timestep parameter

Δ (x)

ek learned linear layer ke zariye per-token compute hota hai; bada

Δ

slow decay deta hai (long memory), chota

Δ

fast decay deta hai (quick forgetting)

Mamba S4 ka FFT-based convolution trick kyun use nahi kar sakta?

Kyunki kernel ab input-dependent hai (har sequence ke liye alag), jo FFT convolution ke liye zaroori time-invariance tod deta hai

SSM inference vs Transformer inference mein per-token computational cost kya hai?

SSMs:

O (1)

(sirf fixed-size state update karo). Transformers:

O (n)

(saare previous tokens attend karo)

Kaun si initialization strategy SSMs ko long-term memory maintain karne mein help karti hai?

HiPPO initialization, jo

A

ke eigenvalues ko specific values par set karta hai jo stable memory channels banate hain

Kis sequence length par SSMs typically Transformers se faster ho jaate hain?

Hardware aur implementation details ke hisaab se around 4K-8K tokens par; <2K tokens ke liye, optimized Transformers (FlashAttention) aksar faster hote hain