nlpie/modernalbert-medium-v1.0

ModernALBERT-Medium

ModernALBERT-Medium 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, in a very small footprint.

Model repo: `nlpie/modernalbert-medium-v1.0`

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

Table of Contents

1. 1.Overview
2. 2.Model Architecture
3. 3.How to Use
4. 4.Training & Dataset
5. 5.GLUE Benchmark Results
6. 6.SQuAD-v2 Results
7. 7.Inference Efficiency
8. 8.Key Features and Design Choices
9. 9.Limitations
10. 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-Medium is a shallow, efficient variant: 12 layers organised into just 3 shared groups, with a single MoL layer at the end of the recursion, aimed at latency- and memory-constrained deployments.

Model Architecture

• —Parameter sharing: layers are grouped (group depth = 4); all layers within a group share attention and FFN weights, so the model behaves like a 12-layer network while storing far fewer unique parameters.
• —Mixture of LoRAs: the final group replaces the shared FFN with a router over 8 low-rank LoRA experts (top-2 routing), letting different tokens activate different experts.
• —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.
• —~54M parameters total (paper-reported figure: 55M).

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-medium-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-medium-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-medium-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-medium-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-medium-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-2 token routing, 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

SQuAD-v2 Results

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):

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 small; MoL restores per-token expressivity with low-rank experts.
• —Conditional computation: only the top-2 experts activate per token in the MoL layer.
• —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 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 shallowest 1024-hidden-dim variant besides Tiny, Medium trades some accuracy (particularly on CoLA and RTE) for a smaller, faster model — see the Base and Large cards if you need higher accuracy.
• —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.