Manifold Projection and Iterative Autoencoder Refinement for Masked Language Modeling

Narges Mokhtari; Farzan Haddadi; Ebrahim Rezaii

cs.CLSep 9, 2026

Manifold Projection and Iterative Autoencoder Refinement for Masked Language Modeling

Authors: Narges Mokhtari, Farzan Haddadi, Ebrahim Rezaii

Organizations: School of Electrical Engineering, Iran University of Science & Technology, Tehran, Iran.

Abstract

In Transformer-based masked language models, attention is the primary mechanism for context mixing, but there are other ways to mix data across tokens. Recent attention-free mixers replace attention with fixed or hypernetwork-generated MLPs, alternating their dynamic, content-dependent weighting for computational simplicity. We build an alternative that gets the same property from a low-rank bottleneck autoencoder. We replace attention with a stack of autoencoder-based mixing modules, one operating over local neighborhoods, one over the full sequence, and one across attention heads, each compressing and reconstructing its input through a bottleneck, and its width is a hyperparameter rather than a training effect. In masked positions, we introduce an iterative refinement procedure that has two distinct steps. A pulling step that pulls an embedding representation toward a weighted average of its neighbors, and a correcting step that projects the result back to the learned manifold via an autoencoder. Our architecture achieves a significant portion of attention's performance at about $1.9 \times$ fewer FLOPs when pretrained on C4 and evaluated with parameter-matched BERT baselines. Our model equals parameter-matched BERT and TinyBERT baselines on the rarest-token frequency bucket using a frequency-aware training schedule that samples rare tokens more than uniformly for the masking tasks.

Figures & tables

Fig. 1: The model architecture

Fig. 2: Internal structure of the mixing modules.

Panel of Figure 3 — Fig. 3: The iterative refinement procedure consists of two steps: pulling and correcting. Starting from $p^{(t)}$ , the pulling step moves the embedding toward the neighbor mean $\mu$ , landing off the manifold at $\hat{p}$ . Then in the correcting step $\mathrm{AE}_{\mathrm{proj}}(\hat{p})$ project this back onto the manifold, giving $p^{(t+1)}$ .

Config.	$D$	$L$	Params (M)	Time/step	Score (%)	[100,2k)	[2k,5k)	[5k)	PPL
Base	680	512	32	51s	38.1	51.2	37.3	23.4	33.1
Base ( $D{=}384$ )	384	512	16.8	39s	34.4	47.4	33.4	19.9	41.4
B: ChunkedContentBias	680	512	32	59s	41.2	58.3	38.9	24.9	27.4
C: Extended training dynamics	680	128	31	52s	37.8	50.6	37.4	22.2	36.6

TABLE I: Comparison across AE-mixing configurations, all on C4, 100k steps. Bucket columns report masked-token accuracy (%) for frequency ranges

[100,2\mathrm{k})

[2\mathrm{k},5\mathrm{k})

, and

[5\mathrm{k})

Model	$D$	$L$	Batch	Params (M)	Time/step	Score (%)	[100,2k)	[2k,5k)	[5k)	PPL
Our model	680	512	128	32	51s	38.1	51.2	37.3	23.4	33.1
TinyBERT	384	128	192	31	61s	43.1	70.2	40.9	22.0	25.8
BERT *	680	512	128	28	65s	46.6	71.9	42.1	24.4	20.9
BERT *	680	128	192	32	71s	46.9	74.0	44.3	24.8	18.0

TABLE II: Our AE base configuration versus three BERT-based baselines, matched in parameter count. * BERT here is standard self-attention Transformer encoders trained from scratch under our own setup (same optimizer, masking scheme, and steps as our model).

Task	Metric	AE (+ContentBias)	Attention ( $D{=}384,L{=}128$ )	Attention ( $D{=}680,L{=}128$ )	Attention ( $D{=}680,L{=}512$ )
CoLA	Matthews corr.	25.9 $\pm$ 1.5	17.7	19.0	21.9 $\pm$ 1.3
SST-2	Accuracy	85.5 $\pm$ 0.8	80.6	84.0	84.2 $\pm$ 0.6
MRPC	Accuracy / F1	69.5 $\pm$ 0.7 / 79.9 $\pm$ 0.3	72.4 / 80.9	73.7 / 81.8	82.1 $\pm$ 0.8 / 87.1 $\pm$ 1.2
STS-B	Pearson / Spearman	65.5 $\pm$ 0.9 / 65.3 $\pm$ 1.1	79.9 / 79.6	82.4 / 82.1	83.0 $\pm$ 0.4 / 82.8 $\pm$ 0.4
QQP	Accuracy / F1	79.8 $\pm$ 0.2 / 73.8 $\pm$ 0.8	63.1 / 60.2	83.9 / 78.4	83.3 $\pm$ 0.1 / 76.7 $\pm$ 0.4
MNLI	Matched / Mismatched	53.2 $\pm$ 0.6 / 53.8 $\pm$ 0.8	31.8 / 31.8	66.2 / 67.1	63.7 $\pm$ 0.1 / 64.7 $\pm$ 0.2

TABLE III: GLUE validation results over 3 seeds: our best AE configuration (base

+

ChunkedContentBias) versus the three attention-based baselines of Table II .

Configuration	FLOPs	FLOPs
	( $L{=}128$ )	( $L{=}512$ )
Full model	114.7 G	477 G
WindowMixAE $\times 3$	63.95 G	255.79 G
GlobalMixAE	1.40 G	22.36 G
ChannelMixAE	16.3 G	65.3 G
ManifoldAE reconstruction (training pass)	3.47 G	13.85 G

TABLE IV: FLOPs per forward pass, base + ChunkedContentBias configuration (batch size 32).

Refinement iterations $T$
Value	Loss	PPL	Acc. (%)
2	10.4	33k	6.4
3	6.9	1k	15
4	4.5	96	28
5	3.6	40	36
6	3.5	33	38

TABLE V: Ablations over the refinement loop’s three hyperparameters: number of iterations

T

, mixing ratio

\alpha

, and step size

\gamma

Masked Language Modeling

Masked language modeling trains text encoders to infer hidden tokens, producing contextual representations for language tasks and specialized domains.

See all 47 papers

Typhoon: Gradient-Based Task-Specific Masking for Masked Language Models
Mar 27, 2023Bruce Changlong Xu, Muhammed Shahir Abdurrahman, Hashem ElezabiMasked Language Modeling Fine-Tuning Language Models
Department of Computer Science, Stanford University
Looped Diffusion Language Models
May 25, 2026Sanghyun Lee, Chunsan Hong, Seungryong Kim +3Masked Diffusion Models Diffusion Language Model Inference
1KAIST · 2KRAFTON · University of California, Berkeley
Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
Jun 9, 2026Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher +1Masked Diffusion Models Parallel Decoding
University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland

Mar 27, 2023J/K move · Enter open · S save

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

Bruce Changlong Xu, Muhammed Shahir Abdurrahman, Hashem Elezabi

Department of Computer Science, Stanford University

Masked Language Modeling Fine-Tuning Language Models

Open paper AskPDF arXiv

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.

Abstract

Masked Language Modeling

Explore similar work

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

Looped Diffusion Language Models

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models