CRUMB abl_pure_attn

Model Overview

abl_pure_attn is a pure Transformer baseline for the CRUMB ablation: 12 GQA (Grouped-Query Attention) layers and zero Mamba layers. It is pre-trained exclusively on Python source code.

This model serves as the architectural control against which the Mamba / hybrid variants are measured. It is the smallest model in the ablation (129.0M parameters) because pure attention layers carry fewer parameters per layer than pure Mamba layers at the same d_model.

Architecture

Property	Value
Total parameters	128,994,048 (129.0M)
d_model	768
n_layers	12
n_heads	12
n_kv_heads	4 (GQA)
d_head	64
d_ff	3072
vocab_size	32768
seq_len	4096
Tie embeddings	yes
Pos. encoding	RoPE (base = 10000)
Mamba layers	0 (pure attention baseline)

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

Placement — N/A

All 12 layers are attention layers. Layer order: A A A A A A A A A A A A

Training

Property	Value
Training data	Python subset of bigcode/the-stack-dedup-v2
Tokens seen	5,367,455,730 (~5.37 B)
Steps	163,840
Context length	4096
Training time	34 h 31 m 20 s ⭐ (fastest of the 11)
Final learning rate	3.00e-05
Peak GPU memory	5,432 MB
Training throughput	43,337 tok/s ⭐ (fastest 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.2792
Eval perplexity	3.5939
Eval time	3,064.11 s (~51 min)
Syntax validity (n=200)	50 / 200 → 25.0 %
Inference gen. time (200×128 tok)	121.60 s ⭐ (fastest 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_attn ranks 8th overall — the second-worst pure baseline behind all hybrids and abl_pure_mamba. Despite being the fastest model to train and infer (43,337 tok/s training, 121.6 s for 200 generations), its pure-attention quality lags every hybrid except the frontloaded ones. This is the central efficiency-vs-quality finding of Phase 1: at seq_len=4096, attention's throughput advantage is preserved, but hybrid quality is consistently higher.

A separate parameter-equalized pilot (abl_pure_attn_pilot, d_ff=3832 → 150.1M params, 2.68 B training tokens) reached eval PPL 3.34 and confirms that pure attention's disadvantage is not driven by its smaller 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_attn.yaml")
model = CRUMBModel(config)
state = torch.load("saved/model/abl_pure_attn/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)