4521, what is actually in the vector? People say magnitude is linearly decodable. Ok fine. But does the model keep a real quantity like a number line? Or just rank order? Or digit symbols glued together? Everyone cites probes, nobody lines the stories up on the same stimuli.
So I spent a long stretch building constraint tests across a bunch of models (small Qwen up to DeepSeek 70B, Llama, Phi, GPT-2, BERT, etc). Not to confirm my favorite theory. Mostly to break theories.
This is my informal writeup of that work. Numbers live in Experiment2/results/*_verdict.json if you want the raw logs.
If you sample integers uniformly from 1 to a million, value, log value, and rank all move together. Monotone.
So a probe that only cares about ordering cant tell you if the model encoded cardinal size, ordinal rank, or digit pattern. I formalized that in the writeup as a non-identifiability thing. Practical fix: use gapped number sets where rank and log value disagree, plus cross val, shuffle nulls, bootstrap CIs. Dont trust single split R² from a uniform grid.
Also intrinsic dimension on 100 points in 1500d looks “thin” for fake reasons. Small sample. I had to stop trusting my own earlier manifold plots until I reran with controls.
Six rough accounts from the literature:
I wrote down 16 constraints any story has to survive. Things like: log decodable, no vector arithmetic, digit identity readable, magnitude weak vs total variance, no shared axis across domains, causal necessity, notation robustness, etc.
The scorecard table in my notes is just a map. Experiments carry the weight, not the tick count.
Log value probes work across models and domains. R² often 0.94 to 1.00 on the setups I ran.
But geometry is compressed. Between/within cluster ratios κ around 2 to 4.8. Isometry to log distance is moderate (ρ ~ 0.2 to 0.5). So yes metric-ish in log space, not a clean ruler.
Partial distance tests on gapped sets: log value still beats pure rank in the large encoders (bootstrap CI excludes 0 for Llama3 / DeepSeek on several domains). Small effects on Qwen 7B integer and GPT-2 distance were marginal. I report those without pretending they’re huge.
This was the mechanism result that stuck.
Decompose log10(n):
So the model is mostly counting digits in the string, not holding an abstract scalar. Log compression falls out of positional notation.
Notation test backed this: same value as plain digits, comma form, scientific, words. Log still decodable 0.90 to 0.99.
Zero pad to fixed width (kills raw token length cue): R² drop ≤0.04, mean ~0.015. Magnitude still there. So it’s where significant digits sit, not raw token count per se.
Same numeric value in different templates (population vs distance vs price vs weight, eight domains total).
Probe directions for magnitude barely align. Cross domain cosine mean ~0.12 to 0.25. Above random noise in high D, far below “same axis.”
Pooled leave-one-domain-out R² weak (up to ~0.46 on Qwen, negative on Llama3 in one setup). So you dont get one universal magnitude direction.
But magnitude information still transfers somewhat (calibration free Spearman 0.53 to 0.97 depending model/domain). My read: context picks a new readout direction each time, order info is partly portable, not a shared register.
Addition composition in embedding space doesnt work like vector add. Product from concat embeddings fails OOD. Group symmetries absent (that constraint is noisy but directionally right).
So dont expect h(347)+h(892) ≈ h(1239).
On Qwen2.5-1.5B comparison task (“is a > b?”):
| condition | accuracy |
|---|---|
| baseline | 0.91 |
| erase magnitude subspace all layers | 0.53 (chance) |
| erase random rank-8 subspace | 0.67 ± 0.15 |
Erase length (#digits) subspace: same-order pairs mostly ok (0.90), cross-order hurt (0.62).
Erase leading digit subspace: same-order collapses (~0.50).
Magnitude subspace and length subspace overlap heavily (cos ~0.89). Leading digit more distinct (cos ~0.43). Double dissociation at 1.5B.
Replicated magnitude collapse at 7B (0.55 vs baseline ~0.92). Length-only selective effect was cleaner at 1.5B; at 7B length erasure looked more redundant. I dont oversell that as universal.
Point: the comparison behavior needs the magnitude/length directions. Geometry isnt decorative.
I tried adding one domain’s magnitude direction into another domain’s residual stream mid layers. Hoped for weak cross domain shift.
Within domain moved judgments. Cross domain looked smaller. Then random direction control moved things similarly. After correction, contrast unreliable. Generic perturbation bias. I dont use steering as evidence. Correlational orthogonality still stands.
Contextual Compositional Magnitude is just my label for what didnt die:
CCM isnt claiming I discovered log compression or probing. Lot of that exists in 2025 papers. My bit is adjudication + identifiability caution + length mechanism with erasure + no-shared-direction result.
On the internal scorecard CCM marks all 16 constraints. Closest competitor is the string/context hybrid account at 13. Take the table as organization, not proof.
Integer templates, eight ordered domains, gapped geometries for identifiability, repeated CV probes, bootstrap where noted, twelve models total (nine core + three for notation battery), causal all-layer subspace projection on Qwen 1.5B and 7B.
Code under Experiment2/ experiments. Key json: ablation_verdict.json, notation_verdict.json, domains_verdict.json, transfer_refine.json, bootstrap_ci.json, geometry_verdict.json, steering_verdict.json.
I was running multiply prompts on our local GPUs one evening. Same template every time, greedy decode, no chain of thought. And I kept getting annoyed at the answers.
The model sounds sure of itself. The number looks fine at a glance. About the right length. Like someone who knows the answer is “big” but never did the multiplication.
So I started saving outputs properly. Then I wondered what the internal state looks like right before it prints anything. That turned into a bunch of probe runs on the same box.
Writing this down for myself mostly. If you read the verdict jsons later, this is the english version.
Prompt:
Question: What is 347 times 892? Give only the number.
Answer:
No CoT. Greedy. Max 12 tokens. Mostly 2 to 3 digit operands. Few hundred pairs per model.
Two scores:
Exact multiply is bad. Scale is oddly good.
| Model | Exact | OOM (right scale) |
|---|---|---|
| Qwen2.5-1.5B | 7% | 93% |
| Qwen2.5-7B | 22% | 100% |
| Llama3-8B | 12% | 100% |
| Phi-3-mini | 13% | 74% |
| Mistral-7B | 3% | 79% |
Qwen 1.5B: 93% get the scale right. 7% get the product.
I call it Fermi decoding bc its like a Fermi guess. Right size, wrong value.
Token limit?
Bumped max_new_tokens 12 to 48. Exact still ~9%. So not truncation.
All ops broken?
Same prompt style for addition on Qwen 1.5B: 92% exact, 99% OOM. Add is fine. This is mostly a multiply thing with this template.
Wrong answers shape
On 1.5B wrong multiply answers, 90.6% still OOM-correct. Usually right length, wrong digits. Not random garbage.
Scale in hidden state before decode?
Probed last prompt token. Linear probe on product scale R² ~0.82 to 0.85, about same as digit count probe. So scale info is often already there internally. Feels like the model fails at spitting out exact digits, not at having zero idea of size.
If your benchmark only checks “did it output a plausible sized number” you will overrate these models on multiply.
After logging outputs I wanted to know if multiply vs add look different inside when operands are the same.
Linear probes on hidden states. Ridge regression, cross val, shuffle nulls (shuffle labels, probe should collapse). Mainly Qwen2.5-1.5B, some reruns on 7B and Phi-3.
Binary probe for operation word: R² ~1.0 on three models. Shuffle nulls negative.
Then cosine between mul and add last-token states for same (a,b): ~0.9999. Basically the same vector.
So yeah you can read off times vs plus. But geometrically its not two separate clusters. More like a tiny shift on top of almost identical states.
Layer 4:
Both prompts still have a and b in the text. Probe can pull operands and combine. Dont read that as “add mode stores multiplication.”
Operand digit hiddens -> log10(product) R² ~0.98. Last token only adds ~0.01 on top.
Bilinear h_a * h_b does worse than concat. So not some clever multiplicative geometry, mostly linear readout.
Static embeddings of operands predict mul prompt state R² ~0.73. Context changes the vector. Magnitude still easy to read once youre in the prompt.
Digit identity R² ~0.97, place R² ~0.95 on operand digit tokens in mul prompts.
But token index on same rows: R² ~0.99. So “place” might just be where you are in the sequence. I wouldnt claim abstract place value from this alone.
Static Integer:n vs mul prompt cosine ~0.32.
Static vs dumb filler prompt (“summarize weather in Paris”) cosine ~0.21. So low cosine is partly just different prompts, not magic arithmetic geometry.
Reword mul template, same numbers: cosine ~0.72. Wording changes the state.
Shift add state along mean(mul - add). Classifies op fine. log10 product vs log10 sum probe scores basically dont move. Direction tags the op, doesnt act like a magnitude switch.
Qwen 1.5B layer 2: log10(product) probe R² ~0.996
Same setup greedy multiply: 7% exact, 93% OOM
Hidden state looks informative. Output often wrong anyway. Decodable doesnt mean the model uses it correctly at decode time.
| output logging | probes | |
|---|---|---|
| exact multiply | low on Qwen | didnt focus on this |
| scale / OOM | high on Qwen | log product R² ~0.99 in prompt |
| where it fails | generation | not early scale readout |
| mul vs add | Fermi mostly on mul | op readable, states nearly same |
| causal | no | no, correlational |
How I explain it to someone in the lab:
Fixed templates. Forward pass for probes, greedy for outputs. Usually last prompt token. Ridge, 5-fold CV repeated, shuffle nulls. Numbers in cursor-work/results/, code in cursor-work/experiments/.
Split exact and OOM on multiply. Dont use probe R² as a competence score. If you care about operation understanding use filler prompts and paraphrases and position baselines. Add under same terse format is much easier, dont assume all ops look like mul.
Started bc multiply answers looked like good guesses with wrong exact digits. Still does across the models I tried. Inside, scale probes easy, op word decodable but tiny on shared state, magnitude doesnt split clean by operation.
Gap between probe R² ~0.99 and 7% exact is what stuck with me. Feels like readout/generation is where it breaks, not that magnitude is totally absent from representations.
Wouldve helped me to read something like this before prompt 200.
Vishesh
]]>I wrapped up my work and i am sitting in front of the main building of the Indian institute of science…its so soothing cold the winds of november are giving my mind thrills…
Ok so we are about to talk about tokenizers today!!! Why does it come to mind??? As we go through reading foundational models …”How LLMs understand english is one of the many questions readers come around!”
And tokenizer is what breaks sentences into tiny tokens, tokens are not words, not characters not syllables, tokens are subword units chosen to make the entire text space efficient to represent. The cleverness and strategy while tokenizing the corpus we take decides how well our model is able to grasp the context and generate the next token with a good tpc (tokens per character)(lower is better) ohhhh what’s tpc now!!!! …calm down we will talk about this now, it’s just a simple measure of how much total token your model can generate divided by the total number of characters in the text…
For example : if the gpt-pro-max can generate 11 tokens for the sentence “Who is Ramanujan??” and on the other side gpt-potato generates 23 tokens for the same sentence..can you spot the time and space complexity which version of gpt will take more??? Yes you got it, that’s gpt-potato… because to represent a entity he is taking more chars then gpt-pro-max….
Yeah moving forward while keeping the tone of discussion light we will talk and in between i will keep making you know the further concepts tokenizers are associated with or famous tokenizers how they cleverly make text space efficient.
You must be wondering how they look, is there any pattern or symmetry between them…
If you are not thinking this, okay be calm i will make you think. You see every tokenizer when get trained it makes a finite list of tokens that finite list of tokens we called vocabulary and using that we make every words which is in the text space so just like 26 alphabet is vocab of english lang…Tokenizers invent an alphabet, but the letters are allowed to be as big or as weird as they want, as long as the final reconstruction cost stays small.
According to some notable papers , they reveal promising benchmarks on upgrading and working on optimization of tokenizers as we know earlier researchers were only interested in thinking about optimizing the architecture or inventing one… This is how impactful tokenizers can be!!!!!
Tokenizers are bound by a deterministic rule which is very senseful, just like when you translate a sentence in hindi , it should not get changed when you translate it back to english. Similarly when we tokenize the corpus of text we expect the same text corpus if we “"”detokenize”””” it in future. detokenize wired words right!!! that’s why we introduce encode and decode nomenclature … when you tokenize something you are encoding it and decoding it when you need your corpus back to its original state.
Great now you become a geek who knows basis of toeknizers…lets take this article further and i will discuss you about some of the coolest topics, my friends who 1 year back sit on a placement they were telling me interviewers are really interested if one knows about modules which makes generative models work …everyone knows how gpt works right transformers transformers transformers but knowing what qkv scoring is , knowing how multi head attention is being computed or knowing the working of positional encoding is what keeps your answers and replies in placement healthy…
Just like i told you about tpc…there are so many entities to judge these tokenizers we call it Tokenization cost metrics….you are sitting in a MAANG interview for gen ai JD and interviewer is not asking TCM….thats where your dream will break and you are not really giving an interview…lol
Ok back to the topic! I told you how tpc can judge which tokenizer is efficient and cost effective similarly we are blessed with more TCMs like: Token Per Word (TPW), Tokenization Entropy, OOV Rate (Out-of-Vocabulary)…
We will discuss one by one ; just like TPC tells us how expensive a tokenizer is for a language…TPW helps readers gauge efficiency quickly but according to me it’s less scientific. “WHYWHY VISHESH???” I dont know!!!! It just doesn’t feel scientific to me.
I am so unpredictable lol..just like me there is a term to judge how unpredictable the token boundaries are for that language and tokenizer that is called Tokenization Entropy. And for domain specific writing like medicine or cook book to make veg food, How often tokenizers fail to represent a term and have to back off to characters is called its OOV Rate aks Out-of-Vocabulary Rate.
Now there are some more concepts which do matter in terms of tokenizers but works under the hood , i have written their definition you can read if you like:
NOW NOW NOW We will talk …oh who switched my speaker off! Ok …. NOW NOW NOW WE WILL TALK ABOUT THE MAIN TOKENIZER FAMILIES AND THEIR WORKING..huh that was lowd.
One of the earliest used tokenizers was Byte-Pair Encoding(BPE) used in GPT-2/3, BLOOM…and working on these will cost me a whole new article so I will give you an intuition!
BPE starts with characters -> counts which pair of characters appear together the most -> merge those pairs and treat them as a single symbol(merge tables) -> repeat 100000+ times ->Final merges -> become the vocabulary.
Merge table is an ordered list of pair-merge decisions the tokenizer learned during training.
It’s simple, deterministic, and compresses well, but it struggles with languages where words morph heavily.
It looks like this:
0: ("e", "r") → "er"
1: ("pl", "ay") → "play"
2: ("play", "er") → "player"
3: ("player", "s") → "players"
Start text = token
Merge table it formes:
1: ("t", "o") → "to"
7: ("k", "e") → "ke"
10: ("to", "k") → "tok"
52: ("ke", "n") → "ken"
100: ("tok", "en") → "token"
possible_doubt : why not (“e”,”n”) -> “en”
Because merging a pair in a merge table is not logical! BPE never decides merges based on logic.It only follows the statistics of the training corpus.
Encodings happens like:
Google has a famous library called Sentence piece where they have introduced Unigram Language Model Tokenizer, using which all the LLaMA family is trained on!!!
So unlike BPE instead of merging it starts with a giant list of candidate sub-words and trains a probabilistic model. The goal of the model “what is the most likely set of subwords that explains all sentences???” and iteratively prunes the vocabulary.
It handles multilingual and weird morphology better than BPE. It generates multiple possible tokenizations and picks the best.
But google didn’t come up with sentence pieces in the first place Early google models like BERT, DistilBERT were trained on different tokenizers known as WordPiece. Its where Google was transitioning their aura form BPE to Unigram LM Tokenizer.
So that transition development is what they call WordPiece. It’s very similar to BPE but with different scoring objectives, they asked “if I add this subword to vocabulary will it help me to explain my text better?????” and thats all maximize likelihood really meant.
WordPiece aften ads “##” markers to show subword continuations. People still talk about this because it introduced the idea that tokenization can be tightly coupled with the LM’s training data.
I told you BPE was used in GPT-2 / 3…OpenAI really took tokenizers seriously and while hunting they found Byte-Level BPE they have used in GPT4 and Mistral also comes into account for using Byte-Level BPE.
How does it work ?? same as bpe, but it starts from raw bytes(0-255) instead of characters.that means every piece of text is representable thus no OOVs. It became dominant because now it can handle emojis, accents, programming languages, and rare symbols effortlessly.
But now OpenAI uses there own tokenizers they call it tiktoken / GPT-4’s Encodings As it introduces modern pretokenization rules, special tokens, optimized merges, extreme speed by memory mapped tables; i will talk about this now …
People says tiktoken is nothing but “BPE but fast” that’s very surface level…because underneath, there are a few design moves that make these tokenizers behave cleaner, tighter, and more predictable than classical BPE.
Pretokenization Rules ie OpenAI tokenizers don’t treat raw text as-is. They apply collapsing weird whitespaces patterns into std ones, normalizing some unicode edge cases, treating punctuation in more consistent way, spitting boundaries around certain characters…its not like full normalization (like NFKC) it’s more like “Fix the obvious inconsistencies but don’t destroy meaning.”
Special Tokens : Gpt model rely on a fairly rich set of reserved tokens i.e. end-of-text, beginning-of-assistant message, beginning-of-user-message, system prompts, document seperators, bye-level fallbacks and some metadata tokens for chat formatting… It’s not at all impressive right, but together they let the model understand where a conversation turn begins, what role a segment plays, and how to differentiate data from instructions.
Older tokenizers basically ignored structure, tiktoken has just baked structure into vocabulary itself
Optimized Merge Table : You know BPE generates very large, redundant, slow scan merge tables.OpenAI reworks this table into something extremely compact and efficiently indexable.Every merge is packed in a way that reduces cache misses and speeds up lookup time during encoding.
Memory-Mapped Tables for Speed : I can say this is pure system engineering, Instead of loading merge rules into RAM and parsing them at runtime, tiktokenn uses memory mapped files: the OS loads the pages you actually touch, random access is extremely fast, zero copy overhead.
Soft Normalization : This is the part most people are not aware of. Older BPE tokenizers got wrecked by slight Unicode variations: smart quotes, accent marks, zero-width spaces, exotic whitespace… All these caused fragmentation, meaning the tokenizer produced way too many tokens for the same concept.
]]>Let’s start with something obvious that somehow no one treats as obvious: language isn’t just communication. It’s compression. It’s the lossy encoding of an entire culture’s worldview into syntax, idioms, metaphors, jokes. When you train a large language model, you’re not just training it to predict the next word — you’re training it to predict reality as seen through a particular civilization’s eyes.
Now think about India. 1.4 billion people, hundreds of languages, thousands of dialects, infinite cultural contexts — all stacked on top of one another like some cosmic recursion of Babel. Every village, every caste, every generation has its own dataset. Yet our AI, our so-called “intelligent” systems, are trained mostly on English internet text, scraped from Reddit, Wikipedia, and Twitter. That means our models speak like Silicon Valley. They reason like Silicon Valley. They hallucinate like Silicon Valley.
But India doesn’t think like Silicon Valley.
Indian reasoning — in its philosophical, linguistic, and even mathematical roots — is nonlinear. Sanskrit grammar was formalized two thousand years ago by Pāṇini in what is effectively a programming language: recursive, compositional, modular. Indian logic, from Nyāya to Vedānta, doesn’t separate emotion from inference the way Western logic does. And the cadence of Indian languages — the way they intertwine context, relationship, and implication — is structurally richer than English tokenization ever captured.
When you realize that, it becomes obvious: training India’s future on Western data is like trying to simulate the Vedas on a typewriter.
India needs its own LLMs — not just translated, but native-born. Models that grew up on Indian text, learned from Indian discourse, and internalized Indian ways of thinking. Models that understand that “acha” can mean yes, no, maybe, or “I heard you, but I’m still deciding.” Models that don’t treat Hindi and Tamil and Bengali as “low-resource” languages, but as primary operating systems of thought.
This isn’t just a matter of linguistic justice. It’s a matter of cognitive sovereignty. Whoever trains the models, trains the minds. The next billion users won’t read philosophy; they’ll chat with it. They won’t learn history; they’ll ask it questions. If the models answering those questions are trained on Western priorities, they’ll inherit Western biases — subtle, invisible, but deeply real.
The danger is quiet. You won’t notice when a model starts preferring certain histories, certain heroes, certain tones of reasoning. You’ll just wake up one day in a digital world that speaks your language but not your soul.
So yes, India needs its own LLMs — not as a nationalist project, but as a cultural survival mechanism. Imagine a model that can code-switch between Sanskrit’s precision and Hinglish’s chaos; one that understands the humor of Mumbai and the poetry of Madurai; one that can translate between languages not just semantically, but philosophically.
The technical challenge is immense. We’ll need curated corpora from local sources, multilingual tokenizers that respect morphology, and training pipelines tuned for code-mixed text. We’ll need to rethink benchmarks — because why should “commonsense reasoning” be defined by American common sense?
But the reward is equally immense: a model that’s not just Indian in data, but Indian in cognition. A system that embodies our way of seeing the world — plural, paradoxical, context-dependent, resilient.
So if we build it right, India’s LLMs won’t just serve India. They’ll expand what AI itself can be. They’ll teach the machines that there are other ways to think — that the human mind isn’t a monolith, but a kaleidoscope.
Because in the end, language isn’t just a tool we use to talk to AI. It’s the substrate through which AI learns what it means to be human. And humanity, in all its accents and contradictions, deserves to be represented.
That’s why India needs its own LLMs. Not to compete with the West. To complete the story of intelligence.
]]>The Mamba model, introduced in the paper Mamba: Linear-Time Sequence Modeling with Selective State Spaces by Albert Gu and Tri Dao, represents a significant advancement in sequence modeling. Unlike traditional transformer-based architectures that scale quadratically with sequence length, Mamba achieves linear-time complexity using selective state space models (SSMs). This blog post explores the technical depth of implementing Mamba in PyTorch, as provided in the given model.py and scan.py files, breaking down the architecture, key components, and challenges encountered in the implementation.
State Space Models (SSMs) provide a framework for modeling sequences by representing them as a system of differential equations. Traditionally, SSMs like S4 (Structured State Space Sequence Models) use fixed parameters, limiting their adaptability to input data. Mamba introduces selective SSMs, where parameters like ( \Delta ), ( B ), and ( C ) are input-dependent, enabling the model to adapt dynamically to the sequence context. This selectivity, combined with a linear-time scanning mechanism, makes Mamba both efficient and expressive for long-sequence modeling tasks.
The implementation provided is a minimal, single-file PyTorch version of Mamba, designed for clarity and ease of understanding. It includes the core components of the model, such as the Mamba block, selective scan, and RMS normalization, while also supporting pretrained model loading from HuggingFace.
ModelArgsThe ModelArgs dataclass defines the hyperparameters for the Mamba model, such as:
d_model: The hidden dimension of the model.n_layer: The number of Mamba layers.vocab_size: The size of the vocabulary for the embedding layer.d_state: The latent state dimension (( N ) in the Mamba paper).expand: The expansion factor to compute the inner dimension (( d{\text{inner}} = d{\text{model}} \times \text{expand} )).dt_rank: The rank of the input-dependent step size (( \Delta )), with an “auto” option to set it as ( \lceil d_{\text{model}} / 16 \rceil ).d_conv: The kernel size for the 1D convolution.scan_mode: The mode for the selective scan algorithm (cumsum or logcumsumexp).The __post_init__ method ensures that d_inner and dt_rank are computed appropriately and that vocab_size is padded to a multiple of pad_vocab_size_multiple for efficient processing.
The Mamba class encapsulates the full model, consisting of:
nn.Embedding layer maps input tokens to dense vectors of size d_model.n_layer ResidualBlock modules, each containing a normalization layer and a MambaBlock.RMSNorm layer applied before the final output projection.lm_head) that projects the hidden states back to the vocabulary size, with tied weights to the embedding layer for efficiency (as described in the “Weight Tying” paper).The forward method processes input token IDs through the embedding layer, applies the residual blocks, normalizes the output, and projects it to logits over the vocabulary.
The ResidualBlock wraps a MambaBlock with a normalization layer (RMSNorm) and a residual connection. The forward pass follows the structure: [ \text{output} = \text{MambaBlock}(\text{RMSNorm}(x)) + x ] This differs from the official Mamba implementation, which uses a fused Add-Norm-Mamba structure for performance. The provided implementation prioritizes simplicity and numerical equivalence, applying normalization before the Mamba block and adding the residual connection afterward.
The MambaBlock is the core of the model, implementing the selective SSM as described in Figure 3 of the Mamba paper. Its components include:
in_proj) projects the input from d_model to 2 * d_inner, splitting the output into x and a residual (res).conv1d) with kernel size d_conv processes the input sequence, introducing local context.x_proj linear layer generates input-dependent parameters ( \Delta ), ( B ), and ( C ).dt_proj layer projects ( \Delta ) from dt_rank to d_inner.A_log (logarithm of the state transition matrix) to ensure stability, and ( D ) is a learnable vector.ssm method computes the state space dynamics using the selective_scan function.out_proj) maps the SSM output back to d_model.The forward pass applies the convolution, SiLU activation, the selective SSM, and a gated residual connection before the final projection.
scan.py)The selective_scan function implements Algorithm 2 from the Mamba paper, computing the SSM output given input ( u ), step size ( \Delta ), and parameters ( A ), ( B ), ( C ), and ( D ). It supports two modes:
The provided scan.py includes a test script that compares the two modes, showing that the mismatch between them grows with sequence length due to numerical precision differences.
The RMSNorm class implements Root Mean Square Normalization, a variant of layer normalization that scales the input by the inverse of its root mean square value. It is used in both the ResidualBlock and the final normalization step of the Mamba model.
The from_pretrained method in the Mamba class enables loading pretrained weights from HuggingFace models (e.g., state-spaces/mamba-2.8b). It loads the model configuration and state dictionary, mapping weights to the corresponding model parameters while handling differences in naming conventions (e.g., removing the backbone. prefix).
Implementing Mamba in PyTorch involves several challenges:
logcumsumexp mode addresses this by operating in log space, but it introduces additional complexity.Mamba, ResidualBlock, MambaBlock, and RMSNorm. This makes it easier to understand and extend compared to a monolithic implementation.scan_mode parameter allows switching between cumsum and logcumsumexp, enabling experimentation with numerical stability trade-offs.from_pretrained method facilitates integration with pretrained models, allowing researchers to leverage existing weights for fine-tuning or evaluation.scan.py to cover edge cases (e.g., very long sequences or extreme values) would ensure robustness.The provided Mamba implementation is a concise yet powerful demonstration of selective state space modeling in PyTorch. By breaking down the architecture into modular components and addressing numerical stability, it offers a clear entry point for researchers and practitioners interested in exploring linear-time sequence models. The ability to load pretrained weights further enhances its utility, making it a valuable tool for advancing research in efficient sequence modeling.
For further reading, refer to: