State Space Models बनाम Transformers: एक व्यावहारिक गाइड

ट्रांसफ़ॉर्मर अब ज़्यादा देर तक अकेले खिलाड़ी नहीं रहेंगे। यह दावा थोड़ा बड़ा लगता है, मानता हूँ। हम सालों से देख रहे हैं कि language models से लेकर protein folding तक, हर जगह transformer architecture छाया हुआ है। लेकिन पिछले कई महीनों से production workloads पर state space models की transformer baselines से तुलना करने के बाद मुझे यकीन हो गया है कि बदलाव असली है। इसका मतलब यह नहीं कि SSMs हर मामले में बेहतर हैं। वे नहीं हैं। लेकिन वे ऐसी खास समस्याएँ हल करते हैं जिन्हें transformers मूल रूप से नहीं कर सकते, और नई पीढ़ी ने गुणवत्ता का अंतर इतना कम कर दिया है कि अब इन्हें नज़रअंदाज़ करना एक तकनीकी कर्ज़ वाला फ़ैसला बन गया है।

यह कोई हाइप वाला लेख नहीं है। मैं बताऊँगा कि state space models असल में हैं क्या, Mamba-3 कहाँ सच में सुधार लाता है, SSMs अब भी कहाँ कमज़ोर पड़ते हैं, और आपकी टीम को अपनाने के बारे में कैसे सोचना चाहिए। एक practitioner से दूसरे practitioner की बात।

बड़े पैमाने पर Transformers की दीवार से टकराव क्यों

Quadratic scaling की कहानी आप पहले से जानते हैं। Self-attention लंबाई N के sequence के लिए N×N matrix बनाता है, इसलिए context length दोगुना करने पर compute और memory चार गुना हो जाती है। लंबे समय तक यह ज़्यादा मायने नहीं रखा, क्योंकि models कुछ हज़ार tokens पर चलते थे और hardware उनके साथ चल जाता था।

वह दौर अब खत्म हो चुका है। आज हम जो workloads बना रहे हैं, उनमें अक्सर 100K+ token contexts चाहिए। Code assistants को पूरे repositories देखने होते हैं। Multimodal pipelines घंटों का video प्रोसेस करती हैं। Agents कई दिनों तक चलने वाली बातचीत का इतिहास संभालते हैं। इन पैमानों पर quadratic attention सिर्फ महँगा नहीं, बल्कि एक दीवार बन जाता है।

KV cache की समस्या इसे और बिगाड़ देती है। Autoregressive generation के दौरान हर transformer layer उस हर token के लिए key-value pairs स्टोर करती है जो उसने देखा है। यह cache हर layer में linearly बढ़ता है और GPU memory को तेज़ी से खा जाता है। मैंने देखा है कि एक 7B transformer सिर्फ 128K context पर KV cache के लिए 40GB VRAM खा जाता है। यह वह memory है जिसे आप ज़्यादा requests को batch करने में इस्तेमाल नहीं कर सकते।

Quadratic memory scaling के कारण standard attention के साथ दस लाख-token contexts लगभग असंभव हैं
KV cache की वृद्धि प्रति GPU समवर्ती users की संख्या सीमित करती है, जो production में सीधे लागत बढ़ाने वाला कारक है
लंबे-context transformer inference की ऊर्जा खपत अब उचित ठहराना मुश्किल होता जा रहा है
Real-time applications (robotics, edge AI) को sub-millisecond token generation चाहिए, जो attention नहीं दे सकता
'Attention sink' की घटना बहुत लंबे sequences पर गुणवत्ता घटाती है, भले ही आपके पास memory हो

ये सिद्धांत की बातें नहीं हैं। पिछले साल मेरे साथ काम करने वाली तीन अलग-अलग टीमों ने इन्हीं कारणों से विकल्पों का गंभीरता से मूल्यांकन शुरू किया।

State Space Models असल में कैसे काम करते हैं

State space models control theory से आते हैं, जहाँ इनका उपयोग दशकों से dynamical systems को model करने में होता रहा है। मूल विचार सरल है: attention की तरह हर आउटपुट की गणना के लिए पिछले सभी tokens देखने के बजाय, आप एक संकुचित hidden state रखते हैं जो समय के साथ बदलता है। नए tokens state को अपडेट करते हैं और state से आउटपुट निकलता है। बस, इतना ही।

