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

> Satyajit Ghana — Head of Engineering @ Inkers Technology
> canonical: https://ai.thesatyajit.com/articles/ultrabert-embed
> date: 2026-10-14
> tags: retrieval, encoders, long-context, small-models, benchmarks

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](https://arxiv.org/abs/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.

<ModelCard repo="dmis-lab/UltraBERT-Embed-150M" />

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](/architectures/encoder-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):

| Masking | GLUE | BEIR nDCG@10 |
|---|--:|--:|
| 6% | **89.5** | 43.5 |
| 15% | 88.9 | 43.8 |
| 30% | 88.3 | 44.2 |
| 50% | 87.5 | **45.0** |
| decay 50% → 6% | 88.8 | 43.9 |
| uniform 5% to 50% | 88.0 | 44.2 |
| 5% and 50% in one batch, shared head | 88.9 | 44.7 |
| 5% and 50% in one batch, **separate heads** | 89.4 | 45.1 |

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.

<Figure
  src="https://media.thesatyajit.com/ai/articles/ultrabert-embed/fig1.png?v=e316f5a994824a8c"
  alt="Diagram: batches of 5% masked sequences (four times as many) and 50% masked sequences feed UltraBERT, which was pruned from a 1.3B teacher. Two LM heads come off the encoder: the 5% head is trained against the teacher with a KL distillation loss, and the 50% head with the standard MLM loss."
  caption="The UltraBERT pre-training recipe: 5% and 50% masked sequences at 4:1 in each batch, one LM head per ratio, distillation from the 1.3B teacher on the 5% head only (UltraBERT paper, Figure 10)."
/>

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.

<Figure
  src="https://media.thesatyajit.com/ai/articles/ultrabert-embed/fig2.png?v=bb7e18efe501d5e3"
  alt="Two scatter plots of classification average against retrieval average, one for base encoders up to 200M and one for medium encoders up to 500M. In both, UltraBERT sits alone in the top-right corner, above and to the right of Ettin, ModernBERT and, in the medium group, LFM2.5."
  caption="Classification average (GLUE, SQuAD 2.0, two NER sets, RewardBench) against retrieval average (BEIR, RAR-b code and math, MLDR) for each backbone after the paper's own fine-tuning; UltraBERT is the starred point in both size groups (UltraBERT paper, Figure 1)."
/>

## 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`:

```python
# 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](/architectures/rope) 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.

<EncoderBlockCompare />

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

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

| Workload | Wall time |
|---|--:|
| 64 passages × 128 tokens, batch 32 | 12.07 s (best of 3) |
| one 2,048-token document | 4.72 s |
| one 8,192-token document | 38.69 s |

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

<SimilarityMatrix />

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.

<Figure
  src="https://media.thesatyajit.com/ai/articles/ultrabert-embed/fig3.png?v=a3338f9d5aa04ac2"
  alt="Heatmap of passkey retrieval accuracy against document length from 0.25K to 32K tokens for base and medium encoders. Most baselines turn red beyond their context length; ModernBERT and Ettin drop at 16K and 32K. UltraBERT's row is green at every length in both size groups."
  caption="Passkey accuracy by document length for each backbone after the paper's MS MARCO and MLDR fine-tune; UltraBERT holds 99% at every length to 32K at both sizes (UltraBERT paper, Figure 11)."
/>

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

| Average | Card | Recomputed |
|---|--:|--:|
| UltraBERT-Embed, MTEB(eng, v2) | 67.24 | 67.236 |
| Tarka-Embedding-150M-V1 | 66.39 | 66.391 |
| gte-modernbert-base | 66.26 | 66.259 |
| UltraBERT-Embed, MTEB(Code, v1), 12 tasks | 73.30 | 73.296 |
| UltraBERT-Embed, LongEmbed, 6 tasks | 79.30 | 79.302 |

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](/articles/embeddinggemma-2) spends 740M on five modalities;
[pplx-embed-v2-late](/articles/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](/articles/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.