| --- |
| 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. |
| |
