CRUMB abl_pure_mamba

Model Overview

abl_pure_mamba is a pure Mamba baseline for the CRUMB ablation: 12 Mamba-3 selective state-space layers and zero attention layers. It is pre-trained exclusively on Python source code.

This model serves as the second architectural control (alongside abl_pure_attn) against which the hybrid variants are measured. It is the largest model in the ablation (154.6M parameters) because pure Mamba layers carry more parameters per layer than pure attention layers at the same d_model (due to the input-dependent SSM state matrices and the expand=2 projection).

Architecture

Property	Value
Total parameters	154,616,640 (154.6M)
d_model	768
n_layers	12
n_heads	12 (unused, kept for config compatibility)
n_kv_heads	4 (unused)
d_head	64 (unused)
d_ff	3072
vocab_size	32768
seq_len	4096
Tie embeddings	yes
Pos. encoding	RoPE (base = 10000)
Mamba layer type	Mamba-3 (d_state=64, expand=2, headdim=64, ngroups=1, chunk=64)
Attention layers	0 (pure Mamba baseline)

Mamba : Attention ratio — 12 : 0 (pure Mamba)

Placement — N/A

All 12 layers are Mamba layers. Layer order: M M M M M M M M M M M M

Training

Property	Value
Training data	Python subset of bigcode/the-stack-dedup-v2
Tokens seen	5,367,468,015 (~5.37 B)
Steps	163,840
Context length	4096
Training time	49 h 11 m 44 s (longest of the 11)
Final learning rate	3.00e-05
Peak GPU memory	6,089 MB
Training throughput	30,409 tok/s (slowest of the 11)

Evaluation Method

Perplexity (primary metric)

Per-token cross-entropy loss with BF16 autocast, computed over the full held-out evaluation set.

Setting	Value
Eval sequences	20,063 batches
Eval tokens	328,631,940
Implementation	src/evaluation/perplexity.py

Generation-based metrics

• Python syntax validity — 200 free-form completions generated per model from 49 diverse Python prompts at temperature=0.8, top_k=50, max_new_tokens=128; each completion checked with ast.parse(). Implementation: src/evaluation/syntax_validity.py.
• Qualitative side-by-side completions — 10 fixed prompts at temperature=0.6, top_k=50, max_new_tokens=200, identical random seed per prompt. Implementation: src/evaluation/qualitative_comparison.py.

Evaluation Results

Metric	Value
Eval loss	1.2595
Eval perplexity	3.5237
Eval time	3,144.42 s (~52 min)
Syntax validity (n=200)	52 / 200 → 26.0 %
Inference gen. time (200×128 tok)	193.02 s (slowest inference)

Rank Summary

Out of 11 ablation configurations evaluated at the same token budget:

Rank	Model	Perplexity
1	abl_2_1_interleaved	3.4182
2	abl_3_1_interleaved	3.4359
3	abl_3_1_backloaded	3.4493
4	abl_2_1_backloaded	3.4683
5	abl_1_1_backloaded	3.4763
6	abl_pure_mamba	3.5237
7	abl_1_1_interleaved	3.5407
8	abl_pure_attn	3.5939
9	abl_3_1_frontloaded	3.6798
10	abl_2_1_frontloaded	3.7078
11	abl_1_1_frontloaded	3.7315

abl_pure_mamba ranks 6th overall — the best pure baseline but still beaten by all five backloaded and interleaved hybrid configurations. Despite carrying 20 % more parameters than abl_pure_attn, its perplexity is only 0.07 lower. At seq_len=4096, pure Mamba is also the slowest model to train and infer (30,409 tok/s vs 43,337 tok/s for attention, a 30 % gap).

A separate parameter-equalized pilot (abl_pure_mamba_pilot, d_ff=2904 → 150.0M params, 2.68 B training tokens) reached eval PPL 3.39 and confirms that pure Mamba's disadvantage in the standard ablation is not driven by its larger parameter count.

Intended Use & Limitations

• Domain: Python source-code language modelling.
• Base model only: no instruction tuning, no chat alignment, no safety filtering. Outputs are unconstrained code completions.
• Repetitive degeneration: all base CRUMB models tend to repeat function signatures / docstrings during free-form generation; this is expected behaviour for unaligned base models.

Citation / Context

This model is part of the CRUMB Phase-1 ablation study:

Efficient Architectural Hybrids for Small-Scale Language Models in Python Program Synthesis — Department of Computer Science and Engineering, Daffodil International University. Findings documented in documents/phase1_ablation_findings.md.

How to Load

from tokenizers import Tokenizer
import torch
from src.model.config import CRUMBConfig
from src.model.model import CRUMBModel

config = CRUMBConfig.from_yaml("configs/model/abl_pure_mamba.yaml")
model = CRUMBModel(config)
state = torch.load("saved/model/abl_pure_mamba/model.pt", map_location="cpu")
model.load_state_dict(state)
model.eval()

tok = Tokenizer.from_file("saved/tokenizer/crumb_tok_hf/tokenizer.json")
ids = tok.encode("def fibonacci(n):\n").ids
x = torch.tensor([ids])
with torch.no_grad():
    y = model(x)