MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training

• URL: http://arxiv.org/abs/2606.08788v1
• Date: Sun, 07 Jun 2026 19:11:22 GMT
• Title: MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training
• Abstract summary: MaskAlign is a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training.<n>We show that MaskAlign reduces the dependence of representation alignment on the complete token set and encourages alignment behavior that is more stable under token-subset perturbations.
• Score: 53.777757218665386
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Representation alignment with pretrained vision models has recently shown strong potential for accelerating diffusion transformer training. By aligning intermediate diffusion features with clean-image representations from self-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-token representation alignment, tokens with large alignment-gradient norms exhibit 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 propose MaskAlign, a token-subset representation alignment method that applies alignment to randomly sampled token subsets during training. By exposing the model to different token subsets across iterations, MaskAlign reduces the dependence of representation alignment on 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 lightweight pre-mask token mixing block that shares information across tokens before masking.

Related papers

• Where Does Generative Difficulty Reside? An Empirical Study of Target Representations [62.54876287800644]
  The target representation defines the distribution an image generator must learn, yet it is often treated as an interchangeable interface.<n>We study raw pixels, SD-VAE latents and DINOv2 as well as MAE representation-autoencoder features within a unified masked autoregressive rectified-flow model.<n>Our results show that compression, reconstruction fidelity, token dimensionality, and visible semantic clustering do not individually predict generative behavior.
  arXiv  Detail & Related papers  (2026-08-01T12:32:31Z)
• Token-to-Token Alignment of Text Embeddings for Semantic Blending [55.88554156801394]
  We introduce Token-to-Token alignment, a framework that establishes explicit semantic correspondence between tokens across prompts.<n>Our results show that text embedding spaces in text-to-image models implicitly encode a continuous semantic structure that becomes accessible once representations are properly aligned.
  arXiv  Detail & Related papers  (2026-06-22T23:54:40Z)
• Residual Connections and the Causal Shift: Uncovering a Structural Misalignment in Transformers [9.617245548268437]
  Large Language Models (LLMs) are trained with next-token prediction, implemented in autoregressive Transformers.<n>This creates a subtle misalignment: residual connections tie activations to the current token, while supervision targets the next token.<n>We propose a lightweight residual-path mitigation based on residual attenuation, implemented either as a fixed-layer intervention or as a learnable gating mechanism.
  arXiv  Detail & Related papers  (2026-02-16T14:04:42Z)
• Composable Visual Tokenizers with Generator-Free Diagnostics of Learnability [30.139325285692568]
  We introduce CompTok, a training framework for learning visual tokenizers whose tokens are enhanced for compositionality.<n>By employing an InfoGAN-style objective, we train a recognition model to predict the tokens used to condition a diffusion decoder.<n>We show in experiments that CompTok can improve on both of the metrics as well as supporting state-of-the-art generators for class conditioned generation.
  arXiv  Detail & Related papers  (2026-02-03T10:02:51Z)
• Improving Flexible Image Tokenizers for Autoregressive Image Generation [53.238708824055664]
  textbfReToK is a flexible tokenizer with underlineRedundant underlineToken Padding and Hierarchical Semantic Regularization.<n>Our method achieves superior generation performance compared with both flexible and fixed-length tokenizers.
  arXiv  Detail & Related papers  (2026-01-04T14:11:45Z)
• Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models [40.902681492117786]
  RemeDi is a mask-based DLM that predicts token distributions and per-token confidence scores at each step.<n>We train a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens.<n>Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
  arXiv  Detail & Related papers  (2025-09-28T05:39:49Z)
• SeeDiff: Off-the-Shelf Seeded Mask Generation from Diffusion Models [16.109077391631917]
  We show that cross-attention alone provides very coarse object localization, which however can provide initial seeds.<n>We also observe that a simple-text-guided synthetic image often has a uniform background, which is easier to find correspondences.<n>Our proposed method, dubbed SeeDiff, generates high-quality masks off-the-shelf from Stable Diffusion.
  arXiv  Detail & Related papers  (2025-07-26T05:44:00Z)
• Bridging Continuous and Discrete Tokens for Autoregressive Visual Generation [85.82112629564942]
  We propose TokenBridge, which maintains the strong representation capacity of continuous tokens while preserving the modeling simplicity of discrete tokens.<n>We introduce a dimension-wise quantization strategy that independently discretizes each feature dimension, paired with a lightweight autoregressive prediction mechanism.<n>Our approach achieves reconstruction and generation quality on par with continuous methods while using standard categorical prediction.
  arXiv  Detail & Related papers  (2025-03-20T17:59:59Z)
• "Principal Components" Enable A New Language of Images [79.45806370905775]
  We introduce a novel visual tokenization framework that embeds a provable PCA-like structure into the latent token space.<n>Our approach achieves state-of-the-art reconstruction performance and enables better interpretability to align with the human vision system.
  arXiv  Detail & Related papers  (2025-03-11T17:59:41Z)
• Learning to Mask and Permute Visual Tokens for Vision Transformer Pre-Training [55.12082817901671]
  We propose a new self-supervised pre-training approach, named Masked and Permuted Vision Transformer (MaPeT)<n>MaPeT employs autoregressive and permuted predictions to capture intra-patch dependencies.<n>Our results demonstrate that MaPeT achieves competitive performance on ImageNet, compared to baselines and competitors under the same model setting.
  arXiv  Detail & Related papers  (2023-06-12T18:12:19Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.