गणितीय रूप से SSM को चार matrices, A, B, C और D, से परिभाषित किया जाता है, जो तय करते हैं कि input x के जवाब में hidden state h कैसे बदलता है। Sequential data के लिए इन continuous equations को discretize करने पर आपको एक recurrence मिलती है, जिसे inference के समय कंप्यूट करना बेहद आसान है।

import torch
def ssm_step(A_bar, B_bar, C, D, h, x_t):
"""Single SSM step: O(1) memory, O(1) compute.
Compare this to attention, which needs to look at
every previous token. The SSM just updates its state.
"""
h_new = A_bar @ h + B_bar @ x_t  # Update hidden state
y_t = C @ h_new + D * x_t         # Compute output
return h_new, y_t
def ssm_generate(A_bar, B_bar, C, D, tokens, embed):
"""Autoregressive generation with constant memory.
Whether you've processed 100 tokens or 500,000,
this uses the same amount of memory.
"""
h = torch.zeros(A_bar.shape[0])
outputs = []
for t in tokens:
x_t = embed(t)
h, y_t = ssm_step(A_bar, B_bar, C, D, h, x_t)
outputs.append(y_t)
return torch.stack(outputs)

खूबसूरती कोड में ही दिख जाती है। Inference में हर token के लिए O(1) memory और O(1) compute लगता है, sequence की लंबाई चाहे जो हो। कोई KV cache नहीं। कोई quadratic विस्फोट नहीं। पूरा इतिहास hidden state vector में संकुचित हो जाता है।

पकड़ यह है: training के दौरान इस recurrence को क्रमिक रूप से चलाना बहुत धीमा होगा। चाल यह है कि वही computation convolution या parallel scan के रूप में दोबारा लिखी जा सकती है, जिसे GPUs कुशलता से संभाल लेते हैं। यानी आपको parallel training और recurrent inference दोनों मिलते हैं, दोनों दुनियाओं का सबसे अच्छा।

Mamba से Mamba-3 तक: हर पीढ़ी ने क्या ठीक किया

S4 जैसे शुरुआती SSMs ने अवधारणा साबित की, पर उनकी एक बड़ी कमज़ोरी थी: वे content-based reasoning अच्छी तरह नहीं कर पाते थे। State transition matrices सभी inputs के लिए तय थीं, इसलिए model यह तय नहीं कर सकता था कि वास्तव में पढ़ी जा रही सामग्री के आधार पर क्या याद रखना है और क्या भूलना है। यह ऐसा है जैसे नोट्स लेने का नियम हो 'हर तीसरा शब्द लिखो'। कुछ उपयोगी जानकारी मिल जाएगी, पर आप यह नहीं बदल सकते कि क्या ज़रूरी है।

Albert Gu और Tri Dao द्वारा 2023 के अंत में पेश किए गए Mamba ने इसे एक सुंदर विचार से ठीक किया: SSM parameters को input पर निर्भर बनाओ। तय A, B, C matrices की जगह, Mamba उन्हें मौजूदा token के फ़ंक्शन के रूप में निकालता है। Model प्रासंगिक जानकारी चुनकर स्टोर करना और शोर को छोड़ना सीखता है। इस 'selective' तंत्र ने SSMs को वह content-awareness दी जो उनमें नहीं थी।

Mamba-2 ने सैद्धांतिक अंतर्दृष्टि दी कि structured SSMs और linear attention गणितीय रूप से dual हैं, State Space Duality (SSD) framework के रूप में। यह सिर्फ अकादमिक नहीं था। इससे hardware-aware implementations संभव हुईं, जो GPU tensor cores का बेहतर उपयोग करती हैं और training throughput को काफी बढ़ाती हैं।

Deployment के नज़रिए से Mamba-3 सबसे दिलचस्प है। तीन नवाचार सबसे ज़्यादा मायने रखते हैं:

Multi-scale state tracking: model एक साथ कई समय-रिज़ॉल्यूशन पर state रखता है, जिससे स्थानीय पैटर्न और लंबी दूरी की निर्भरता दोनों पकड़ी जाती हैं, किसी एक को कुर्बान किए बिना
Adaptive state compression: जटिल reasoning वाले अनुच्छेदों के लिए hidden state गतिशील रूप से फैलता है और अनुमानित टेक्स्ट के लिए सिकुड़ता है, जिससे गुणवत्ता खोए बिना compute बचता है
Improved initialization और gating: बड़े पैमाने पर training stability में नाटकीय सुधार हुआ, जो तब बहुत मायने रखता है जब आप training run पर लाखों खर्च कर रहे हों

