mask-align checkpoints
Checkpoints from a study of top-K masks on SwiGLU gates and the V2 similarity-alignment loss. V2 pushes tokens whose MLP outputs are similar to also share similar sets of active gate units. The aim is to make individual units and groups of units easier to interpret.
All files are raw PyTorch state_dicts (torch.load, pickle). Load only if you trust this repo.
d10/: nanochat depth-10 models trained from scratch
nanochat (commit 92d63d4, patched nanochat/gpt.py and scripts/base_train.py): 10 layers, d_model 640, MLP hidden size 1792, vocab 16384, sequence length 2048, 1402 steps × 524k tokens on ClimbMix. topk_* models keep the top 25% of gates per token (topk_frac 0.25) and add λ · (1 − V2) to the loss. Each folder has model_001402.pt and meta_001402.json (full config). d10/tokenizer/ is the nanochat tokenizer.
| Model | λ (V2) | val bpb |
|---|---|---|
| swiglu (dense baseline) | – | 0.913 |
| topk_l0 | 0 | 0.943 |
| topk_l0p1 | 0.1 | 0.947 |
| topk_l0p3 | 0.3 | 0.953 |
| topk_l1 | 1 | 0.972 |
| topk_l3 | 3 | 0.999 |
| topk_l10 | 10 | 1.033 |
lfm_posthoc/: post-hoc sparsified LFM2.5-230M-Base
These are fine-tunes of LiquidAI/LFM2.5-230M-Base, so they inherit the LFM Open License v1.0. See the base model's LICENSE. The weights were modified by the training described below.
In layers 2–13 the MLP keeps the top 25% of units by |gate| (silu·u values are kept; layers 0–1 stay dense). Each model is trained for 20M ClimbMix tokens by KL distillation from the frozen dense model, plus λ · (1 − V2). The V2 target is relu(cosine similarity) of the teacher's batch-centred MLP outputs, and the mask is sign-split. Settings: lr 5e-5, batch 8×1024. bpb is measured on 64 held-out documents: the dense model scores 0.835, and the sparsified model before training scores 1.116.
| Model | λ | bpb | V2 |
|---|---|---|---|
| lam0 | 0 | 0.901 | 0.525 |
| lam10 | 10 | 0.949 | 0.743 |
| lam30 | 30 | 1.009 | 0.757 |
student.pt is a full LFM2ForCausalLM state dict in fp32. Load it into AutoModelForCausalLM.from_pretrained("LiquidAI/LFM2.5-230M-Base"), then patch the sparse MLPs: see Sparse / load_student in lfm_posthoc.py. Without the patch the model runs dense, which these weights were not trained for. log.json is the training log.