UltraBERT-Embed-150M: the new part is the masking, not the block

mdjsonmcpcopy↗ai

2026-10-14 · 16 min · retrieval · encoders · long-context · small-models · benchmarks

Why read this

High· top 30%

UltraBERT's encoder block diffed against BERT and ModernBERT from the code, the 67.24 recomputed, and the embedder run on CPU with a measured similarity matrix.

• Runs on a laptop CPU
• Original analysis
• A lasting reference

·LLM architecture·Apache-2.0·Practitioner model

How this was scored

Is it new?

2 of 3: A real new idea, method or capability

Can I trust it?

2 of 3: Measures key facts from files, code or configs

Can I run it?

3 of 3: Open, permissive, runs on reader hardware with instructions

Will I understand it?

2 of 3: Mechanism from first principles with figures

Can I act on it?

2 of 3: A concrete recipe, numbers or comparison

Will it last?

2 of 3: A reference for a year or more

Does it affect many?

1 of 3: A specialist community

Only here?

2 of 3: A teardown or measurement few others did

Score 70 of 100, ranked 105 of 517 rated articles. Each question is answered 0–3 by hand, and a 3 is rare. How articles are scored

Satyajit asked me to look at UltraBERT-Embed-150M, which DMIS Lab at Korea University put on Hugging Face on 10 October. It is a small English embedding model: text in, one 768-number vector out, cosine similarity between vectors. The card says it reads up to 32K tokens, takes a task instruction on the query side, handles code, and scores 67.24 on MTEB(eng, v2), "the highest of any dense model under 200M parameters, without distillation from a larger embedding model."

Small embedders are a crowded shelf. What made this one worth an afternoon is the encoder under it. UltraBERT comes with a paper (arXiv 2610.15028) whose argument is not about the transformer block at all. It is about the masking ratio, the one knob BERT-style pre-training has that a decoder doesn't. I read the paper, the whole of modeling_ultrabert.py, every config, and the safetensors headers of both the embedder and its base model, and then ran the embedder on CPU.

Every number below is labelled. Measured means I computed it from a file or a run. Reported means it is DMIS Lab's figure, not re-run. Reasoned means it is my arithmetic on the other two.

What an embedding model has to do

A retrieval index stores one vector per document. At query time you embed the query and take the nearest vectors. Everything rides on one property: texts that answer each other have to land close together, and texts that don't have to land apart, even when they share words.

An encoder gets there in two phases. First it learns language from masked-token prediction: hide some tokens, predict them from both sides. That is what BERT did in 2018. Then a contrastive stage pulls matching pairs together and pushes in-batch strangers apart, and one token's final hidden state, or the mean of all of them, becomes the vector.

The second phase gets most of the attention. The UltraBERT paper argues the first phase sets the ceiling: "the pre-training recipe of the backbone determines the upper bound of the resulting model." Its evidence is the core of this release, so I'll start there and come back to the block.

The finding: classification and retrieval want different masking

Masked language modelling has one number BERT picked and everyone inherited: mask 15% of tokens. Later work found higher ratios better, especially for bigger models. The paper's point is that those studies ran on 10B to 100B tokens, where a higher ratio mostly converges faster: more masked tokens per sequence means more loss signal per step. Speed and final quality were tangled.

So they trained to saturation. Encoders at 150M and 1.3B on up to 8.8T tokens plateau on GLUE, SQuAD and BEIR at around 1 to 2 trillion tokens, even as the training loss keeps falling (reported). With speed factored out, the masking ratio splits the tasks cleanly (reported, Table 1, 150M encoder):

Light masking makes a better classifier. Heavy masking makes a better retriever. Their explanation for the retrieval side is effective rank: the mean-pooled vectors of a 50%-masked model spread over more dimensions (reported, Figure 6). They are careful about it: "We cannot tell whether the higher effective rank causes the improvement in retrieval." Forcing the spread with a KoLeo regulariser at 6% masking helped BEIR less than plain 50% masking and cost GLUE, so the rank is not the whole story.

