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

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

Placement — N/A

All 12 layers are Mamba layers.

Training

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).

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

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()