Mamba-3 हर benchmark पर transformers को नहीं हराता। उसे ऐसा करने की ज़रूरत भी नहीं है। ज़्यादातर standard evaluations पर यह गुणवत्ता में बराबरी करता है, और इसके लिए inference compute का एक छोटा हिस्सा लगता है। ज़्यादातर production workloads के लिए यही वह trade-off है जो मायने रखता है।

Linear Attention और SSMs का मिलन

एक समानांतर धारा भी समझने लायक है। Linear attention उसी efficiency समस्या पर काम करता है, लेकिन transformer ढाँचे के भीतर से। Standard attention पूरा N×N matrix बनाता है। Linear attention softmax की जगह एक decomposable kernel function लगाता है और गणित को इस तरह पुनर्व्यवस्थित करता है कि वह quadratic matrix कभी बने ही नहीं।

# Standard attention: O(N^2 * d)
# score = softmax(Q @ K.T / sqrt(d)) @ V
# Linear attention: O(N * d^2)
# Replace softmax with kernel feature map phi()
# Rearrange: compute K^T @ V first (d×d), then multiply by Q
def linear_attention_step(q_t, running_kv, running_k, k_t, v_t, phi):
"""Incremental linear attention — runs like a recurrence.
This is why SSMs and linear attention are duals:
both compress history into a fixed-size state.
"""
k_feat = phi(k_t)
q_feat = phi(q_t)
running_kv = running_kv + k_feat.unsqueeze(-1) * v_t.unsqueeze(-2)
running_k = running_k + k_feat
y_t = (q_feat @ running_kv) / (q_feat @ running_k + 1e-6)
return y_t, running_kv, running_k

उस कोड को ध्यान से देखिए। Linear attention, जब क्रमिक रूप से चलाया जाए, एक running state रखता है और हर नए token के साथ उसे अपडेट करता है। जाना-पहचाना लगा? लगना भी चाहिए, क्योंकि यह लगभग वही काम कर रहा है जो एक SSM करता है। SSD framework ने इस संबंध को औपचारिक रूप दिया, और यह हाल के sequence modeling शोध की सबसे महत्वपूर्ण सैद्धांतिक अंतर्दृष्टियों में से एक है।

GLA (Gated Linear Attention) और RetNet variants जैसी architectures ने इसे और आगे बढ़ाया है, data-dependent gating जोड़कर, जो linear attention और selective SSMs के बीच की रेखा लगभग पूरी तरह धुंधली कर देता है। व्यावहारिक निष्कर्ष: इन्हें प्रतिस्पर्धी तरीके न समझें। ये एक-दूसरे की ओर बढ़ रहे हैं।

Hybrid Architectures: production में असल में क्या जीत रहा है

जब कोई टीम पूछती है कि क्या उन्हें SSMs पर जाना चाहिए, तो मैं यही कहता हूँ: किसी एक शुद्ध तरीके पर मत जाओ। अभी सबसे अच्छे नतीजे वे architectures दे रहे हैं जो SSM layers को कुछ attention layers के साथ मिलाते हैं। अलग-अलग computational primitives अलग-अलग कामों में अच्छे हैं, और इसे अनदेखा करने से प्रदर्शन छूट जाता है।

SSM layers sequential जानकारी को कुशलता से संकुचित करने और आगे बढ़ाने में माहिर हैं। Attention layers अब भी सटीक, content-based retrieval में बेजोड़ हैं: 'पेज 47 की वह सटीक पंक्ति ढूँढो जो इस सवाल का जवाब देती है।' अच्छी तरह डिज़ाइन किया गया hybrid अपनी 80-90% layers में SSMs रखता है और attention को वहीं छिड़कता है जहाँ उसकी सबसे ज़्यादा ज़रूरत है।

