---
license: other
license_name: mixed
tags:
- interpretability
- sparsity
- swiglu
- top-k
---

# 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_dict`s (`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](https://huggingface.co/LiquidAI/LFM2.5-230M-Base), so they inherit the **LFM Open License v1.0**. See the base model's [LICENSE](https://huggingface.co/LiquidAI/LFM2.5-230M-Base/blob/main/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.