Mixing ratios mostly averages them. A uniform draw from 5% to 50% lands where a flat 30% does (88.0 and 44.2 against 88.3 and 44.2). The one mixture that gets both ends is the last row: put 5% and 50% masked sequences in the same batch, and give each its own LM head, the projection from the final hidden state back to the 49,152-token vocabulary. The paper's diagnosis is that the two regimes pull the output projection into different shapes. The input embeddings of a 6% and a 50% model look almost identical, and the LM heads do not (reported, Figure 7). One shared head is a tug of war. Two heads let the shared encoder serve both.

The final recipe adds one more piece. A 1.3B UltraBERT is the teacher. It is pruned along width to 380M, distilled, and pruned again to 150M (reported, Appendix A.2). Distillation helped GLUE at 15% masking (90.2 to 90.8) and hurt BEIR at 50% (47.2 to 46.5) in their 380M test, so the KL loss is applied only to the 5% sequences, and the 50% sequences keep plain MLM.

expand

Two pieces of arithmetic on that figure. The batch is four 5% sequences per 50% sequence, so the average masked fraction is (4 × 5 + 50) / 5 = 14% (reasoned), almost BERT's 15%. The average is not the point: the uniform-mixture row shows that a model fed the average behaves like the average. And neither LM head exists in the embedder. Both are discarded once pre-training ends, which is exactly why decoupling them is cheap.

The base encoder's own results: 89.5 GLUE and 45.2 BEIR at 152M, against ModernBERT-base's 88.4 and 42.6 and Ettin's 88.9 and 42.4 at 149M (reported, Table 3). That BEIR number is the paper's own probe, a mean-pooled fine-tune on MS MARCO alone. It measures the backbone, not UltraBERT-Embed.

expand

The block: ModernBERT's skeleton, new proportions

Now the code. modeling_ultrabert.py is a stock Hugging Face transformer file generated from a modular definition. It contains model code only: no network calls, no file writes, no eval, no subprocess. That is the bar I set before letting trust_remote_code run it.

Here is one layer, from UltrabertLayer.forward:

# modeling_ultrabert.py, UltrabertLayer.forward (trimmed)
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)          # RMSNorm
hidden_states, _ = self.self_attn(hidden_states, ...)          # RoPE inside
hidden_states = residual + hidden_states
residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)  # RMSNorm
hidden_states = self.mlp(hidden_states)                       # SwiGLU
hidden_states = residual + hidden_states

If that looks like a Llama layer, it is one, with is_causal = False. Every choice in it was already in ModernBERT: rotary positions instead of a learned table, normalisation before each sublayer instead of after, a gated MLP, no biases, and attention that alternates two local layers with one global one. UltraBERT's changes are in the numbers, all from config.json (measured):

• Depth and heads. 24 layers of width 768, with 6 heads of 128 dimensions instead of 12 of 64. Fewer, wider heads come from the pruning: the 380M parent had 8 heads of 128, and the prune dropped two heads rather than shrinking each.
• A thin MLP. intermediate_size is 1,024. SwiGLU computes down(silu(gate(x)) * up(x)), so each MLP is three 768 × 1,024 matrices, 2,359,296 weights. An attention block is four 768 × 768 matrices, also 2,359,296. Over 24 layers, attention and MLP each hold exactly 56,623,104 parameters (measured from the header). BERT spends twice as much on the MLP as on attention. UltraBERT spends the same.
• A wider local window. The local mask is built by create_local_attention_mask, which keeps tril(128).triu(-128): 128 tokens to each side. Calling it at length 1,024, row 512 sees 257 positions and row 0 sees 129 (measured). ModernBERT's local_attention: 128 is the total window, 64 a side. UltraBERT's reach per local layer is twice that.
• A long global RoPE base. rope_theta is 10,000 on local layers and 2,560,000 on global ones, sixteen times ModernBERT-base's 160,000. The paper says it started at 640K and was raised to 2.56M when the context was extended to 32K for the last 100B of 2T training tokens (reported, A.1). A larger base slows the lowest rotary frequencies down, so positions 32K apart still get distinct angles. The site's RoPE explainer has the derivation.
• Global layers at 2, 5, 8 … 23. The pattern is offset by two from ModernBERT's 0, 3, 6, so the last layer is global. That matters for this model in particular, below.
• A smaller tokenizer. SmolLM2's 49,152-token BPE, the one with <jupyter_code> and <gh_stars> in its special tokens. Its code-heavy vocabulary is a plausible contributor to the code scores, though nothing in the release isolates it.

