CoolFace
Modelpublic

crumb-ai/abl_2_1_interleaved

sourceHugging Facemitupdated 2mo agoView on Hugging Face
0likes3downloads
Model Card

CRUMB abl_2_1_interleaved ⭐ Best in Phase-1 Ablation

Model Overview

abl_2_1_interleaved is the best-performing model in the CRUMB Phase-1 ablation. It is a hybrid decoder-only language model that interleaves Mamba-3 selective state-space layers with GQA (Grouped-Query Attention) layers at a 2:1 Mamba-majority ratio with interleaved placement, and was pre-trained exclusively on Python source code.

Across the eleven ablation variants, this configuration achieves the lowest evaluation perplexity (3.4182), confirming that a moderate Mamba majority (67 % of layers) with attention distributed throughout the stack is the optimal design choice at the 150M-parameter scale.

Architecture

PropertyValue
Total parameters146,075,776 (146.1M)
d_model768
n_layers12
n_heads12
n_kv_heads4 (GQA)
d_head64
d_ff3072
vocab_size32768
seq_len4096
Tie embeddingsyes
Pos. encodingRoPE (base = 10000)
Mamba layer typeMamba-3 (d_state=64, expand=2, headdim=64, ngroups=1, chunk=64)

Mamba : Attention ratio — 2 : 1

8 Mamba layers + 4 GQA attention layers (2 Mamba blocks per attention block).

Placement — Interleaved

Mamba and attention layers alternate throughout the network according to a 2:1 schedule. Layer order: M A M M A M M A M M A M

This placement distributes the attention layers as evenly as the 2:1 ratio allows while keeping every attention layer close to a Mamba layer that provides the broad sequential context.

Training

PropertyValue
Training dataPython subset of bigcode/the-stack-dedup-v2
Tokens seen5,367,398,400 (~5.37 B)
Steps163,840
Context length4096
Training time47 h 24 m 43 s
Final learning rate3.00e-05

Evaluation Method

Perplexity (primary metric)

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

SettingValue
Eval sequences20,063 batches
Eval tokens328,631,940
Implementationsrc/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

MetricValue
Eval loss1.2291 ⭐ (best)
Eval perplexity3.4182 ⭐ (best)
Eval time3,175.94 s (~53 min)
Syntax validity (n=200)55 / 200 → 27.5 %
Inference gen. time (200×128 tok)180.33 s

Rank Summary

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

RankModelPerplexity
1`abl_2_1_interleaved`3.4182
2abl_3_1_interleaved3.4359
3abl_3_1_backloaded3.4493
4abl_2_1_backloaded3.4683
5abl_1_1_backloaded3.4763
6abl_pure_mamba3.5237
7abl_1_1_interleaved3.5407
8abl_pure_attn3.5939
9abl_3_1_frontloaded3.6798
10abl_2_1_frontloaded3.7078
11abl_1_1_frontloaded3.7315

abl_2_1_interleaved is recommended as the base configuration for Phase 2 of the CRUMB project. The overall spread across all 11 configurations is 0.314 perplexity points (9.2 %); this model is 0.176 PPL (4.9 %) better than the worst model (abl_1_1_frontloaded) and 0.105 PPL (3.0 %) better than the best pure baseline (abl_pure_mamba).

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

python
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_2_1_interleaved.yaml")
model = CRUMBModel(config)
state = torch.load("saved/model/abl_2_1_interleaved/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)