This is a version of
ModernBERT-base distilled down to 16 layers out of 22.
This reduces the number of parameters from 149M to 119M; however, practically speaking, since the embedding params
do not contribute greatly to latency, the effect is reducing the "trunk" of the model from 110M params to 80M params.
I would expect this to reduce latency by roughly 25% (increasing throughput by roughly 33%).
The last 6 local attention layers were removed:
Unfortunately the HuggingFace modeling code for ModernBERT relies on global-local attention patterns being uniform throughout the model,
so loading this bad boy properly takes a bit of model surgery. I hope in the future that the HuggingFace team will update this
model configuration to allow custom striping of global+local layers. For now, here's how to do it:
1import torch.nn as nn
2from transformers import AutoTokenizer, AutoModelForMaskedLM
3
4model_id = "answerdotai/ModernBERT-base"
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForMaskedLM.from_pretrained(model_id)
1layers_to_remove = [13, 14, 16, 17, 19, 20]
2model.model.layers = nn.ModuleList([
3 layer for idx, layer in enumerate(model.model.layers)
4 if idx not in layers_to_remove
5])
1state_dict = torch.load("model.pt")
2model.model.load_state_dict(state_dict)
This model was distilled from ModernBERT-base on the
MiniPile dataset,
which includes English and code data. Distillation used all 1M samples in this dataset for 1 epoch, MSE loss on the logits,
batch size of 16, AdamW optimizer, and constant learning rate of 1.0e-5.
The embeddings/LM head were frozen and shared between the teacher and student; only the transformer blocks were trained.
I have not yet evaluated this model. However, after the initial model surgery, it failed to correctly complete
"The capital of France is [MASK]", and after training, it correctly says "Paris", so something good happened!