arianraje/qwen3-4b-mamba3-hybrid-stage2a-kd-bias1

qwen3-4b-mamba3-hybrid-stage2a-kd-bias1

Stage-2a endpoint, bias-1 arm -- identical recipe to -stage2a-kd-bias0 (600M tokens, 287 steps, LR 2e-4 / 2e-5, 5 h 58 m on 4x B200) from the bias-1 stage-1 checkpoint.

Final val KL 0.1258, top1-agreement 0.8601, wikitext-2 PPL 12.41, from a starting point of 0.5058 / 0.7465 / 24.99. Led the bias-0 arm at every one of the 29 evaluations, finishing 3% lower in KL; the best teacher agreement of any Qwen arm at this stage (GDN 0.1437 / 0.8515, Mamba2 0.1819 / 0.8295). Input to stage-2b long-context KD.

Loading

This is a custom container (model_type: qwen3_mamba3), not a stock architecture, so trust_remote_code=True is required. The modeling code is bundled in this repo (modeling_qwen3_mamba3.py) and needs no other checkout:

from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained(REPO, trust_remote_code=True, dtype="bfloat16")
tok = AutoTokenizer.from_pretrained(REPO)

Kernel. The mixers call upstream's fused Mamba-3 MIMO kernel, mamba_ssm.ops.tilelang.mamba3.mamba3_mimo, from state-spaces/mamba at commit e9594ce (TileLang + Triton; the pip mamba_ssm 2.2.6.post3 predates Mamba-3, so install from source). If that import fails the model warns once and every mixer runs mamba3_mimo_reference, a sequential fp32 torch transcription of the kernel's semantics: numerically equivalent (fwd rel err 6e-3 vs the kernel in bf16), and slow -- fine for a decode-parity check, not for throughput. Set Qwen3Mamba3Mixer.force_reference = True to opt into it explicitly.

Training with the kernel. Upstream's mamba3_mimo.py backward at e9594ce reads ctx.saved_tensors twice, which PyTorch's non-reentrant activation checkpointing (use_reentrant=False, HF's default) rejects with CheckpointError: Unpack is being triggered for a tensor that was already unpacked once. The fix is a one-line change (read it once); the training runs below used a local commit with that change. Inference is unaffected.

The study

Full-attention Qwen3-4B converted to a Mamba-3 hybrid: 27 of 36 layers have their attention replaced by a Mamba-3 MIMO mixer, with every 4th layer (indices 3, 7, 11 ... 35) retained as full attention -- the u4 pattern. Retained layers are literal Qwen3DecoderLayers, so retention is bit-exact. Mixer geometry: 32 heads x 128 headdim, 8 B/C groups, `dstate 128, **MIMO rank 4**, expand 1.6, chunk 8, RoPE on half of the state dims, the pre-gate grouped RMSNorm fused into the kernel. Parameter layout and math are upstream's exactly (inproj` rows `[z | x | B | C | dddt | dd_A | trap | angles], rank-major (r g n)` B/C), so the upstream kernel is called unchanged. 4,741,428,416 params, of which the 27 mixers are 1,426,756,032 (30.1%).

Capability is recovered by staged distillation rather than pretraining: surgery/init -> stage-1 per-layer hidden-state alignment (100M tok) -> stage-2a forward-KL KD @4k (600M) -> stage-2b long-context KD @32k (294M) -> stage-3 on-policy distillation. Teacher: Qwen/Qwen3-4B. This arm exists to isolate the mixer: it shares the teacher, the packed data, the token budgets and the schedules with the Gated DeltaNet and Mamba2 arms trained identically.

Inherited init

The mixer starts as the teacher's attention head, not as a random SSM (inherited fraction 99.95%): per converted layer, x head h <- v_proj head h // 4; B[r, g] <- k_proj head g; C[r, g] <- q_proj head 4g + r; out_proj <- o_proj; C_norm / B_norm <- q_norm / k_norm. The dd_dt, dd_A, trapezoid and angle rows of in_proj are zero (per-head constant dt, A = -1, trapezoid 1/2, no rotation; all keep gradient). dt_bias is the softplus-inverse of exp(U(log 1e-3, log 1e-1)) (the repaired dt init of the Mamba2 arm; horizons of 10-1000 tokens). mimo_x is a random simplex over ranks per (head, dim) so the ranks get distinct gradients while the init state equals upstream's uniform 1/R; mimo_o selects each head's own rank; mimo_z, D and the norm weights are 1. At init a converted head equals the teacher's head with softmax replaced by causal decay (tested to 3e-8).

B/C bias. Upstream initializes the B/C biases to 1.0. On an inherited head that adds sum(q) + sum(k) + 128 to every score, which looked like it would swamp the teacher's structure, so the primary arm used 0 and a second arm kept upstream's 1.0. Measured: bias 1 led at every stage-1 and stage-2a evaluation (below). Both arms are published; bias 1 is now the default.

Results so far (trainer metrics only)

Val KL is the forward KL to the teacher over the full vocabulary on held-out packed text at the stage's sequence length; PPL is wikitext-2 at 4096 context. Mamba-3 enters stage 2a at KL 0.50 (Mamba2: 1.08) and passes Mamba2's final KL by step 20 of 287.

Not yet evaluated. No downstream battery (commonsense, RULER needle-in-a-haystack, GSM8K / MATH / AIME, MMLU) has been run on any Mamba-3 checkpoint; the numbers above are the trainers' own validation metrics. Stage 2b (32k KD) is complete for both arms. Stage 3 (on-policy distillation) has not been run. Intermediate checkpoints, published for the ladder.