abl_pure_mamba

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

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

Evaluation Method

Perplexity (primary metric)

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

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

Rank Summary

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

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)