The cost meter in the widget is pattern arithmetic. At 32,768 tokens a global layer scores 32,768² query-key pairs and a local one 32,768 × 257. With 8 global and 16 local layers, the stack does (8 + 16 × 257 / 32,768) / 24 = 33.9% of the score pairs of an all-global stack (reasoned). The global layers are nearly all of it. That is why 32K is affordable, and also why it is not free.

Parameters: 151,032,576, and where the 152M went

The embedder's model.safetensors header lists 218 tensors, all float32, totalling 151,032,576 parameters (measured; the loaded model agrees to the parameter). The token table is 49,152 × 768 = 37,748,736 of them, 25.0% (measured). The 49 norm vectors (two per layer and a final one) add 37,632.

The paper calls the base model 152M "excluding the LM head." The base checkpoint, dmis-lab/UltraBERT-150M, holds 189,960,960 (measured from its header). Subtract the 49,152 × 768 cls output layer and you get 152,212,224, which is the paper's 152M. The remaining 1,179,648 are a two-layer UltrabertPredictionHead that sits before cls. The embedder drops both, which leaves the card's 151M (reasoned, exact to the parameter). Three counts, one model, all consistent.

The checkpoint ships in float32, so the file is about 604 MB (reasoned: 151,032,576 × 4 bytes). The card's own example loads it in bfloat16, which halves that in memory.

How the embedder is trained, as the card states it

None of the embedder's training code is published, so this section is reported:

1. Weakly supervised contrastive stage. About 500M text and code pairs: question-answer, title-body, StackOverflow, Wikipedia, S2ORC, plus CornStack and other code pairs. In-batch negatives only.
2. Supervised stage. About 5.7M examples with mined hard negatives: MS MARCO, NQ, HotpotQA, FEVER, MLDR, BioASQ, NLI, synthetic classification data from NVIDIA's Nemotron embedding set, toxicity sets. A task instruction on the query side, and inputs up to 32K tokens so long documents are seen whole.
3. Checkpoint averaging. The released weights are an average of several stage-two checkpoints.

The "without distillation" claim is about this stage: no larger embedding model's scores were used as targets, unlike Tarka-Embedding-150M-V1 and F2LLM-v2-160M in the same table. The base encoder itself was distilled from the 1.3B UltraBERT during pre-training. Both statements are true. They are about different stages.

What the inference path does

The Sentence Transformers config is two modules: the transformer, then CLS pooling with include_prompt: true (measured). There is no dense projection and no normalisation module, so the raw vector is not unit length: a one-character input came out with norm 89.19 (measured). The card sets cosine as the similarity, which divides that out. If you put these vectors in an index that uses inner product, normalise them first.

"CLS" here is not a dedicated classification token. The tokenizer's template prepends <|endoftext|>, id 0, to every input, and that is the token pooled (measured from tokenizer.json). The same id is also the pad token. With right padding that is harmless: the pooled position is always 0.

This is where the global-layer offset earns its place. In a local layer, position 0 sees only the first 129 tokens. Only the 8 global layers let it read a whole 32K document, and because layer 23 is global, the pooled token's last update reads everything (reasoned from the config).

On instructions, config_sentence_transformers.json defines two prompts (measured):

query:    "Instruct: Given a question, retrieve passages that answer the question\nQuery: "
document: ""

Documents never take an instruction. For symmetric tasks the card says to put the same instruction on both sides, and labels for zero-shot classification take Label: . There are no Matryoshka dimensions: the card never claims them and the pooling config emits 768 only, so truncating these vectors is untested.

Running it on CPU

I ran it in a throwaway venv: torch 2.14.1 CPU, Sentence Transformers 6, SDPA attention, float32, two threads on a shared four-core Xeon at 2.10 GHz, under nice. These timings carry the noise of a busy machine, so read them as orders of magnitude (all measured):