Jamba-style models: Mamba और attention layers को बारी-बारी से रखना, MoE feed-forward blocks के साथ, जो कुशल SSM processing और सटीक attention के बीच गतिशील रूप से routing करते हैं
Griffin-family designs: recurrent gated linear units को local sliding-window attention के साथ जोड़ना, कम full attention के साथ मज़बूत नतीजे
Mamba-attention hybrids: ज़्यादातर layers में Mamba-3 blocks, और वैश्विक जानकारी के routing के लिए रणनीतिक गहराइयों पर full attention layers
StripedHyena के उत्तराधिकारी: gated convolutions, SSM layers और sparse attention को NAS-अनुकूलित पैटर्न में interleave करना

आंकड़े इसकी पुष्टि करते हैं। कई स्वतंत्र समूहों ने दिखाया है कि 85/15 SSM-से-attention अनुपात समान parameter count पर शुद्ध transformer की गुणवत्ता से मेल खाता है, और साथ ही inference FLOPs को 40-60% घटा देता है। लंबे-context workloads के लिए memory बचत और भी बड़ी है। यह मामूली सुधार नहीं है। यह आपका GPU बिल आधा करने जैसा है।

Production Benchmarks: SSMs कहाँ चमकते हैं और कहाँ नहीं

आंकड़ों के बारे में विशिष्ट होने दीजिए, क्योंकि अस्पष्ट efficiency दावे deployment के फ़ैसले लेने वालों के किसी काम के नहीं होते।

Inference throughput: 8B parameters वाला Mamba-3 आधारित मॉडल, चाहे context 1K tokens का हो या 500K का, एक ही रफ़्तार से tokens generate करता है। तुलनीय transformer, KV cache बढ़ने के साथ धीरे-धीरे धीमा होता जाता है। 500K context पर SSM model प्रति GPU 5-8 गुना अधिक throughput देता है। यह सिद्धांत नहीं है, मैंने इसे मापा है।

समवर्ती users: KV cache के बिना SSM models कहीं ज़्यादा एक साथ आने वाली requests संभाल सकते हैं। एक A100 पर, जहाँ transformer 32K context पर शायद 8 समवर्ती streams संभालता है, वहीं समकक्ष SSM model 30+ संभाल सकता है। बड़े पैमाने पर inference चलाने वालों के लिए यही वह संख्या है जो अर्थशास्त्र बदल देती है।

Training speed: यहाँ लाभ अधिक सामान्य हैं। H100 clusters पर Mamba-3 समकक्ष transformer से लगभग 1.4 गुना throughput पर train होता है। लंबे sequences के साथ अंतर बढ़ता है: 32K tokens से ऊपर SSM training 2-3 गुना तेज़ चलती है, क्योंकि यह quadratic attention से पूरी तरह बचता है।

लेकिन सीमाओं के बारे में मुझे ईमानदार होना पड़ेगा। लंबे contexts से सटीक शब्दशः recall वाले कामों पर, जैसे 'लाइन 4,382 पर सटीक error message क्या था?', शुद्ध SSMs अब भी पीछे रहते हैं। तय आकार की संकुचित state एक lossy representation है। Attention मूल tokens पर वापस देख सकता है। इसीलिए hybrid architectures काम करती हैं: attention layers वही retrieval संभालती हैं जो SSMs नहीं कर सकते।

जहाँ SSMs अब भी कमज़ोर हैं

मैं बाकी कमियों के बारे में साफ़ नज़र रखना चाहता हूँ, क्योंकि अधूरी जानकारी के आधार पर नई architecture अपनाना छह महीने बर्बाद करने का आसान तरीका है।

In-context learning: Transformers prompt में few-shot उदाहरणों के आधार पर अपना व्यवहार बदलने में अब भी बेहतर हैं। SSMs यह कर सकते हैं, पर कम भरोसेमंद तरीके से। अगर आपका application उदाहरणों वाली prompt engineering पर भारी निर्भर है, तो शुद्ध SSMs निराश करेंगे।
Ecosystem maturity: Transformer tooling को वर्षों का optimization मिला है। SSM-specific kernels, serving infrastructure और fine-tuning libraries तेज़ी से सुधर रही हैं, पर अभी बराबरी पर नहीं हैं। Integration के लिए अतिरिक्त समय रखें।
70B से ऊपर scaling अनिश्चितता: 70B parameters तक के Mamba-3 मॉडल अच्छे scaling curves दिखाते हैं, पर 200B+ के frontier पर हमारे पास मज़बूत आंकड़े नहीं हैं। क्या SSM scaling laws अत्यधिक आकारों पर भी लागू होते हैं, यह वास्तव में अज्ञात है।
Fine-tuning techniques: Transformers के लिए LoRA और QLoRA अच्छी तरह समझे जाते हैं। SSM architectures पर इन्हें लागू करने के लिए अलग तरीके चाहिए, और best practices अभी तय हो रही हैं।
Hardware mismatch: मौजूदा GPUs उन matrix multiplies के लिए optimized हैं जिन्हें attention पसंद करता है। SSMs भारी रूप से parallel scans पर निर्भर करते हैं, जो आधुनिक hardware पर ठीक-ठाक चलते हैं, पर वह operation नहीं हैं जिसके लिए GPUs बनाए गए थे।

