nlpie/modernalbert-tiny-v1.0

ModernALBERT-Tiny

ModernALBERT-Tiny is a compact, recursive transformer for natural language understanding. It combines ALBERT-style cross-layer parameter sharing with Mixture of LoRAs (MoL) — a lightweight, token-conditional routing mechanism that restores the expressivity normally lost when transformer layers share weights — plus a set of modern architectural upgrades (RoPE, GeGLU, FlashAttention, Pre-Norm).

It is well suited to text classification, natural language inference, paraphrase/semantic-similarity detection, extractive QA, and dense retrieval, and is the smallest, fastest model in the family.

Model repo: nlpie/modernalbert-tiny-v1.0

Other sizes in this family: tiny · medium · base · large

---

Table of Contents

1. Overview
2. Model Architecture
3. How to Use
4. Training & Dataset
5. GLUE Benchmark Results
6. SQuAD-v2 Results
7. Inference Efficiency
8. Key Features and Design Choices
9. Limitations
10. Citation

---

Overview

ModernALBERT builds on ALBERT's cross-layer parameter sharing, which reduces model size but can cap representational capacity when layers are fully tied. ModernALBERT addresses this with:

• Mixture of LoRAs (MoL): low-rank LoRA "experts" injected directly into the weights of the shared feed-forward network, with sparse router-driven activation, simulating a Mixture-of-Experts layer at a fraction of the parameter cost.
• Modern architecture: Pre-Norm, GeGLU activations, rotary position embeddings (RoPE), and FlashAttention (with an automatic PyTorch SDPA fallback).
• Distillation-based initialisation: weights are seeded from a fully-parameterised ModernBERT teacher via layer-mapped initialisation, and training uses knowledge distillation from that same teacher — critical for reaching strong performance on a comparatively small pretraining budget.

ModernALBERT-Tiny is the smallest variant: 14 layers organised into 7 shared groups (a shorter group depth of 2, unlike the other variants), with a smaller 768 hidden size, and a lighter MoL configuration (4 experts, top-1 routing instead of 8 experts / top-2). Despite this, it outperforms similarly-sized baselines like MiniLM on GLUE.

---

Model Architecture

Variant	Layers	Groups	MoL Groups	Hidden Dim	FFN Intermediate Dim	Expert (LoRA) Dim	Experts	Top-K
Tiny	14	7	6, 7	768	1152	2624	4	1

• Parameter sharing: layers are grouped in pairs (group depth = 2); all layers within a group share attention and FFN weights, so the model behaves like a 14-layer network while storing far fewer unique parameters.
• Mixture of LoRAs: the last two groups replace the shared FFN with a router over 4 low-rank LoRA experts with top-1 routing — a lighter configuration than the other ModernALBERT variants, matched to Tiny's smaller hidden size.
• Attention: rotary embeddings for position information, FlashAttention (unpadded/varlen) when available, otherwise scaled-dot-product attention.
• Embeddings: ALBERT-style factorised token embeddings (small embedding dimension projected up to the hidden size), reducing the size of the embedding matrix.
• ~51M parameters total (paper-reported figure: 50M).

---

How to Use

ModernALBERT ships with custom transformers-compatible modeling code (ModernALBERTConfig, ModernALBERTModel, ModernALBERTForMaskedLM, ModernALBERTForSequenceClassification, ModernALBERTForQuestionAnswering). Load it with trust_remote_code=True. It runs on both CPU and GPU: FlashAttention is used automatically when installed, and the model falls back to PyTorch's built-in SDPA attention otherwise — no extra configuration needed either way.

pip install transformers torch
# Optional, for the fastest attention path on supported GPUs (auto-detected; falls back to SDPA if absent):
pip install flash-attn --no-build-isolation

Masked language modeling

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

model_id = "nlpie/modernalbert-tiny-v1.0"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(model_id, trust_remote_code=True)
model.eval()

