Typhoon: Gradient-Based Task-Specific Masking for Masked Language Models

arXiv:2303.15619 · cs.CL, cs.AI · Submitted 2023-03-27 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "Typhoon: Gradient-Based Task-Specific Masking for Masked Language Models".

Jane: Here's the gist of the paper. The gist: gradient-based task-adaptive masking, Typhoon, performs on par with random and whole-word masking when seed variance is accounted for, suggesting that this approach does not reliably outperform resource-free methods at this scale.

Tom: Transformer-based language models have driven rapid progress across natural language processing from machine translation to applications in the life sciences Transformer-based language models have driven rapid progress across natural language processing from machine translation to applications in the life sciences In a standard MLM, special

MASK: tokens replace a uniformly random 15% of the input, and the model is trained to recover the hidden tokens from their bidirectional context Recent work shows that content-aware schemes that mask spans, whole words, named entities, or statistically correlated tokens can outperform random masking on knowledge- and reasoning-intensive tasks A natural question follows: if the best masking target depends on the task, can we learn it from the task itself Typhoon Strategy Typhoon is a masking strategy that uses the gradient of the task loss with respect to one-hot token inputs to estimate, online, how much each token type contributes to the objective It computes the gradient of the loss with respect to the one-hot input representation, accumulates a per-token-type saliency score with an exponential moving average for stability across batches, and converts these scores into a masking distribution that respects a target masking budget under a token-independence approximation The key intuition is that tokens whose presence most affects the loss are the tokens whose reconstruction is most informative, and should therefore be masked preferentially The method involves several steps to generate a calibrated masking distribution)&Setup and notation Let V be the vocabulary with V = V.

Jane: An input is a token sequence x = (x1,..., xn) with each xi ∈ V.

Tom: We write ei ∈ 0, 1 for the one-hot encoding of xi and E ∈ R V ×d for the embedding matrix, so that the embedded input is E⊤ei ∈ R d A masking operation selects a subset M ⊆ 1,..., n of positions, replaces x i: i ∈ M with

MASK: , and the MLM is trained to reconstruct them under loss L We target an expected masking rate p ∈ (0, 1): in expectation a fraction p of tokens are masked)&Gradient-based token saliency For each position we take the partial derivative of the loss with respect to the one-hot coordinate of the true token occupying that position, s i = ∂L/∂ei, t i, where t i = x i if position i is unmasked, y i if position i is masked)&Token-type weights via exponential moving average A single batch gives a noisy estimate of token saliency, so we accumulate a weight w v for every token type v ∈ V across batches The running weight is then updated with an exponential moving average that places weight λ ∈ (0, 1) on the new estimate)&From weights to a calibrated masking distribution Weights are not probabilities: a large w v does not directly say how often type v should be masked, and naively masking high-weight types would violate the target budget p To enforce (ii) we maintain a frequency estimate f v, the relative frequency of token type v over the dataset seen so far, with P v f v = 1 Under a token-independence approximation—each position is masked independently given its type—the expected fraction of masked tokens is P v f v π v)&Training pipeline We run the same pipeline on three BERT-family backbones of increasing size: TinyBERT

Jiao et al., 2020: (a 4-layer distilled model), DistilBERT

Sanh et al., 2019: , and BERT-base

