modernbert-base-yat-mlp-swap

ModernBERT-base Yat MLP Swap

This repo contains Yat student replacements for all 22 MLP blocks of answerdotai/ModernBERT-base.

The checkpoints were trained on Kaggle 2xT4 with a two-phase pipeline:

1. Phase 1: data-free random/on-shell distillation. Random Gaussian probes are projected through each layer's mlp_norm, then each Yat block learns the corresponding ModernBERT MLP delta map.
2. Phase 2: real-activation fine-tuning. Real mlp_norm activations are collected from Wikitext text, then every Phase-1 Yat block is fine-tuned on its layer's real activations.

The model is distributed as a lightweight patch over the base ModernBERT weights. yat_modernbert.py downloads answerdotai/ModernBERT-base, loads the Phase-2 Yat checkpoints, and replaces every model.model.layers[i].mlp module.

Results

Phase 1 mean random/on-shell rho: 0.281783

Phase 2 mean real-activation rho: 1.654114 -> 0.074821

Usage

from yat_modernbert import load_model

model, tokenizer = load_model("azettaai/modernbert-base-yat-mlp-swap")
inputs = tokenizer("Yat blocks replace every ModernBERT MLP.", return_tensors="pt")
inputs = inputs.to(next(model.parameters()).device)
out = model(**inputs)
print(out.logits.shape)

Files

• phase2/block0.safetensors ... phase2/block21.safetensors: final Yat MLP replacements.
• phase1/: random/on-shell warm-start checkpoints.
• scripts/: Kaggle scripts used to train and bundle the model.
• yat_modernbert.py: loader and patching code.
• manifest.json: bundle metadata copied from the Kaggle artifact.