MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Summary
MaskAlign proposes a token-subset representation alignment method that improves diffusion transformer training by reducing reliance on complete token sets and maintaining stable alignment under perturbations.
View Cached Full Text
Cached at: 06/12/26, 06:51 AM
Paper page - MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
Source: https://huggingface.co/papers/2606.08788 Published on Jun 7
·
Submitted byhttps://huggingface.co/pangly
panglyon Jun 12
Abstract
Token-subset representation alignment method called MaskAlign improves diffusion transformer training by reducing reliance on complete token sets and maintaining stable alignment behavior under perturbations.
Representation alignmentwith pretrainedvision modelshas recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations fromself-supervised vision encoders, existing methods improve convergence and generation quality. However, such alignment also introduces a non-trivial constraint: diffusion models operate on noisy inputs whose usable information varies across timesteps, while the reference features are extracted from clean images. In this paper, we revisit this mismatch from a token-level perspective. We find that, under full-tokenrepresentation alignment, tokens with largealignment-gradient normsexhibit a stable spatial preference, suggesting that the alignment objective does not affect all tokens uniformly and may encourage the model to rely on the complete set of clean-image tokens. To address this issue, we proposeMaskAlign, atoken-subset representation alignmentmethod that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations,MaskAlignreduces the dependence ofrepresentation alignmenton the complete token set and encourages alignment behavior that is more stable under token-subset perturbations. To mitigate the information loss caused by directly dropping tokens, we further introduce a lightweightpre-mask token mixingblock that shares information across tokens before masking.
View arXiv pageView PDFAdd to collection
Get this paper in your agent:
hf papers read 2606\.08788
Don’t have the latest CLI?curl \-LsSf https://hf\.co/cli/install\.sh \| bash
Models citing this paper0
No model linking this paper
Cite arxiv.org/abs/2606.08788 in a model README.md to link it from this page.
Datasets citing this paper0
No dataset linking this paper
Cite arxiv.org/abs/2606.08788 in a dataset README.md to link it from this page.
Spaces citing this paper0
No Space linking this paper
Cite arxiv.org/abs/2606.08788 in a Space README.md to link it from this page.
Collections including this paper0
No Collection including this paper
Add this paper to acollectionto link it from this page.
Similar Articles
Remask, Don't Replace: Token-to-Mask Refinement in Masked Diffusion Language Models
Introduces Token-to-Mask (T2M) remasking to fix generation errors in masked diffusion LMs by resetting suspect tokens to mask state instead of overwriting, yielding up to +5.92 accuracy on CMATH without extra training or parameters.
PreDiff-LM: Pretrained Discrete Masked Diffusion Language Modeling with Hybrid Attention
PreDiff-LM proposes a hybrid attention mechanism that preserves causal attention for prompt tokens and bidirectional attention for masked target tokens, enabling adaptation of pretrained autoregressive models for discrete masked diffusion language modeling, achieving improvements in perplexity and downstream tasks over prior diffusion baselines.
Nemotron-Labs-Diffusion-Image: Advancing Masked Discrete Diffusion for High-Resolution Image Synthesis
This paper proposes Nemotron-Labs-Diffusion-Image, a masked discrete diffusion model for high-resolution text-to-image synthesis, introducing a token-editing mechanism and grouped cross-entropy objective to improve token refinement and training efficiency.
Don't Retrain, Align: Adapting Autoregressive LMs to Diffusion LMs via Representation Alignment
This paper introduces Repr-Align, a method to adapt autoregressive language models into diffusion language models via representation alignment, achieving up to 4x training acceleration without retraining representations from scratch.
Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models
This paper proposes TIE, a knowledge fusion framework for masked diffusion language models that tracks confidence dynamics to identify reliable decoding trajectories and iteratively transfers partially denoised sequences between models, improving generation quality on reasoning tasks.