Devlin et al., 2019: The pipeline is: Load pretrained backbone weights for masked language modeling, continue MLM training on the target GLUE task’s text, masking with the strategy under test (random, whole-word, or Typhoon), transfer the MLM encoder into a sequence-classification head, and fine-tune the classifier on the target task and evaluate Experimental Results The evaluation was conducted across two GLUE tasks MRPC and CoLA on three BERT-family backbones: TinyBERT, DistilBERT, and BERT-base The comparison involved random masking, whole-word masking, and Typhoon Across all three backbones the three masking strategies land close to one another on the primary metric On MRPC, Typhoon edges out the baselines for BERT-base (0.901 vs.

Paper summary: Tom: So, we’re talking about this paper called Typhoon: Gradient-Based Task-Specific Masking for Masked Language Models. Basically, they look at how you choose which tokens to mask in a language model and they say that standard random masking isn't always the best way to train the model.

Jane: That sounds like a really important question because most of us just use one fixed percentage, usually fifteen percent, for the MASK tokens and hope it works well enough. So, what’s this Typhoon strategy actually trying to do?

Lu: Typhoon uses the gradient of the task loss with respect to one-hot token inputs to estimate online how much each token type contributes to the objective. They build on existing ideas where people mask spans or whole words because those things are often more informative for knowledge and reasoning tasks

Joshi et al., two thousand twenty Cui et al <ref:2303.15619#pg1>., two thousand twenty-one Sun et al <ref:2303.15619#pg1>., two thousand nineteen Levine et al <ref:2303.15619#pg1>., two thousand twenty-one Sadeq et al <ref:2303.15619#pg1>., two thousand twenty-three: <ref:2303.15619#pg1>.

Meng: From a deployment standpoint, it sounds like they are trying to make the masking adaptive based on what the model is currently struggling with during fine-tuning. I wonder if that’s practical when you’re running thousands of inferences per second.

Lalam: The paper says Typhoon maintains an exponential moving average of per-token-type saliency and calibrates those scores into a masking distribution whose expected masking rate matches a target budget under a token-independence approximation <ref:2303.15619#pg2>.

Tom: So, the core idea is that tokens whose presence most affects the loss are the ones that should be masked more often because reconstructing them is more informative for learning <ref:2303.15619#pg2>. It’s using the training signal to guide what we mask.

Jane: But they don't just pick one type of token; they track a frequency estimate, f v, which is the relative frequency of each token type over the data seen so far, and then use that to calculate a masking distribution where the expected fraction of masked tokens is P v f v pi v under a token-independence approximation <ref:2303.15619#pg2>.

Lu: That part about using that frequency estimate, f v, is interesting because it ties the gradient signal to the actual data distribution they’re seeing during training. It moves beyond just looking at the current batch and gives a running weight for every token type <ref:2303.15619#pg2>.

Paper summary: Meng: If we are tracking frequencies, does that mean you have to keep a massive dictionary in memory to track those weights across different batches? I don't want me needing a super large lookup table just to decide what gets masked.

Tom: They use an exponential moving average for stability across batches so the estimates aren't too noisy, which keeps the whole process stable during training <ref:2303.15619#pg2>. It’s about balancing that noise with getting a good signal from the gradients, which is what they are measuring.

Jane: They formalize this by saying that under a token-independence approximation, if you use this budget-preserving min-max calibration, the expected masking rate matches a target budget p <ref:2303.15619#pg2>. So, it's trying to hit a specific masking rate without wildly overshooting or undershooting your plan.

Lu: The key intuition they mention is that tokens whose presence most affects the loss are the tokens whose reconstruction is most informative, and therefore should be masked preferentially <ref:2303.15619#pg2>. It’s a way to learn adaptive masking beyond just fixed heuristics <ref:2303.15619#pg3>.

Meng: I see how it’s designed to respect the budget, but what about those correlations between adjacent tokens? They are ignoring them when they make this token-independence approximation <ref:2303.15619#pg4>. That sounds like a potential weakness for real-world text where context is everything.

Tom: Yeah, they explicitly state that the token-independence approximation ignores correlations between adjacent tokens, which is what span or PMI methods try to exploit directly <ref:2303.15619#pg4>. It’s a limitation they acknowledge upfront.

Jane: So, if we look at the results mentioned in the summary, they tested this strategy against random masking and whole-word masking on two tasks, MRPC and CoLA, across three BERT backbones: TinyBERT, DistilBERT, and BERT-base <ref:2303.15619#pg1>.

Lu: The evaluation showed that when accounting for seed variance—meaning they ran the experiment multiple times with different starting points—no masking strategy was reliably better than the others on these tasks <ref:2303.15619#pg4>.

Meng: So, even though Typhoon seemed promising in some single runs, when you test it rigorously across different random seeds, it doesn't beat the simpler methods like random masking or whole-word masking at this scale <ref:2303.15619#pg4>. That’s a big deal for deployment if we’re trying to save compute on the data augmentation side.

Tom: I agree with Meng, and I think it points to how much overhead is required for these gradient-based methods when you are just comparing them against standard baselines <ref:2303.15619#pg4>. The paper shows that the gap between Typhoon and the best baseline stays within zero point zero zero four F1 on MRPC, and across all twelve Typhoon comparisons, no paired test reached significance <ref:2303.15619#pg4>.

Paper summary: Jane: So even though they built a sophisticated gradient-based system, when you factor in the randomness of the training process, it doesn't reliably outperform uniform random masking on MRPC <ref:2303.15619#pg4>. What does that imply about learning from the task itself?

Lu: It implies that for this specific setup and scale, gradient-based task-adaptive masking isn't necessarily better than just picking a fixed rule, but it opens up the idea of learning those rules adaptively <ref:2303.15619#pg3>. The mechanism itself is still useful for understanding what tokens matter most.

Meng: For me, if I were deploying this, I’d be skeptical about the complexity versus the performance gain on a standard production model running on typical traffic <ref:2303.15619#pg4>. It adds a lot of complexity to the training pipeline without delivering a reliable lift in accuracy over simpler methods.

Tom: Exactly, and looking at the explanation for why it doesn't win, they show that budget-preserving min-max calibration compresses the gradient ranking into a narrow band <ref:2303.15619#pg4>. It’s practically operationally indistinguishable from uniform random masking because its Kullback-Leibler divergence from uniform random masking is seven times ten-four nats per token <ref:2303.15619#pg4>.

Jane: So, even when you look inside the mechanism, the calibration is actually very close to random masking in terms of information content per token, which means widening that calibration spread by a factor of ten still didn't help downstream accuracy within seed noise <ref:2303.15619#pg4>. That’s a pretty strong piece of evidence.

Lu: It suggests the advantage they showed in single-run experiments doesn't survive multi-seed evaluation, which is a crucial piece of evidence for whether this approach is fundamentally better or just lucky in one trial <ref:2303.15619#pg4>. The authors are being cautious about overstating the findings.

Meng: If we’re shipping models, we need reliability, not just high scores on a few random runs <ref:2303.15619#pg4>. So, the reality is that gradient-based task-adaptive masking is competitive with random masking at this scale because it doesn't reliably outperform it <ref:2303.15619#pg4>.

Tom: It’s a neat illustration of how much research in this area still has to go, showing that learning the optimal masking target from the task itself is an idea, but not necessarily a winning strategy at this specific implementation level <ref:2303.15619#pg1>. We can see value in publishing this corrected picture rather than the original optimistic one <ref:2303.15619#pg4>.

Paper summary: Jane: It’s a fair conclusion because it sets a benchmark for future work, showing where we need to focus next when we ask if the best masking target depends on the task <ref:2303.15619#pg2>. So, while Typhoon didn't beat random masking here, it gives us a concrete way to probe that relationship.

Lu: That’s what I think—it provides a solid baseline for where we need to look next when we explore learned and adaptive masking <ref:2303.15619#pg3>. It shows the intuition is sound, even if the current calibration method doesn't deliver a win yet.

Meng: For production, that means sticking to simpler, more robust methods unless we can find a way to make that gradient signal more reliable under real traffic conditions <ref:2303.15619#pg4>. It’s not worth the engineering complexity right now if it doesn't give us a measurable gain over the established baselines.

Tom: So, to wrap up, Typhoon is a way to use gradients to estimate token importance for masking, and while it’s interesting theoretically, its practical performance doesn't beat random masking when we account for seed variance <ref:2303.15619#pg4>. It proves the concept is valid but the specific implementation isn't superior yet.

Jane: And that means for now, if you want to improve performance on reasoning tasks, you might need to look at other ways to define informativeness than just the task gradient <ref:2303.15619#pg3>. What we learned here is that we have a clear idea of what the current method does and where its limitations lie.

Lu: Yeah, it shows us that while we can learn masking policies, they need to be carefully calibrated to avoid being overly sensitive to noise in the training process <ref:2303.15619#pg4>. It’s a step in the right direction for exploring learned and adaptive masking <ref:2303.15619#pg3>.

Meng: Right, so for now, we stick to what’s proven reliable on real traffic instead of chasing complex gradient-based heuristics that don't hold up under rigorous testing <ref:2303.15619#pg4>. It keeps the deployment pipeline simpler.

Tom: Exactly, and that’s where I sit—all about numbers and what actually runs reliably <ref:2303.15619#pg4>. The paper lays out the mechanics clearly, but the results show that at this point, it’s competitive but not better than uniform random masking <ref:2303.15619#pg4>.

Jane: So we get a clearer picture of what works and what doesn't for content-aware masking, which is valuable context for anyone working in this space <ref:2303.15619#pg2>. It’s about understanding the signal before trying to build the perfect filter.

Conclusion: Tom: So, Typhoon is this new way to choose which tokens to mask in masked language modeling by using gradients from the loss function itself.

Jane: That’s right, Tom, and it’s essentially trying to make the masking strategy adapt based on what the model is actually learning during fine-tuning.

Lu: The paper introduces Typhoon as a masking strategy that uses the gradient of the MLM loss with respect to one-hot token inputs to estimate online how much each token type contributes to that objective.

Meng: So, it’s not just picking random tokens; it’s trying to figure out which tokens are most important for the task by looking at the gradients.

Lalam: The core idea is that tokens whose presence most affects the loss are the ones whose reconstruction is most informative and should therefore be masked preferentially.

Tom: And they do this online using an exponential moving average of per-token-type saliency to keep things stable across batches.

Jane: But what’s the actual result? Does this adaptive approach actually beat just picking a fixed percentage of tokens?

Lu: Well, the experiments show that when you account for seed variance, no masking strategy, including Typhoon, is reliably better than uniform random masking on the tasks they tested.

Meng: Wait, so after all that gradient work and calibration with budgets and approximations—it doesn't actually outperform the simple random method in a rigorous test?

Tom: Exactly. The paper shows that budget-preserving min-max calibration compresses the gradient ranking into a very narrow band, making it practically indistinguishable from uniform random masking in terms of information content per token.

Jane: So, even though they built this complicated system based on gradients and task objectives, the performance gain over a fixed random percentage doesn't reliably show up at this scale.

Lu: The authors point out that their token-independence approximation means they ignore correlations between adjacent tokens, which is what other methods try to exploit.

Meng: And what does that mean for deployment? It adds a lot of complexity to the training pipeline without giving us a measurable lift in accuracy over established methods on real traffic.

Tom: So, the implication is that while learning adaptive masking from the task seems like a good idea, at this stage, we need to be more careful about how we calibrate those gradient estimates and what assumptions we are making.

Jane: It gives us a concrete way to probe that relationship—the question of whether the best masking target depends on the specific task—but it also shows us where the current method stops working reliably without more advanced calibration.

Lu: The paper sets a benchmark here, showing that for now, this gradient-based approach is competitive with random masking but not superior at this scale.

Meng: It’s a lot of theory if it doesn't translate to a better deployed model than what we already have on the table.

Tom: So, for now, we stick to what’s proven reliable on real traffic instead of chasing complex gradient-based heuristics that don't hold up under rigorous testing.

Muhammed Shahir Abdurrahman, Hashem Elezabi, Bruce Changlong Xu

Stanford University

cs.CL, cs.AI

Submitted: 2023-03-27

Updated: 2026-10-20

License: http://creativecommons.org/licenses/by/4.0/

The gist: The gist: gradient-based task-adaptive masking, Typhoon, performs on par with random and whole-word masking when seed variance is accounted for, suggesting that this approach does not reliably

Key concepts

MLM
Masked Language Modeling is the standard pretraining task where special [MASK] tokens replace randomly selected words in text. The model learns to predict these original hidden words based on the surrounding context, which is how most language models are initially trained.
Token Saliency
This refers to a measure of how much a specific type of token (like 'person' or 'verb') contributes to the overall loss during training. Typhoon uses the gradient of this loss to calculate saliency, identifying which tokens are most informative for learning.
Budget-Preserving Calibration
Since masking too many important tokens violates the target budget, this technique adjusts the calculated token importance scores into a distribution that respects the budget constraint. This ensures that only the most relevant tokens are prioritized for masking.
Token-Independence Approximation
This assumption simplifies how Typhoon calculates which positions to mask. It treats each position's masking decision as independent of other positions, making the calibration process more tractable while still aiming to respect the overall budget.

Terminology

Summary

The gist: gradient-based task-adaptive masking, Typhoon, performs on par with random and whole-word masking when seed variance is accounted for, suggesting that this approach does not reliably outperform resource-free methods at this scale.

Introduction

Transformer-based language models have driven rapid progress across natural language processing from machine translation to applications in the life sciences Transformer-based language models have driven rapid progress across natural language processing from machine translation to applications in the life sciences In a standard MLM, special [MASK] tokens replace a uniformly random 15% of the input, and the model is trained to recover the hidden tokens from their bidirectional context Recent work shows that content-aware schemes that mask spans, whole words, named entities, or statistically correlated tokens can outperform random masking on knowledge- and reasoning-intensive tasks A natural question follows: if the best masking target depends on the task, can we learn it from the task itself

Typhoon Strategy

Typhoon is a masking strategy that uses the gradient of the task loss with respect to one-hot token inputs to estimate, online, how much each token type contributes to the objective It computes the gradient of the loss with respect to the one-hot input representation, accumulates a per-token-type saliency score with an exponential moving average for stability across batches, and converts these scores into a masking distribution that respects a target masking budget under a token-independence approximation The key intuition is that tokens whose presence most affects the loss are the tokens whose reconstruction is most informative, and should therefore be masked preferentially

Methodology

The method involves several steps to generate a calibrated masking distribution

**)&Setup and notation Let V be the vocabulary with V = V. An input is a token sequence x = (x1,..., xn) with each xi ∈ V. We write ei ∈ 0, 1 for the one-hot encoding of xi and E ∈ R V ×d for the embedding matrix, so that the embedded input is E⊤ei ∈ R d A masking operation selects a subset M ⊆ 1,..., n of positions, replaces x i: i ∈ M with [MASK], and the MLM is trained to reconstruct them under loss L We target an expected masking rate p ∈ (0, 1): in expectation a fraction p of tokens are masked **