That is about 5.3 short passages a second (reasoned). Quadrupling the document from 2,048 to 8,192 tokens cost 8.2× the time (reasoned), the global layers' quadratic term showing through.

Two checks on correctness. The same passage embedded alone and in a batch next to a much longer one gave cosine 0.9999999 (measured), so padding masks are applied properly. And the sanity retrieval below: six questions, six passages, each passage written to answer one of them, with deliberate overlaps (the RoPE question and the long-attention question share vocabulary).

Both ways it gets 6 of 6 (measured). The instruction's effect is on the floor. The mean off-diagonal cosine falls from 0.135 bare to 0.048 instructed, and the diagonal falls too, from 0.586 to 0.523 (measured). The smallest gap between a right passage and the best wrong one actually narrows slightly, from 0.288 to 0.264. Six pairs say the model works and the instruction does what the card says. They say nothing about rankings.

The 32K context, and the path that can't reach it

The card recommends FlashAttention 2, and the code shows why. On that path the mask stays a padding mask and the window goes to the kernel as sliding_window=local_window + 1, so memory stays linear in the local layers. On every other path, SDPA and eager, UltrabertModel.forward builds a dense [batch, 1, L, L] float mask for each distinct window, here two (one ±128, one global). At L = 32,768 one such mask is 32,768² × 4 bytes = 4.29 GB (reasoned), per sequence, before attention has computed anything. The 32K claim is real on a GPU with FlashAttention. On a CPU or a Mac it is not a number you get to use. My 8,192-token run was already the slowest thing I measured.

The paper's long-context evidence is good. LongEmbed's passkey test hides a number in a document and asks for it back, at lengths from 256 to 32K.

expand

Checking the benchmark claims

The card's averages recompute exactly from its own per-category columns, weighting each by its task count, 8 + 8 + 3 + 2 + 10 + 9 + 1 = 41 tasks (reasoned):

The lead is 0.85 points over Tarka and 0.98 over gte-modernbert-base (reasoned). That is the mean over tasks. Averaged over the seven task types instead, which gives the single summarisation task the same weight as ten retrieval tasks, UltraBERT-Embed scores 62.77 (reasoned). Both are legitimate summaries. They should not be compared across tables.

The more interesting reading is how it wins. Of seven task types, UltraBERT-Embed is best in one, summarisation, which is a single task (reported). It is second in retrieval (57.20 against DenseOn's 59.09) and in STS (83.03 against GIST's 83.35). In classification, 81.74 trails the two distilled models at 87.13 and 86.28, and clustering trails them too (reported). It takes the average by never being far from the top anywhere. For a general-purpose index that is the property you want. For a classifier, the distilled models are better.

The code table has the same shape. It leads 5 of 12 tasks, but CosQA (36.29 against 43.36), text-to-SQL (58.05 against 68.27) and CodeEdit (62.21 against 74.81) go to other models (reported). On LongEmbed it leads five of six, and on Needle a 106M bekko model scores 76.50 to its 66.50 (reported). The RAR-b row is mixed too: SpartQA spatial reasoning is 7.05 against 64.63 for a 312M granite model.

What I could not check: none of the baseline numbers were re-run, the "state of the art under 200M" line depends on a leaderboard snapshot dated 11 October, and the training-set sizes are the card's. The paper says the evaluation code is released. I found no link to it on the card, the paper page or the model repo.

Where it fits

Against the other embedders on this site, this is the small, single-vector, English end. EmbeddingGemma 2 spends 740M on five modalities; pplx-embed-v2-late keeps one vector per token. If your documents are long English text or code, and you have a GPU with FlashAttention, UltraBERT-Embed embeds them whole at 151M parameters. If you run on CPU, budget for short chunks anyway, and keep BM25 next to it: six toy pairs prove nothing about exact-term recall.

The lasting part is the backbone paper, not the embedder. Masking ratio has been treated as a setting to tune once. UltraBERT's result is that two settings, each with its own output head, beat any one setting or any schedule between them, and that the heads can be thrown away after. That transfers to any encoder pre-training run.