इनमें से कोई भी रुकावट नहीं है। ये ज्ञात समाधान वाली engineering समस्याएँ हैं। पर ये वास्तविक हैं, और आपकी timeline में इन्हें गिनना चाहिए।

व्यावहारिक सिफ़ारिशें: कब अपनाएँ और कैसे शुरू करें

कई production workloads पर SSMs का मूल्यांकन करने के बाद, टीमों को सलाह देने के लिए मैं यह ढाँचा इस्तेमाल करता हूँ।

अगर आपके workload में लंबा-context inference (नियमित रूप से 32K+ tokens), उच्च concurrency की ज़रूरत, या latency-sensitive edge deployment शामिल है, तो आक्रामक तरीके से अपनाएँ। ROI बड़ा और तुरंत मिलने वाला है। Pure SSM से शुरू करने के बजाय Jamba या Griffin-family जैसे hybrid architecture से शुरू करें, ताकि ज़्यादातर efficiency लाभ कम जोखिम के साथ मिल जाएँ।

अगर आपका workload मुख्यतः छोटे-context वाला है, in-context learning पर भारी निर्भर है, और inference लागत का दबाव नहीं है, तो इंतज़ार करें और देखें। यहाँ transformers की बढ़त अब भी है, और उनका ecosystem ज़्यादा परिपक्व है।

फ़ैसला करने से पहले अपने वास्तविक inference workload को प्रोफ़ाइल करें: median context length और समवर्ती users की संख्या मुख्य चर हैं
Pure SSMs की जगह hybrid architectures से शुरू करें, क्योंकि ये कम जोखिम वाले हैं और फिर भी 40-60% inference लागत में कमी देते हैं
अपने विशिष्ट कामों पर benchmark करें: SSMs summarization और लंबी-दूरी reasoning में अच्छे हैं, पर सटीक retrieval में पीछे रहते हैं
अभी से architecture-comparison infrastructure बनाएँ: सिर्फ accuracy नहीं, आपको latency, throughput, memory और प्रति query लागत मापनी होगी
SSM tooling ecosystem को हर तिमाही ट्रैक करें, क्योंकि सुधार की रफ़्तार इतनी तेज़ है कि आज जो अव्यावहारिक है, वह तीन महीनों में production-ready हो सकता है

Architecture परिदृश्य बँट रहा है, और यह अच्छी बात है

सबको एक ही architecture से चलाने का दौर खत्म हो रहा है। हम ऐसी दुनिया की ओर बढ़ रहे हैं जहाँ टीमें computational primitives चुनती हैं, जैसे full attention, linear attention, selective SSMs, gated convolutions, और अपनी विशिष्ट बाधाओं के आधार पर उन्हें संयोजित करती हैं। परिपक्व engineering disciplines ऐसे ही काम करती हैं। हर संरचना स्टील से नहीं बनती। आप सामग्री उस भार के आधार पर चुनते हैं जो उसे उठाना है।

Transformer मरा नहीं है। कई workloads के लिए यह अब भी सबसे सिद्ध architecture है, और आने वाले वर्षों में यह महत्वपूर्ण AI systems को चलाता रहेगा। लेकिन sequence modeling में state-of-the-art पर इसका एकाधिकार खत्म हो चुका है। SSMs और hybrids ने अपनी जगह production के प्रथम श्रेणी के औज़ारों के रूप में कमाई है, शोध की जिज्ञासा के रूप में नहीं।

हम में से जो असली सिस्टम बना रहे हैं, उनके लिए अधिक architectural विकल्पों का मतलब है खास समस्याओं के लिए बेहतर औज़ार। यह डर की वजह नहीं, बल्कि इस्तेमाल करने लायक engineering leverage है।