**)&Gradient-based token saliency For each position we take the partial derivative of the loss with respect to the one-hot coordinate of the true token occupying that position, s i = ∂L/∂ei, t i, where t i = x i if position i is unmasked, y i if position i is masked **

**)&Token-type weights via exponential moving average A single batch gives a noisy estimate of token saliency, so we accumulate a weight w v for every token type v ∈ V across batches The running weight is then updated with an exponential moving average that places weight λ ∈ (0, 1) on the new estimate **

**)&From weights to a calibrated masking distribution Weights are not probabilities: a large w v does not directly say how often type v should be masked, and naively masking high-weight types would violate the target budget p To enforce (ii) we maintain a frequency estimate f v, the relative frequency of token type v over the dataset seen so far, with P v f v = 1 Under a token-independence approximation—each position is masked independently given its type—the expected fraction of masked tokens is P v f v π v **

**)&Training pipeline We run the same pipeline on three BERT-family backbones of increasing size: TinyBERT [Jiao et al., 2020] (a 4-layer distilled model), DistilBERT [Sanh et al., 2019], and BERT-base [Devlin et al., 2019] The pipeline is: Load pretrained backbone weights for masked language modeling, continue MLM training on the target GLUE task’s text, masking with the strategy under test (random, whole-word, or Typhoon), transfer the MLM encoder into a sequence-classification head, and fine-tune the classifier on the target task and evaluate **