text = f"Paris is the capital of {tokenizer.mask_token}."
inputs = tokenizer(text, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

mask_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
predicted_id = outputs.logits[0, mask_index].argmax(dim=-1)
print(tokenizer.decode(predicted_id))

Sentence / token embeddings

from transformers import AutoTokenizer, AutoModel
import torch

model_id = "nlpie/modernalbert-tiny-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
model.eval()

inputs = tokenizer("Example sentence for embeddings.", return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

attention_mask = inputs["attention_mask"]
last_hidden = outputs.last_hidden_state
# Mean pooling over valid tokens
embedding = (last_hidden * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(1, keepdim=True)

Fine-tuning for sequence classification

from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer

model_id = "nlpie/modernalbert-tiny-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(
    model_id, trust_remote_code=True, num_labels=2
)

# tokenized_train_dataset, tokenized_eval_dataset = ...  # your tokenized datasets

training_args = TrainingArguments(
    output_dir="./modernalbert-tiny-finetuned",
    per_device_train_batch_size=16,
    num_train_epochs=3,
    learning_rate=2e-5,
    eval_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train_dataset,
    eval_dataset=tokenized_eval_dataset,
)
trainer.train()

Extractive question answering

from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch

model_id = "nlpie/modernalbert-tiny-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForQuestionAnswering.from_pretrained(model_id, trust_remote_code=True)

question, context = "What does MoL stand for?", "ModernALBERT introduces Mixture of LoRAs (MoL), a routing mechanism over low-rank experts."
inputs = tokenizer(question, context, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

start = outputs.start_logits.argmax()
end = outputs.end_logits.argmax() + 1
print(tokenizer.decode(inputs["input_ids"][0][start:end]))

If Auto-class mapping isn't yet wired up on the Hub

If trust_remote_code=True doesn't resolve the classes automatically, either import the classes directly from the repo's Python files, or add an auto_map block to config.json:

{
  "auto_map": {
    "AutoConfig": "configuration_modernalbert.ModernALBERTConfig",
    "AutoModel": "modeling_modernalbert.ModernALBERTModel",
    "AutoModelForMaskedLM": "modeling_modernalbert.ModernALBERTForMaskedLM",
    "AutoModelForSequenceClassification": "modeling_modernalbert.ModernALBERTForSequenceClassification",
    "AutoModelForQuestionAnswering": "modeling_modernalbert.ModernALBERTForQuestionAnswering"
  }
}

Compatibility

This code has been verified end-to-end — building the model, running a forward pass, and a save_pretrained → from_pretrained round trip with bit-identical weights — on CPU (PyTorch's SDPA attention path), and structurally validated against the FlashAttention code path. It targets a recent transformers release (tested against 5.x); on much older transformers versions you may need to upgrade, since the weight-tying and rotary-embedding buffer conventions it relies on changed across versions.

Efficient (merged) inference

The routing_strategy field in ModernALBERTConfig controls how the MoL layer behaves:

• "standard" (default): sparse top-1 token routing (Tiny uses top-1, unlike the top-2 routing in the other variants), as used during pretraining.
• "uniform": experts are averaged with equal weight into a single static LoRA adapter (no per-token routing), matching the "Vanilla" merge strategy in the paper.
• "ema": experts are merged using an exponential moving average of the router's historical activations, matching the paper's dynamic EMA-merging strategy — this recovers accuracy close to the unmerged model while removing routing overhead at inference time (see Inference Efficiency).

---

Training & Dataset

• Corpus: two-stage curriculum — warm-up on RedPajama-1T (20k–30k steps), then continued training on RefinedWeb (70k–80k further steps).
• Budget: ~30B tokens total, versus 1.7T tokens for the ModernBERT teacher.
• Initialisation: step-wise, layer-mapped initialisation from a fully-parameterised ModernBERT teacher.
• Distillation: ModernBERT's predictions are used as soft targets alongside the MLM objective.
• Optimisation: AdamW, global batch size 384, max sequence length 1024, linear warmup to a peak learning rate of 5×10⁻⁴ or 5×10⁻⁵, followed by linear decay.

---

GLUE Benchmark Results

Task Category	Task	Score
Single Sentence	CoLA	58.4
Single Sentence	SST-2	93.0
Paraphrase / Similarity	MRPC	90.2
Paraphrase / Similarity	STS-B	90.4
Paraphrase / Similarity	QQP	90.5
Natural Language Inference	MNLI	84.6
Natural Language Inference	QNLI	91.3
Natural Language Inference	RTE	81.2
Average		84.95

This is the strongest reported result among compact ~30–70M-parameter models in the paper, ahead of MiniLM-33M (84.49 avg), DistilBERT, and TinyBERT.

---

SQuAD-v2 Results

Metric	Score
F1	90.0
Exact Match	82.9

For BEIR retrieval results, see the ModernALBERT-Large model card, where a subset of BEIR datasets is reported in the paper.

---

Inference Efficiency

Latency and throughput measured with and without the expert-merging procedure described in the paper (batch inference, single GPU):

Model	Latency (ms) ↓	Throughput (tok/s) ↑	Memory (GB) ↓
ModernALBERT-tiny, no merging	13.40	72,810	0.196
ModernALBERT-tiny, merged experts	9.46	106,527	0.196

Merging collapses the dynamic MoL router into a single static LoRA adapter at deployment time (see Efficient (merged) inference above), cutting latency substantially while keeping the same memory footprint and most of the accuracy gains from routing.

---

Key Features and Design Choices

• Compact but flexible: parameter sharing keeps the model very small; MoL restores per-token expressivity with a lightweight pool of low-rank experts.
• Conditional computation: a single (top-1) expert activates per token in MoL layers, keeping compute minimal.
• Modern training stack: Pre-Norm, GeGLU, RoPE, and FlashAttention (with SDPA fallback) for training stability and speed.
• Distillation-based warm start: initialised and distilled from ModernBERT, enabling strong results on a fraction of ModernBERT's pretraining token budget.
• Deployment-friendly: an optional expert-merging step (routing_strategy="ema" or "uniform") compresses MoL into a single dense adapter for even lower-latency inference.

---

Limitations

• MoE-style routing still carries more computational overhead than a fully dense/shared model, even with the merging optimisation; further work on load balancing and expert selection could close this gap.
• The model uses global attention only (no local/sliding-window attention), so it may underperform on tasks requiring very long context or fine-grained long-range reasoning compared to architectures with hybrid attention patterns.
• As the smallest variant (768 hidden size, 4 experts, top-1 routing), Tiny trades some accuracy for size and speed relative to Medium, Base, and Large.
• Benchmark numbers above are as reported in the accompanying paper; results can vary with fine-tuning setup, hardware, and library versions.

---

Citation

This model accompanies the paper "Improving Recursive Transformers with Mixture of LoRAs" (currently under anonymous review). Formal citation details will be added once the paper is published — check back on this model card or the paper's repository for an updated BibTeX entry.