What it actually does
A causal model sees only the left context because it has to: the chain rule imposes an order. If you are not going to generate, that constraint buys nothing — a classifier reading a sentence has the whole sentence. So hide fifteen percent of the tokens, let every position attend in both directions, and predict what was hidden. The representations are strictly better for understanding tasks because each token is informed by what follows it as well as what precedes it. What you give up is the ability to generate and the ability to compute a likelihood, and both losses follow from the same fact.
- The whole input is available at inference, which is what makes bidirectional context legitimate
- Generation is not required — the objective provides no mechanism for it and no ordering to generate along
- A likelihood is not required, since the conditionals do not compose into a joint density
- The input fits inside the pretrained sequence length, which is a hard maximum
- The pretraining domain is close enough to the target that the representations transfer, which the tokenizer largely determines
Loss function
Classified as a loss rather than a likelihood, and that classification is the central technical content of the entry. Each term looks like a log-probability, so the sum looks like a log-likelihood, and it is not one: the model produces the conditional distribution of each masked token given the rest, and a set of conditionals given different conditioning sets does not multiply into a joint. There is no factorization theorem here of the kind the chain rule provides for a causal model. Three things follow, and all three are usually stated as separate facts when they are one fact. The model cannot generate, because there is no ordering to generate along. It cannot score a sequence, because there is no density to evaluate — the pseudo-perplexity people compute is a heuristic and not a likelihood. And the loss value is not comparable to a causal model’s perplexity, so the two numbers should never be put on the same axis.
- the masked positions — chosen by the training loop, typically fifteen percent of tokens
- everything not masked, on BOTH sides; there is no causal restriction anywhere
- a conditional over the vocabulary — but these conditionals do NOT compose into a joint
- the count of masked positions; the loss is averaged over them and ignores the rest
How it is minimized
AdamW with a long warmup, then discarded: pretraining is done once and the practical work is fine-tuning
- the masking rate; a tuned constant balancing signal per example against how much context survives
- replaced by a sentinel token that appears at pretraining and never at fine-tuning
- forces the model to check every position rather than trusting unmasked ones
- the correction for the sentinel mismatch — the model must still predict a token that is already there
The corruption recipe is the part worth understanding, because every clause in it is a correction for a problem the previous clause created. Masking fifteen percent is a balance: more gives more training signal per example but leaves less context to predict from, and less wastes forward passes on positions that contribute nothing to the loss. Then the eighty-ten-ten split addresses a mismatch the sentinel token introduces — that token appears during pretraining and never during fine-tuning, so a model that has only ever predicted at sentinel positions has learned to condition on something it will never see again. Replacing some masked positions with random tokens forces the model to evaluate every position rather than trusting unmasked ones, and leaving some unchanged forces it to predict tokens that are already correct. That is three mechanisms to patch one design decision, and it is a good illustration of how much of a method can be corrections rather than the core idea. The efficiency objection is the other thing to know: only fifteen percent of positions contribute to the loss, so the objective uses each forward pass far less than a causal model’s does, and replacement-detection objectives exist precisely to recover that.
masking rate
Signal per example against surviving context. Fifteen percent is the convention; higher works with longer sequences (15% to 40%)
corruption split
The eighty-ten-ten recipe, correcting for the sentinel token that never appears at fine-tuning (80/10/10)
masking granularity
Token, whole-word or span. Span masking is consistently better because predicting one subword of a word is nearly free (token / word / span)
sequence length
Attention is quadratic in it, and bidirectional attention has no cache to amortize anything (128 to 512)
fine-tuning learning rate
One to two orders of magnitude below pretraining; the pretraining rate destroys the representations (1e-5 to 5e-5)
pooling strategy
Sentinel token, mean or max over positions. Not interchangeable between checkpoints even when all are available (CLS / mean / max)
Pretraining is stable and unremarkable, and almost nobody does it — the practical question is which checkpoint to fine-tune, and fine-tuning is where the failures are. The characteristic one is catastrophic forgetting from too high a learning rate: the pretraining rate destroys the representations within a few hundred steps, and the symptom is a model that trains to a plausible-looking loss while performing worse than a linear probe on frozen features. The second is instability on small datasets, where different seeds give materially different results and the published number is frequently the best of several runs. The third is subtler and worth naming: the masking rate and the corruption recipe are tuned constants, and the fifteen percent figure in particular has been shown to be far from optimal at longer sequence lengths — treating it as a law rather than a default leaves real performance unclaimed.
Pretraining is O(n²·d) for attention plus O(n·d²) for projections, fully parallel across positions. Inference is one forward pass per sequence with no sequential structure at all, which is the operational advantage over a causal model — and there is no key-value cache because there is nothing to cache, since every position depends on every other and nothing is reused across calls.
How, where, and why it is used
Time-Series Forecasting
AdaptedMasked reconstruction as a pretraining objective rather than as the final task: patch the series, mask a fraction of the patches, and train an encoder to reconstruct them, then attach a supervised forecasting head. The pretraining uses unlabelled history, which in forecasting is abundant, and the head is what actually forecasts.
- Pretraining a forecasting encoder on a large unlabelled panel before supervised fine-tuning
- Representation learning for downstream classification or clustering of series
- Imputation, where masked reconstruction is literally the task rather than a proxy for it
- Transfer across related panels, where the pretrained encoder carries structure the target lacks data for
The pretraining argument transfers cleanly: unlabelled series are abundant and masked reconstruction extracts structure from them without any labels, which is exactly the situation that made this objective valuable in language. Imputation is the one case where the objective and the task coincide, and it is the strongest fit here. But the model does not forecast — a bidirectional encoder has no mechanism for producing values beyond its input, so a supervised head is doing the forecasting and the pretraining is only initializing it. That distinction matters because it bounds what the pretraining can contribute, and because published gains from masked pretraining on forecasting benchmarks are frequently small once a properly tuned supervised baseline is included.
- Patch the series before masking; masking individual timesteps is nearly free to predict from neighbours and teaches almost nothing
- Mask contiguous spans rather than scattered points, for the same reason span masking beats token masking in language
- Normalize per window, and keep the scaler with the model as for any series encoder
- Fine-tune the head with a much lower learning rate than pretraining used, or the representations are destroyed
Rolling-origin backtesting of the fine-tuned model against a supervised model of the same architecture trained from scratch — that comparison isolates what the pretraining contributed and it is frequently less than claimed. For imputation, reconstruction error on held-out masked spans directly.
- Expecting the encoder to forecast; a supervised head does that and the pretraining only initializes it
- Masking individual timesteps, which are trivially interpolable and teach nothing
- Fine-tuning at the pretraining learning rate and destroying the representations
- Omitting the from-scratch baseline, so the pretraining benefit is never actually measured
Anomaly Detection
Not applicableThere is no likelihood to threshold. The model produces the conditional distribution of a masked token given the rest, and a set of conditionals over different conditioning sets does not compose into a joint density — so there is no probability of a sequence to compute and nothing to call surprising. The pseudo-perplexity that is sometimes reported requires one forward pass per position and is a heuristic rather than a likelihood, and it is not comparable to a causal model’s perplexity; for likelihood-based detection the decoder-only entry is the correct tool.
Optimization
AdaptedTwo separate optimization stories. The corruption recipe is a tuned design where each clause corrects a problem created by the previous one, which makes it a good study in how much of a method can be patches rather than core idea. And the pretrain-then-fine-tune structure is a compute-allocation decision: pretraining is expensive and amortized across every downstream task, fine-tuning is cheap and per-task, and the split between them is the actual choice.
- Masking rate as an explicit trade between signal per example and surviving context
- The corruption recipe as a chain of corrections, each addressing a flaw introduced by the last
- Pretrain-then-fine-tune as amortization: one expensive run serving many cheap adaptations
- Parameter-efficient fine-tuning, capturing most of the benefit at a fraction of the parameters
Worth studying for the sample-efficiency argument, which is sharper here than almost anywhere else in this reference. Only the masked positions contribute to the loss, so at a fifteen percent rate the objective extracts signal from fifteen percent of a forward pass it paid for in full — a causal model gets a prediction at every position. That inefficiency is measurable, it motivated replacement-detection objectives that train on every position instead, and it is the clearest case here of an objective being improved by changing what it asks for rather than how it is optimized. The other transferable point is that the fifteen percent figure has been shown to be far from optimal at longer sequences, which is a reminder that a widely copied constant is not a law.
- Treat the masking rate as a tunable rather than as a constant; fifteen percent is a default from one paper at one sequence length
- Mask spans rather than tokens, since predicting one subword of a known word is nearly free
- Fine-tune at one to two orders of magnitude below the pretraining learning rate
- Prefer parameter-efficient adaptation where the task set is large, since full fine-tuning stores a whole model per task
Sweep the masking rate on the target sequence length rather than inheriting it, and measure downstream task performance rather than pretraining loss — the two diverge, and pretraining loss is not the objective. Compare full against parameter-efficient fine-tuning at matched task performance to see what the extra parameters bought.
- Treating fifteen percent as optimal when it is a default that does not transfer across sequence lengths
- Optimizing pretraining loss rather than downstream performance, which is what actually matters
- Full fine-tuning for every task, which stores a complete model per task for a marginal gain
- Comparing masked-model loss against causal-model perplexity, which are different quantities
Natural Language
Primary fitPretrain a bidirectional encoder on a large corpus, then attach a task head and fine-tune. For classification the pooled representation feeds a linear layer; for token-level tasks each position feeds its own classifier; for retrieval the representations are used directly as embeddings.
- Text classification at high volume, where a small fine-tuned encoder beats a large generative model on both accuracy and cost
- Token-level tasks — named-entity recognition, part-of-speech tagging, extraction — where bidirectional context is decisive
- Sentence and passage embeddings for retrieval and similarity
- Extractive question answering, where the answer is a span in the input rather than generated text
For understanding tasks this remains the right choice and the reasoning is structural rather than historical. Bidirectional context genuinely helps when the whole input is available, which for classification it always is. Inference is a single forward pass with no sequential generation, so throughput is orders of magnitude above a generative model. And the models are small enough to fine-tune and serve cheaply. Generative models have absorbed much of this territory through prompting, and the honest comparison is that they are more convenient and considerably more expensive — on a single high-volume classification task a fine-tuned encoder usually wins on accuracy as well as cost, and that is worth checking rather than assuming.
- Match the tokenizer to the checkpoint exactly; a different one produces different subwords and degrades everything silently
- Use the pooling the checkpoint was trained for — sentinel-token and mean pooling are not interchangeable
- Fine-tune at a low learning rate with a short warmup, since the pretraining rate destroys the representations
- Run several seeds on small datasets, where variance between runs is large enough to dominate any comparison
Task F1 or accuracy on a held-out split, averaged over several seeds rather than reported from the best run — fine-tuning variance on small datasets is large and single-run numbers are not reproducible. Compare against a prompted generative model at matched latency and cost, which is the comparison that reflects the real alternative.
- A tokenizer mismatch with the checkpoint, which degrades everything with no error
- Fine-tuning at the pretraining learning rate and forgetting the pretraining
- Reporting the best of several seeds, which is not reproducible
- Using a large generative model for a task a small encoder handles better and far more cheaply
Computer Vision
AdaptedMasked autoencoding: divide the image into patches, hide most of them — far more than in language — and reconstruct the missing pixels from the visible ones. Because images are spatially redundant, a high masking ratio is necessary to make the task hard enough to be informative, and it also makes the encoder cheap since it only processes the visible patches.
- Self-supervised pretraining of vision encoders without labels
- Pretraining for dense prediction, where the reconstruction objective aligns with spatial detail
- Domains with abundant unlabelled imagery and expensive annotation, such as medical or satellite data
- Initializing an encoder for a downstream task with too few labels to train from scratch
The masking-ratio contrast with language is the instructive part: language needs fifteen percent and images need seventy-five or more, because a missing patch is largely inferable from its neighbours while a missing word frequently is not. That difference is a statement about redundancy in the two modalities rather than a tuning detail. The efficiency consequence is large and favourable — the encoder sees only the visible quarter of patches, so pretraining is several times cheaper than a comparable contrastive method. Against it: the learned features are reconstruction-oriented, so they transfer better to dense tasks than to linear classification, where contrastive pretraining remains stronger. Which objective to prefer depends on the downstream task rather than on general quality.
- Mask aggressively — seventy-five percent or more, since images are redundant and a low ratio makes the task trivial
- Feed only the visible patches to the encoder, which is where most of the cost saving comes from
- Normalize the reconstruction target per patch, which measurably improves the learned features
- Choose the objective for the downstream task: masked reconstruction for dense prediction, contrastive for linear classification
Fine-tuning and linear-probe performance reported separately, because they rank pretraining objectives differently and masked reconstruction looks much worse under linear probing than it performs after fine-tuning. Compare against a contrastive baseline at matched pretraining compute.
- A low masking ratio making the task trivially solvable by interpolation
- Judging the method by linear probing, which understates it relative to fine-tuning
- Feeding masked patches to the encoder, which discards the main efficiency gain
- Expecting language-style masking ratios to transfer, when the redundancy is entirely different
What it costs to run
Pretraining is expensive and done once by someone else — the practical decision is which checkpoint to start from. Fine-tuning is minutes to hours on a single accelerator, and parameter-efficient methods reduce it further, which makes this among the cheapest routes in this reference to a strong task-specific model.
One forward pass per sequence, fully parallel across positions, with no sequential generation and no cache. That combination gives throughput orders of magnitude above a generative model on the same hardware, and it is the strongest remaining operational argument for the architecture.
The encoder is replaced rather than retrained. Task heads are retrained whenever the label distribution moves, which is cheap and should be frequent.
- Track the rate of out-of-vocabulary or heavily-subworded inputs, which indicates the tokenizer no longer matches the domain
- Watch per-class performance rather than the aggregate, since drift appears in a few classes long before the average moves
- Monitor input length against the pretrained maximum, because truncation is silent and removes the end of the text
- Compare against a frozen-feature linear probe periodically — if fine-tuning is not beating it, the fine-tuning has damaged the representations
- There is no likelihood. Do not build a scoring or anomaly-detection system on a masked model — the conditionals do not compose into a joint, and the pseudo-perplexity sometimes reported is a heuristic that also costs one forward pass per position
- The tokenizer is part of the model. A different one produces different subwords and degrades everything with no error anywhere
- Pooling strategy is part of the model. Sentinel-token and mean pooling are not interchangeable, even when a checkpoint exposes both
- Fine-tuning at the pretraining learning rate destroys the representations within a few hundred steps, and the symptom is a plausible loss curve with poor performance
- The sequence length is a hard maximum and truncation removes the end of the text, which for classification is frequently where the conclusion is
- Fine-tuning variance on small datasets is large. A single-seed result is not reproducible and should not be reported as one
Pros and cons, in context
Bidirectional context, which is free when you are not generating
Every token is informed by what follows as well as what precedes, and for classification the whole input is always available — so the causal restriction buys nothing and dropping it costs nothing that matters.
One parallel forward pass, no sequential generation
Throughput orders of magnitude above a generative model on the same hardware, with no cache and no per-token loop. The strongest remaining operational argument for the architecture.
Small enough to fine-tune and serve cheaply
Minutes to hours on one accelerator for a strong task-specific model. On a single high-volume classification task this beats a prompted large model on accuracy as well as cost, which is worth checking rather than assuming.
Masked reconstruction transfers across modalities
The same objective pretrains vision encoders, with a masking ratio five times higher because images are far more redundant. That contrast is itself informative about the modalities.
No joint likelihood, and therefore no scoring
Conditionals over different conditioning sets do not multiply into a density. This is the central limitation and it is the same fact that prevents generation — the pseudo-perplexity sometimes reported is a heuristic, not a likelihood.
Cannot generate
There is no ordering to generate along. For any task requiring produced text the decoder-only architecture is the only option, which is why generative models absorbed so much of this territory.
Sample-inefficient by construction
Only the masked positions contribute to the loss, so fifteen percent of a fully-paid forward pass produces signal. Measurable, and the direct motivation for replacement-detection objectives that train on every position.
The sentinel token creates a pretrain-to-fine-tune mismatch
It appears during pretraining and never afterwards, which the eighty-ten-ten recipe patches rather than solves. Three mechanisms correcting one design decision is a lot of machinery for a detail.
Fine-tuning is unstable on small datasets
Different seeds give materially different results, so single-run numbers are not reproducible and published ones are frequently the best of several. This affects how much any small-data comparison should be believed.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""Masked language model, transcribed from the objective.
No ML library. Bidirectional self-attention with NO causal mask -- that
single omission relative to a decoder-only model is the whole difference,
and everything else follows from it: no generation, no likelihood, and
better representations for anything that reads a complete input.
The loss averages log p(x_i | x_not-masked) over the masked positions
only. It looks like a log-likelihood and is not one: each term conditions
on a different set, and conditionals over different conditioning sets do
not multiply into a joint.
"""
import math
import random
def softmax_rows(scores):
"""Row-wise softmax with the max subtracted for stability."""
out = []
for row in scores:
top = max(row)
exps = [math.exp(value - top) for value in row]
total = sum(exps)
out.append([value / total for value in exps])
return out
def layer_norm(vector, gamma, beta, eps=1e-5):
width = len(vector)
mean = sum(vector) / width
var = sum((value - mean) ** 2 for value in vector) / width
scale = 1.0 / math.sqrt(var + eps)
return [gamma[i] * (vector[i] - mean) * scale + beta[i] for i in range(width)]
def matvec(matrix, vector):
return [sum(row[i] * vector[i] for i in range(len(vector))) for row in matrix]
def bidirectional_attention(hidden, w_q, w_k, w_v, w_o):
"""Self-attention over the full sequence, both directions.
Compare against a causal model: identical arithmetic except that the
decoder zeroes every score at j > i. Here nothing is zeroed. Position
3 sees position 7. That is the entire architectural difference.
"""
n_positions = len(hidden)
width = len(hidden[0])
queries = [matvec(w_q, hidden[pos]) for pos in range(n_positions)]
keys = [matvec(w_k, hidden[pos]) for pos in range(n_positions)]
values = [matvec(w_v, hidden[pos]) for pos in range(n_positions)]
inv_sqrt_d = 1.0 / math.sqrt(width)
scores = []
for query_pos in range(n_positions):
row = []
for key_pos in range(n_positions):
# No mask. A causal model would append -inf when key_pos > query_pos.
dot = sum(queries[query_pos][k] * keys[key_pos][k] for k in range(width))
row.append(dot * inv_sqrt_d)
scores.append(row)
weights = softmax_rows(scores)
context = []
for query_pos in range(n_positions):
mixed = [0.0] * width
for key_pos in range(n_positions):
weight = weights[query_pos][key_pos]
for k in range(width):
mixed[k] += weight * values[key_pos][k]
context.append(matvec(w_o, mixed))
return context
def gelu(value):
"""The activation the original encoder used; tanh approximation."""
inner = math.sqrt(2.0 / math.pi) * (value + 0.044715 * value ** 3)
return 0.5 * value * (1.0 + math.tanh(inner))
def feed_forward(vector, w_in, w_out):
hidden = [gelu(value) for value in matvec(w_in, vector)]
return matvec(w_out, hidden)
def encode(token_ids, params):
"""Run the encoder stack. Every position attends to every position."""
hidden = [
[
params["token_embed"][token_id][k] + params["pos_embed"][pos][k]
for k in range(params["width"])
]
for pos, token_id in enumerate(token_ids)
]
for layer in params["layers"]:
attended = bidirectional_attention(
hidden, layer["w_q"], layer["w_k"], layer["w_v"], layer["w_o"]
)
hidden = [
layer_norm(
[hidden[pos][k] + attended[pos][k] for k in range(params["width"])],
layer["ln1_gamma"],
layer["ln1_beta"],
)
for pos in range(len(hidden))
]
transformed = [
feed_forward(hidden[pos], layer["w_in"], layer["w_out"])
for pos in range(len(hidden))
]
hidden = [
layer_norm(
[hidden[pos][k] + transformed[pos][k] for k in range(params["width"])],
layer["ln2_gamma"],
layer["ln2_beta"],
)
for pos in range(len(hidden))
]
return hidden
def corrupt(token_ids, vocab_size, mask_token_id, rate=0.15, rng=None):
"""The 80/10/10 recipe, written out clause by clause.
Each clause corrects a problem the previous one created:
- masking at all gives the model something to predict
- but the sentinel token appears at pretraining and never at
fine-tuning, so a model that only ever predicts at sentinel
positions has learned to condition on something it will not see
- so 10% of chosen positions get a RANDOM token, forcing the model
to evaluate every position rather than trusting unmasked ones
- and 10% are LEFT ALONE, forcing it to predict a token that is
already sitting there correctly
Three mechanisms patching one design decision. Worth noticing how
much of a published method can be corrections rather than core idea.
"""
rng = rng or random.Random(0)
corrupted = list(token_ids)
targets = {}
for pos, original in enumerate(token_ids):
if rng.random() >= rate:
continue
targets[pos] = original
draw = rng.random()
if draw < 0.8:
corrupted[pos] = mask_token_id
elif draw < 0.9:
corrupted[pos] = rng.randrange(vocab_size)
# else: leave it unchanged -- deliberately, see docstring.
return corrupted, targets
def masked_loss(token_ids, params, mask_token_id, rng=None):
"""Average negative log p(x_i | rest) over masked positions only.
Note the denominator: len(targets), not len(token_ids). Roughly 85%
of the positions we just paid a full forward pass for contribute
nothing. That inefficiency is real, it is measurable, and it is the
direct motivation for replacement-detection objectives that train on
every position instead.
"""
vocab_size = len(params["token_embed"])
corrupted, targets = corrupt(token_ids, vocab_size, mask_token_id, rng=rng)
if not targets:
return 0.0
hidden = encode(corrupted, params)
total = 0.0
for pos, target_id in targets.items():
logits = matvec(params["head"], hidden[pos])
top = max(logits)
log_z = top + math.log(sum(math.exp(value - top) for value in logits))
total += log_z - logits[target_id]
return total / len(targets)
def fill_mask(token_ids, params, mask_token_id, top_k=5):
"""Predict the tokens at sentinel positions.
This is the model's only native inference mode. It is NOT generation:
the positions are fixed, the length is fixed, and filling several
masks jointly would require an ordering the model does not have.
"""
hidden = encode(token_ids, params)
predictions = {}
for pos, token_id in enumerate(token_ids):
if token_id != mask_token_id:
continue
logits = matvec(params["head"], hidden[pos])
ranked = sorted(range(len(logits)), key=lambda idx: -logits[idx])
predictions[pos] = ranked[:top_k]
return predictions
Every position pays a full vocabulary projection and only the masked fifteen percent are read, so the head alone is O(n·d·V) for O(rate·n·d·V) of useful work. Illustrative, not a measured benchmark: at a 30k vocabulary that wasted projection is the single largest avoidable cost in the file.
How this language actually wants it written.
"""Masked language model with an explicit corruption policy.
The design point: make the 80/10/10 recipe a validated object rather than
three constants buried in a branch, and make the tokenizer contract
checkable, since a tokenizer mismatch is the failure mode that degrades
everything with no error anywhere.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Mapping, Sequence
import numpy as np
from numpy.typing import NDArray
EPSILON: Final[float] = 1e-5
class TokenizerMismatch(ValueError):
"""A token id falls outside the vocabulary the model was trained on.
Almost always means the tokenizer does not match the checkpoint. The
silent version of this bug -- ids that happen to be in range but mean
different subwords -- cannot be caught here, which is why the
tokenizer must be pinned alongside the weights.
"""
class SequenceTooLong(ValueError):
"""Input exceeds the pretrained maximum position count.
Raised rather than truncated. Truncation is silent and removes the END
of the text, which for classification is frequently where the
conclusion sits.
"""
class NoMaskedPositions(ValueError):
"""Corruption selected nothing, so the loss has no terms.
Happens on very short sequences at a low masking rate. Averaging over
an empty set would produce a NaN that propagates into the gradient.
"""
class LikelihoodUnavailable(NotImplementedError):
"""The caller asked for a sequence probability. There isn't one.
The model gives p(x_i | x_not-masked) for a masked position. Those
conditionals are over different conditioning sets and do not multiply
into a joint density. Pseudo-perplexity is a heuristic requiring one
forward pass per position, and it is not comparable to a causal
model's perplexity. Use a decoder-only model for scoring.
"""
@dataclass(frozen=True, slots=True)
class CorruptionPolicy:
"""How positions are chosen and what replaces them.
Frozen because it is part of the pretraining contract: changing the
rate mid-run changes the objective. The proportions are validated so
the recipe cannot drift into a state where a fraction is implicitly
negative.
"""
rate: float = 0.15
prob_sentinel: float = 0.8
prob_random: float = 0.1
span_length: int = 1
def __post_init__(self) -> None:
if not 0.0 < self.rate < 1.0:
raise ValueError(f"masking rate must be in (0, 1), got {self.rate}")
if self.prob_sentinel + self.prob_random > 1.0:
raise ValueError(
"sentinel and random proportions exceed 1.0: "
f"{self.prob_sentinel} + {self.prob_random}"
)
if self.span_length < 1:
raise ValueError(f"span_length must be >= 1, got {self.span_length}")
@property
def prob_unchanged(self) -> float:
"""The remainder: positions selected but left as they are."""
return 1.0 - self.prob_sentinel - self.prob_random
@dataclass(frozen=True, slots=True)
class TokenizerContract:
"""What the checkpoint expects. Validated, because mismatches are silent."""
vocab_size: int
sentinel_id: int
max_positions: int
def __post_init__(self) -> None:
if not 0 <= self.sentinel_id < self.vocab_size:
raise ValueError(
f"sentinel_id {self.sentinel_id} outside vocabulary of "
f"{self.vocab_size}"
)
def validate(self, token_ids: Sequence[int]) -> None:
"""Guard clauses, cheapest first."""
if len(token_ids) == 0:
raise ValueError("empty sequence")
if len(token_ids) > self.max_positions:
raise SequenceTooLong(
f"sequence of {len(token_ids)} exceeds pretrained maximum "
f"{self.max_positions}; truncate deliberately, not implicitly"
)
out_of_range = [tid for tid in token_ids if not 0 <= tid < self.vocab_size]
if out_of_range:
raise TokenizerMismatch(
f"{len(out_of_range)} token ids outside vocabulary "
f"(first: {out_of_range[0]}, vocab_size: {self.vocab_size})"
)
@dataclass(frozen=True, slots=True)
class MaskedBatch:
"""A corrupted sequence alongside what was hidden."""
corrupted: NDArray[np.int32]
positions: NDArray[np.int32]
targets: NDArray[np.int32]
def __len__(self) -> int:
return int(self.positions.size)
class MaskedLanguageModel:
"""A bidirectional encoder with a masked-token head.
Deliberately exposes no sampling or scoring method. Both would be
lies: there is no ordering to generate along and no joint density to
evaluate.
"""
def __init__(
self,
weights: Mapping[str, NDArray[np.float32]],
contract: TokenizerContract,
policy: CorruptionPolicy | None = None,
) -> None:
self._weights = weights
self._contract = contract
# No mutable default; a shared policy object across instances
# would let one caller's tuning leak into another's.
self._policy = policy if policy is not None else CorruptionPolicy()
@property
def policy(self) -> CorruptionPolicy:
return self._policy
def corrupt(
self,
token_ids: Sequence[int],
rng: np.random.Generator,
) -> MaskedBatch:
"""Apply the corruption policy, masking whole spans if configured.
Span masking is consistently better than token masking: predicting
one subword of a word whose other subwords are visible is nearly
free, so token-level masking spends much of its budget on
positions that teach nothing.
"""
self._contract.validate(token_ids)
original = np.asarray(token_ids, dtype=np.int32)
corrupted = original.copy()
n_positions = original.size
span = self._policy.span_length
n_spans = max(1, int(round(n_positions * self._policy.rate / span)))
starts = rng.choice(
max(1, n_positions - span + 1), size=n_spans, replace=False
)
selected: list[int] = []
for start in starts:
selected.extend(range(int(start), min(int(start) + span, n_positions)))
positions = np.unique(np.asarray(selected, dtype=np.int32))
if positions.size == 0:
raise NoMaskedPositions(
f"rate {self._policy.rate} selected nothing from "
f"{n_positions} positions"
)
draws = rng.random(positions.size)
use_sentinel = draws < self._policy.prob_sentinel
use_random = (draws >= self._policy.prob_sentinel) & (
draws < self._policy.prob_sentinel + self._policy.prob_random
)
corrupted[positions[use_sentinel]] = self._contract.sentinel_id
n_random = int(use_random.sum())
if n_random > 0:
corrupted[positions[use_random]] = rng.integers(
0, self._contract.vocab_size, size=n_random, dtype=np.int32
)
# positions where neither flag is set keep their original token.
return MaskedBatch(
corrupted=corrupted,
positions=positions,
targets=original[positions],
)
def encode(self, token_ids: Sequence[int]) -> NDArray[np.float32]:
"""Bidirectional encoding. One pass, no cache, nothing sequential."""
self._contract.validate(token_ids)
ids = np.asarray(token_ids, dtype=np.int32)
hidden = (
self._weights["token_embed"][ids]
+ self._weights["pos_embed"][: ids.size]
)
n_layers = int(self._weights["n_layers"][0])
for layer in range(n_layers):
hidden = self._encoder_block(hidden, layer)
return hidden
def loss(self, batch: MaskedBatch) -> float:
"""Cross-entropy over masked positions only.
Not a likelihood, and not comparable to a causal model's
perplexity. Naming it `loss` rather than `nll` is deliberate.
"""
if len(batch) == 0:
raise NoMaskedPositions("batch has no masked positions")
hidden = self.encode(batch.corrupted)[batch.positions]
logits = hidden @ self._weights["head"].T
shifted = logits - logits.max(axis=-1, keepdims=True)
log_z = np.log(np.exp(shifted).sum(axis=-1))
chosen = shifted[np.arange(batch.positions.size), batch.targets]
return float(np.mean(log_z - chosen))
def sequence_log_prob(self, token_ids: Sequence[int]) -> float:
"""Present only to fail loudly, because callers do ask for this."""
raise LikelihoodUnavailable(
f"a masked model has no joint density over the {len(token_ids)} "
"tokens; the conditionals do not compose"
)
def _encoder_block(
self, hidden: NDArray[np.float32], layer: int
) -> NDArray[np.float32]:
queries = hidden @ self._weights[f"w_q_{layer}"].T
keys = hidden @ self._weights[f"w_k_{layer}"].T
values = hidden @ self._weights[f"w_v_{layer}"].T
scores = queries @ keys.T / np.sqrt(hidden.shape[-1])
# No causal mask. The one line a decoder-only model adds here is
# the line that costs it the right half of every context.
weights = _softmax(scores)
attended = (weights @ values) @ self._weights[f"w_o_{layer}"].T
hidden = _layer_norm(hidden + attended)
inner = _gelu(hidden @ self._weights[f"w_in_{layer}"].T)
return _layer_norm(hidden + inner @ self._weights[f"w_out_{layer}"].T)
def _softmax(scores: NDArray[np.float32]) -> NDArray[np.float32]:
shifted = scores - scores.max(axis=-1, keepdims=True)
exps = np.exp(shifted)
return exps / exps.sum(axis=-1, keepdims=True)
def _layer_norm(hidden: NDArray[np.float32]) -> NDArray[np.float32]:
mean = hidden.mean(axis=-1, keepdims=True)
var = hidden.var(axis=-1, keepdims=True)
return (hidden - mean) / np.sqrt(var + EPSILON)
def _gelu(hidden: NDArray[np.float32]) -> NDArray[np.float32]:
inner = np.sqrt(2.0 / np.pi) * (hidden + 0.044715 * hidden ** 3)
return 0.5 * hidden * (1.0 + np.tanh(inner))
What changed: The corruption recipe becomes a frozen dataclass with validated proportions rather than three magic constants inside a branch, so the policy is inspectable and the eighty-ten-ten split cannot silently drift. The tokenizer contract — vocabulary size, sentinel id, maximum length — is checked at the constructor boundary, since a tokenizer mismatch is the failure that degrades everything with no error anywhere. Specific exceptions distinguish the failures that actually occur in production: truncation that would silently discard the end of the text, a sentinel id outside the vocabulary, and an attempt to use the model as though it had a likelihood.
Same asymptotics as the naive version — the vectorization is in the encoder, not the head, so the vocabulary projection still runs at every position. Illustrative, not a measured benchmark: what changed here is that a tokenizer mismatch or an over-length sequence now fails at the boundary instead of indexing out of bounds.
Faster, and the justification for each change.
"""Masked language model with the head confined to masked positions.
The optimization that matters is structural rather than numeric. The
vocabulary projection is the largest matmul in the model -- width times
vocab, typically 768 x 30000 -- and a naive implementation runs it at
every position. Only the masked positions enter the loss. Gathering first
and projecting second cuts that cost by roughly 1/rate, which at the
conventional 15% is close to a 6x reduction on the single most expensive
operation.
Worth being precise about what this does and does not fix. It removes the
wasted projection. It does NOT remove the underlying sample-inefficiency
of the objective: the encoder still ran over every position and only 15%
of them produced a gradient signal. That is a property of the objective,
not of the implementation, and the only fix is a different objective.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Mapping
import numpy as np
from numpy.typing import NDArray
EPSILON: Final[float] = 1e-5
NEG_INF: Final[float] = -1e30
@dataclass(frozen=True, slots=True)
class BatchedMask:
"""Flat masked-position indices into a padded (batch, positions) block.
Flat rather than per-sequence so the gather is one contiguous
take() rather than a loop. row_of keeps the mapping back for anything
that needs per-sequence attribution.
"""
corrupted: NDArray[np.int32]
attention_valid: NDArray[np.bool_]
flat_positions: NDArray[np.int64]
targets: NDArray[np.int32]
row_of: NDArray[np.int32]
def corrupt_batch(
token_ids: NDArray[np.int32],
lengths: NDArray[np.int32],
sentinel_id: int,
vocab_size: int,
rate: float,
rng: np.random.Generator,
) -> BatchedMask:
"""Corrupt a whole padded batch in one vectorized pass.
The subtlety is padding. A uniform draw over the padded rectangle
would select padding positions, which would then contribute garbage
targets to the loss. The valid mask is applied to the draw itself
rather than filtered afterwards, so padding is never selected.
"""
n_sequences, n_positions = token_ids.shape
position_index = np.arange(n_positions, dtype=np.int32)[None, :]
valid = position_index < lengths[:, None]
chosen = (rng.random(token_ids.shape) < rate) & valid
if not chosen.any():
# Force one position per sequence rather than emitting a NaN loss.
chosen[np.arange(n_sequences), np.maximum(lengths - 1, 0)] = True
corrupted = token_ids.copy()
draws = rng.random(token_ids.shape)
use_sentinel = chosen & (draws < 0.8)
use_random = chosen & (draws >= 0.8) & (draws < 0.9)
# chosen & (draws >= 0.9) is left alone, deliberately.
corrupted[use_sentinel] = sentinel_id
n_random = int(use_random.sum())
if n_random > 0:
corrupted[use_random] = rng.integers(
0, vocab_size, size=n_random, dtype=np.int32
)
rows, cols = np.nonzero(chosen)
flat_positions = (rows.astype(np.int64) * n_positions + cols).astype(np.int64)
return BatchedMask(
corrupted=np.ascontiguousarray(corrupted, dtype=np.int32),
attention_valid=valid,
flat_positions=flat_positions,
targets=token_ids[rows, cols].astype(np.int32),
row_of=rows.astype(np.int32),
)
def encode_batch(
corrupted: NDArray[np.int32],
attention_valid: NDArray[np.bool_],
weights: Mapping[str, NDArray[np.float32]],
n_heads: int,
) -> NDArray[np.float32]:
"""Batched bidirectional encoder.
Shapes, asserted rather than inferable from loops:
corrupted (batch, positions)
hidden (batch, positions, width)
heads (batch, n_heads, positions, head_dim)
scores (batch, n_heads, positions, positions)
"""
n_sequences, n_positions = corrupted.shape
width = weights["token_embed"].shape[1]
head_dim = width // n_heads
inv_sqrt_d = np.float32(1.0 / np.sqrt(head_dim))
hidden = (
weights["token_embed"][corrupted]
+ weights["pos_embed"][None, :n_positions, :]
).astype(np.float32, copy=False)
# Padding must be excluded from the KEY axis. Queries at padding
# positions produce garbage, which is harmless only because those
# positions are never gathered for the loss.
key_block = np.where(attention_valid[:, None, None, :], 0.0, NEG_INF).astype(
np.float32
)
n_layers = int(weights["n_layers"][0])
for layer in range(n_layers):
qkv = np.einsum(
"bpw,thw->tbhp",
hidden,
weights[f"qkv_{layer}"].reshape(3, width, width),
optimize=True,
)
queries, keys, values = (
part.reshape(n_sequences, n_heads, head_dim, n_positions).transpose(
0, 1, 3, 2
)
for part in qkv
)
scores = np.einsum("bhqd,bhkd->bhqk", queries, keys, optimize=True)
scores *= inv_sqrt_d
scores += key_block
# Still no causal mask -- only the padding block.
scores -= scores.max(axis=-1, keepdims=True)
np.exp(scores, out=scores)
scores /= scores.sum(axis=-1, keepdims=True)
attended = np.einsum("bhqk,bhkd->bqhd", scores, values, optimize=True)
attended = attended.reshape(n_sequences, n_positions, width)
hidden = _fused_residual_norm(
hidden, attended @ weights[f"w_o_{layer}"].T
)
inner = hidden @ weights[f"w_in_{layer}"].T
_gelu_inplace(inner)
hidden = _fused_residual_norm(hidden, inner @ weights[f"w_out_{layer}"].T)
return hidden
def masked_loss(
batch: BatchedMask,
weights: Mapping[str, NDArray[np.float32]],
n_heads: int,
) -> float:
"""Loss with the vocabulary projection confined to masked positions.
This is the whole optimization. `take` on the flattened
(batch * positions, width) view produces a compact contiguous block of
exactly the masked hidden states, and the head runs on that block
alone. At a 15% masking rate the largest matmul in the model shrinks
by roughly 6x.
"""
hidden = encode_batch(batch.corrupted, batch.attention_valid, weights, n_heads)
width = hidden.shape[-1]
gathered = np.ascontiguousarray(
hidden.reshape(-1, width).take(batch.flat_positions, axis=0)
)
logits = gathered @ weights["head"].T
# Fused log-sum-exp: subtract the max in place, then one log of one
# sum rather than materializing exp(logits) as a separate array.
logits -= logits.max(axis=-1, keepdims=True)
log_z = np.log(np.exp(logits).sum(axis=-1))
chosen = logits[np.arange(batch.targets.size), batch.targets]
return float(np.mean(log_z - chosen))
def _fused_residual_norm(
residual: NDArray[np.float32], delta: NDArray[np.float32]
) -> NDArray[np.float32]:
"""Add and normalize in place so neither intermediate is materialized."""
delta += residual
mean = delta.mean(axis=-1, keepdims=True)
delta -= mean
inv_std = 1.0 / np.sqrt(delta.var(axis=-1, keepdims=True) + EPSILON)
delta *= inv_std
return delta
def _gelu_inplace(hidden: NDArray[np.float32]) -> None:
"""GELU without allocating; the FFN intermediate is 4x width and hurts."""
scratch = hidden * hidden
scratch *= hidden
scratch *= 0.044715
scratch += hidden
scratch *= np.float32(np.sqrt(2.0 / np.pi))
np.tanh(scratch, out=scratch)
scratch += 1.0
hidden *= 0.5
hidden *= scratch
def masking_rate_efficiency(rate: float, n_positions: int, width: int, vocab: int) -> dict[str, float]:
"""What the objective spends versus what it learns from.
Illustrative arithmetic, not a benchmark. The ratio is the point: at
the conventional 15%, the encoder processes every position and 85% of
that work produces no gradient. Sweeping the rate on the actual
sequence length is worth doing -- the 15% figure comes from one paper
at one length and does not transfer.
"""
encoder_flops = float(n_positions) * width * width * 12.0
naive_head_flops = float(n_positions) * width * vocab * 2.0
gathered_head_flops = naive_head_flops * rate
return {
"signal_fraction": rate,
"head_flops_saved_fraction": 1.0 - rate,
"head_share_of_naive": naive_head_flops / (encoder_flops + naive_head_flops),
"head_share_of_gathered": gathered_head_flops
/ (encoder_flops + gathered_head_flops),
}
What changed: The dominant waste in a naive implementation is computing vocabulary logits at every position when only fifteen percent of them enter the loss — the head projection is the single largest matmul in the model, so gathering the masked hidden states before projecting cuts that cost by roughly the inverse of the masking rate. Attention scores are computed once as a batched einsum across heads rather than per-head, layer-norm and the residual add are fused so intermediate arrays are never materialized, and the corruption mask is generated as one vectorized draw over the whole batch instead of a Python loop over positions. The loss is computed via a fused log-sum-exp rather than exponentiating and taking a log separately.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
Attention and projections become batched einsums across heads and examples rather than per-position matvecs
Trade-off: The four-dimensional shape contract is no longer visible line by line; the shapes must be asserted rather than read off the loops
Eliminate Python-level loops over samples
Corruption for a whole batch is one vectorized draw instead of a loop over positions and sequences
Trade-off: Every sequence must be padded to a common length, and the padding has to be excluded from both attention and the loss or it silently contributes
Fuse adjacent operations so intermediates are never materialized
The head projection runs on gathered masked positions only, cutting the largest matmul by the inverse of the masking rate
Trade-off: Per-position logits are no longer available for inspection, so probing what the model predicts at unmasked positions requires a second, deliberate pass
Ensure contiguous memory layout and a single dtype
Gathering masked positions produces a compact contiguous block, so the head matmul hits BLAS at full efficiency
Trade-off: The gather itself copies, which costs bandwidth that only pays off when the masking rate is well below one
The head drops to O(rate·n·d·V), roughly a sixfold reduction at the conventional fifteen percent rate, and the encoder itself is unchanged at O(n²·d + n·d²). Illustrative, not a measured benchmark: the encoder cost is the objective’s inefficiency rather than the implementation’s, and no amount of vectorizing removes it.