Experimental Results

The evaluation was conducted across two GLUE tasks MRPC and CoLA on three BERT-family backbones: TinyBERT, DistilBERT, and BERT-base The comparison involved random masking, whole-word masking, and Typhoon Across all three backbones the three masking strategies land close to one another on the primary metric On MRPC, Typhoon edges out the baselines for BERT-base (0.901 vs. 0.900 F1) but trails whole-word masking for DistilBERT (0.895 vs. 0.899) On CoLA, whole-word masking is marginally better than random masking on DistilBERT MRPC (∆ = 0.006 F1, p = 0.03) Across all three backbones the three masking strategies land close to one another on the primary metric

Mechanistic Explanation

The null result is more informative once we open up the algorithm After three epochs of intermediate MLM training on MRPC, BERT-base’s Typhoon trainer has seen 1.95×10 5 tokens and accumulated weights for 11,563 of the 30,522 vocabulary entries Figure 3 plots the frequency-weighted histogram of πv: under the default spread δ = 0.4, the entire distribution lives in the band [0.22, 0.37] around the budget p = 0.3, with a frequency-weighted standard deviation of 0.017 The budget-preserving calibration squashes the gradient ranking into a distribution operationally indistinguishable from uniform random masking (KL ≈ 7×10-4 nats per token)

