crumb-ai/abl-pure-mamba-pilot

CRUMB abl_pure_mamba_pilot

Model Overview

abl_pure_mamba_pilot is a parameter-equalized pure-Mamba pilot trained as a control for the CRUMB ablation. It is identical to abl_pure_mamba except that its feed-forward dimension has been deflated from d_ff=3072 to d_ff=2904 to compensate for the heavier Mamba layers, bringing the total parameter count to ~150.0 M (vs. 154.6 M for abl_pure_mamba).

It is paired with abl_pure_attn_pilot to test whether the parameter spread (~20 %) across the standard ablation (129M–155M) confounds the Mamba vs. Attention comparison. Both pilots were trained on identical data (2.68 B tokens, 81,920 steps) with the same random seed.

Architecture

Property	Value
Total parameters	~150.0 M (≈ 150,000,000)
d_model	768
n_layers	12
n_heads	12 (unused)
n_kv_heads	4 (unused)
d_head	64 (unused)
d_ff	2904 ← adjusted for parameter equalization (default: 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

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

Placement — N/A

All 12 layers are Mamba layers.

Training

Property	Value
Training data	Python subset of bigcode/the-stack-dedup-v2
Tokens seen	2,683,723,770 (~2.68 B)
Steps	81,920 (half of the standard ablation)
Context length	4096
Training time	24 h 26 m 56 s
Final learning rate	3.00e-05
Peak GPU memory	6,021 MB
Training throughput	~31,000 tok/s

Evaluation Method

Perplexity (primary metric)

Per-token cross-entropy loss with BF16 autocast, computed over the full held-out evaluation set (same eval set used in the standard ablation).

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

Note: this pilot's results are not included in the headline saved/eval/perplexity.json table — only the eleven standard-ablation models are. Pilot eval metrics are reported in documents/phase1_ablation_findings.md (Section 5).

Evaluation Results

Metric	Value
Eval PPL (final, 2.68 B tokens)	3.39
Final train loss (smoothed, last 10 %)	1.213
Final eval loss	1.220
Wall clock	~24.4 h

Why this pilot exists

The standard ablation's parameter counts vary from 129.0 M (pure attention) to 154.6 M (pure Mamba) — a 20 % range caused by Mamba layers being heavier than attention layers at the same d_model. To determine whether this parameter gap, rather than the architectural difference, drives the standard-ablation rankings, both pilots were trained at exactly 150 M parameters by adjusting d_ff:

• abl_pure_attn_pilot: d_ff=3832 (inflated) → 150.1 M
• abl_pure_mamba_pilot: d_ff=2904 (deflated) → 150.0 M

Result: at equal parameters, attention reaches eval PPL 3.34 vs. Mamba's 3.39 — a small but consistent gap. This confirms that the standard-ablation rankings are not an artefact of parameter count. See documents/phase1_ablation_findings.md Section 5 for the full analysis.

Citation / Context

This pilot 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 Section 5.

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_pilot.yaml")
model = CRUMBModel(config)
state = torch.load("saved/model/abl_pure_mamba_pilot/model.pt", map_location="cpu")
model.load_state_dict(state)
model.eval()