Limitations and Conclusion

Typhoon makes a token-independence approximation when calibrating masking probabilities; correlations between adjacent tokens (which span- and PMI-based methods exploit directly) are ignored The null result we report may not hold for longer task-adaptive pretraining, larger models, knowledge-intensive tasks, or higher masking budgets We see value in publishing this corrected picture rather than the original optimistic one The baseline established here should help future work revisiting this question.

A Full Algorithms

The two routines in §3 are run jointly during the intermediate MLM phase Algorithm 1 updates token-type weights from the gradients of each batch; Algorithm 2 samples the mask for each sequence from the current weights subject to the budget Token types never yet observed are assigned the prior masking probability p.

Hyperparameters

The key hyperparameters, shared across all three backbones and both tasks, include a masking budget p of 0.3, an EMA decay λ of 0.8, a probability spread δ of 0.4, and an optimizer AdamW The final column is Typhoon’s per-step overhead relative to random masking.

Table 4

Table 4 reports the paired comparisons summarized in §4.

Improvements for AI systems

  1. Task-adaptive masking for fine-tuning enables models to focus on tokens whose reconstruction is most informative for a specific downstream objective, as Typhoon uses the gradient of the task loss with respect to one-hot token inputs to estimate, online, how much each token type contributes to the objective. This allows for better representation learning during intermediate pretraining.

  2. The system can be improved by using Typhoon during an intermediate masked language modeling phase before fine-tuning on a specific task. This phase is beneficial because MLM loss alone is not predictive of fine-tuned quality, suggesting task-specific masking helps shape representations better than fixed rules.

  3. The improved system will exhibit superior performance on MRPC, as the paper shows that on MRPC the gap between Typhoon and the best baseline stays within 0.004 F1 once seed variance is accounted for, implying it is competitive with random masking for this task under controlled conditions.

  4. The system can be enhanced by using Typhoon to generate a masked set M based on calibrated probabilities πv, where with probability πxi, add i to M, allowing the model to learn from a distribution that respects a target budget p. This provides more targeted learning than uniform random masking.

Abstract

The choice of which tokens to mask is an under-examined design decision in masked language modeling (MLM). Standard pretraining masks tokens uniformly at random, whereas content-aware masking schemes can improve downstream transfer. We study masking as a task-adaptive component of the fine-tuning pipeline and introduce Typhoon, a masking strategy that uses the gradient of the MLM loss with respect to one-hot token inputs to estimate online how much each token type contributes to the objective. Typhoon maintains an exponential moving average of per-token-type saliency and calibrates these scores into a masking distribution whose expected masking rate matches a target budget under a token-independence approximation. We formalize the method and evaluate it against random masking and whole-word masking on two GLUE tasks, MRPC and CoLA, across three BERT-family backbones (TinyBERT, DistilBERT, and BERT-base) and five random seeds per configuration (90 training runs in total). Once seed variance is accounted for, no masking strategy is reliably better than the others on these tasks: on MRPC the gap between Typhoon and the best baseline stays within 0.004 F 1, across all twelve Typhoon comparisons no paired test reaches significance, and every 95% confidence interval contains zero. Inspecting Typhoon's learned distribution explains why: budget-preserving min-max calibration compresses the gradient ranking into a narrow band whose Kullback-Leibler divergence from uniform random masking is 7 times 10-4 nats per token, and widening the calibration spread by an order of magnitude in KL still leaves downstream accuracy unchanged within seed noise. Typhoon's apparent advantage in single-run experiments does not survive multi-seed evaluation, showing that gradient-based task-adaptive masking is competitive with, but does not outperform, uniform random masking at this scale.

Sources

Related papers