Paper deep dive
End-to-End Training for Unified Tokenization and Latent Denoising
Shivam Duggal, Xingjian Bai, Zongze Wu, Richard Zhang, Eli Shechtman, Antonio Torralba, Phillip Isola, William T. Freeman
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 96%
Last extracted: 3/26/2026, 2:48:55 AM
Summary
UNITE is a single-stage autoencoder architecture that unifies tokenization and latent diffusion by using a shared Generative Encoder (GE). By treating tokenization and generation as the same latent inference problem under different conditioning regimes, UNITE jointly optimizes both tasks without adversarial losses or pretrained encoders, achieving near state-of-the-art performance on ImageNet.
Entities (4)
Relation Signals (3)
UNITE â usescomponent â Generative Encoder
confidence 100% ¡ UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator
Generative Encoder â performstask â Tokenization
confidence 95% ¡ Generative Encoder that serves as both image tokenizer and latent generator
Generative Encoder â performstask â Latent Generation
confidence 95% ¡ Generative Encoder that serves as both image tokenizer and latent generator
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.
Tags
Links
- Source: https://arxiv.org/abs/2603.22283v1
- Canonical: https://arxiv.org/abs/2603.22283v1
Trouble viewing inline? Open PDF directly â
Full Text
79,129 characters extracted from source content.
Expand or collapse full text
End-to-End Training for Unified Tokenization and Latent Denoising Shivam Duggal 1 * Xingjian Bai 1 * Zongze Wu 2 Richard Zhang 2 Eli Shechtman 2 Antonio Torralba 1 Phillip Isola 1 William T. Freeman 1 Abstract Latent diffusion models (LDMs) enable high- fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE â an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder.The shared parameters enable gradients to jointly shape the latent space, encouraging a âcommon latent languageâ. Across image and molecule modalities, UNITE achieves near state-of-the-art performance without adversarial losses or any pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet256 Ă 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single-stage joint training of to- kenization and generation from scratch is feasible. Code:https://github.com/Shi vamDuggal4/UNITE-tokenization-g eneration Project Page:https://xingjianbai.co m/unite-tokenization-generation/ * Equal contribution 1 Massachusetts Institute of Technology 2 Adobe. Decoder ... Reconstruction Generation Generative Encoder Images Gaussian Prior Figure 1. Unified Tokenization & Generation via Generative Encoder (UNITE): We propose a single-stage architecture that unifies tokenization and generation through shared parameters. The Generative Encoder operates in two modes: (top) as a tok- enizer, it processes image patches and register tokens to produce a latent representationz 0 ; (bottom) as a generator/denoiser, it evolves latents along a flow-matching trajectory to synthesizez 0 from Gaussian noise. The latent space is jointly shaped by recon. & generative objectives from scratch, without external supervision. 1. Introduction Modern foundation models (Brown et al., 2020; Koroteev, 2021; Radford et al., 2021; Chen et al., 2023; Esser et al., 2024; Polyak et al., 2024; Wan et al., 2025)âfrom language models to video generators, visionâlanguage systems, and scientific generative modelsâare built around two core op- erations: tokenization and generation. Tokenization maps high-dimensional observations into a compact latent space that enables both faithful reconstruction and efficient dis- crimination; generation learns a distribution over this space to synthesize plausible new samples. This division naturally suggests a sequential recipe: first learn a representation space that is easy to reconstruct from and useful for down- stream computation; then learn a generative process that samples from that space. As a result, most systems treat tokenization and generation as separate design problems & train them in stagesâlearning a tokenizer, freezing it & only then fitting a generator on the induced latent distribution. This separation is convenient, but it departs from the prin- ciple of end-to-end learning and leaves a basic question unresolved: should tokenization and generation be trained jointly so that each objective can shape the learned latent space? In a joint setting, generative pressure could sculpt the latent space toward regions that are easier to model, while reconstruction and inference pressure could preserve instance-specific information and semantic structure. Un- 1 arXiv:2603.22283v1 [cs.CV] 23 Mar 2026 End-to-End Training for Unified Tokenization and Latent Denoising derstanding what emerges when these objectives are trained togetherâand whether their interaction helps or hurtsâis the starting point of this work. A natural way to pursue this idea is to start from the standard latent generative pipeline. A tokenizer is typically learned as part of an autoencoder with an encoderEand decoderD: the encoder maps an imagexto a latent sequencez = E(x) and the decoder reconstructsËx = D(z). A generator then models the latent distribution, most commonly by training a diffusion/flow denoiser (Ho et al., 2020; Song et al., 2021; Lipman et al., 2023) on noisy versions ofz: sample a noise levelt, formz t by corruptingz, and learn a network that predicts the clean latent (or an equivalent parameterization) so that new samples can be generated by starting from noise and iteratively denoising in latent space. In a joint training setting, the same latentzmust therefore serve two purposes: it must be decodable byDto preserve instance information, and it must be structured in a way that makes the denoising objective well-posed and easy to learn. Prior works have explored fully end-to-end training of la- tent diffusion models by backpropagating the denoising objective through the tokenized latents and into the encoder. However, when the tokenizer and diffusion model are opti- mized primarily through the denoising objective, this can lead to degenerate solutions and poor performance, as ob- served in REPA-style methods (Yu et al., 2025a; Leng et al., 2025). To address this, these works propose anchoring the tokenizer with an additional objective that aligns diffusion features to pretrained visual encoders. Although effective, this strategy introduces a third componentâa pretrained teacherâto stabilize joint optimization. In contrast, our set- ting relies only on reconstruction and denoising objectives to jointly train the tokenizer and latent generative model, without any external supervision. We propose an alternative perspective on end-to-end train- ing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes (see Fig. 2). Tokenization can be viewed as a generative process under strong observ- ability: given a data pointx, the model induces a highly concentrated (near-single-point) distribution over latents, yielding a latentzthat is consistent with and informative aboutx. Generation corresponds to a weak-observability regime, wherezmust be synthesized from noise (and op- tional conditions) using the learned prior. Under this view, these two operations differ mainly in how much information is availableâfrom the full observationxin tokenization to only a prior in generation. Motivated by this view, we propose UNITE, which jointly trains tokenization and gen- eration end-to-end without external supervision. UNITE ties tokenization & generation through a shared-parameter module we call the Generative Encoder (GE), so that gradi- Image Manifold GE θ z 0 z t GE θ x Noisy Latent Distribution Latent Manifold Latent Gen. Tokenization Figure 2. Tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes. In tokenization, the full observationxstrongly constrains the clean latent,z 0 âź p θ (z | x); in generation, a noisy latentz t provides weaker evidence, and the same target latentz 0 is recovered by denoising. This view motivates using a single shared Generative Encoder (GE θ ) for both tokenization and latent denoising. ents from both objectives directly shape the same weights, pushing the model toward a representation that is jointly op- timal for the two tasks. Hence the name, UNITE:Unifying Tokenization & Latent Generation via shared Generative Encoder. See Fig. 1 for an overview. Concretely, our system consists of only two modules: a Generative EncoderGE θ and a decoderD Ď . The GE oper- ates in two modes: (i) tokenization, mapping an inputxto latent tokensz = GE θ (x), and (i) generation, denoising corrupted latents to produceËz = GE θ (z t ,t)at noise level t. Thus, the same network serves both as the tokenizer and as the multi-step latent denoiser, with parametersθshared across the two objectives. Training proceeds with two for- ward passes throughGE θ . First, we tokenize an input image to obtain clean latentsz. We then corruptzusing a rectified- flow (flow-matching) process to obtainz t , and passz t back throughGE θ to predict the corresponding denoising target. The full pipeline is trained end-to-end in a single stage by jointly optimizing a pixel-space reconstruction objective and a latent-space flow-matching objective. We find that this end-to-end formulation yields a strong la- tent generative model, with near state-of-the-art generation and reconstruction fidelity, while training all modules from scratch rather than relying on large pretrained networks . To understand what drives this behavior, we study other alter- natives to end-to-end training in Sec. 4. This includes an ab- lation that keeps the full training pipeline fixed but removes parameter tying between the encoder and denoiser. Interest- ingly, even without explicit weight sharing, the encoder and denoiser exhibit strong per-layer representational alignment, as measured by centered kernel alignment (CKA) (Kornblith et al., 2019) (See Fig. 6), suggesting that tokenization and denoising are intrinsically compatible tasks in our setting. 2 End-to-End Training for Unified Tokenization and Latent Denoising Further analysis (see Sec. 4) indicates that the model differ- entiates the two modes primarily through normalization: the tokenization and denoising pathways occupy different nor- m/scale regimes, while attention and MLP sublayers remain highly reusable across both. In fact, recent concurrent work on Unified Latent (Heek et al., 2026) investigates a closely related two-module formulation. It can be interpreted as a special case of our end-to-end setting, aligning closely with our separate-weights ablation. While this separate-weights variant performs almost as competitively, we find that pa- rameter tying yields the best overall rFID/gFID trade-off in our experiments (Fig. 5). Overall, these results provide a concrete single-stage recipe in which reconstruction and de- noising objectives can jointly shape the latent space, rather than being optimized in disjoint stages. Practically, this means one training job and one model to store and update, while retaining a near-SOTA tokenizer and generator. 2. Related Work Tokenization & Generation via Auto-Encoding: Vari- ational autoencoders (VAEs) (Kingma & Welling, 2014) introduced a principled framework for learning probabilistic latent representations while enabling generation through a simple Gaussian prior. This foundational work established that reconstruction and generation can be learned within a single model, though the Gaussian prior and likelihood assumptions often limit sample quality. Extensions such as VQ-VAE (Van Den Oord et al., 2017) and VQ-GAN (Esser et al., 2021) improved representation learning by introducing discrete latent spaces and adversarial training, respectively; in practice, many widely used autoencoder- based tokenizers for diffusion are trained with GAN-style (Goodfellow et al., 2020) losses. However, in modern la- tent diffusion pipelines (Peebles & Xie, 2023; Ma et al., 2024), these VAE/VQGAN-style models primarily serve as tokenizers for a downstream diffusion model trained in the resulting frozen latent space; their standalone generative capability is typically weaker and is therefore rarely used in practice. In standard downstream diffusion training, the denoising/generative gradients never flow back into the tok- enization process, preventing the representation from being shaped by the needs of generation. To address this, we cou- ple tokenization with a latent denoising objective and train a single model end-to-end, allowing the encoder to be directly shaped by generative learning. For simplicity, we eliminate adversarial losses in all experiments unless otherwise stated. Self-supervised Visual Encoders for Generation: Re- cent advances in self-supervised learning have produced powerful visual encoders that go beyond naive reconstruc- tion objectives. Masked autoencoders (MAE) (He et al., 2022) show that reconstructing masked patches can learn strong visual representations at scale. DINO-style mod- els (Caron et al., 2021; Oquab et al., 2023) learn semantic features via self-distillation without labels, yielding repre- sentations that capture both local and global image structure. Building on these encoders, recent methods such as REPA (Yu et al., 2025a), REPA-E (Leng et al., 2025), and RAE (Zheng et al., 2025) leverage pretrained SSL models as extra supervision for diffusion model training. REPA improves training efficiency and sample quality by aligning interme- diate diffusion features with SSL representations. REPA-E extends this idea by jointly tuning the VAE and diffusion model to better match the SSL space, while RAE replaces the VAE encoder with an SSL encoder and trains a separate decoder in a subsequent stage for reconstruction. While these approaches achieve strong generation quality, they further increase pipeline staging and do not study how re- construction and generation can jointly shape shared model parameters. In contrast, we focus on a single-stage train- ing approach that learns tokenization and generation jointly, without access to pretrained SSL encoders. Pixel-Space Diffusion Models: Pixel-space diffusion models denoise directly in the RGB domain, avoiding a learned latent space but facing sharper scaling issues at high resolution. As resolution increases, stronger local redun- dancy lets fixed noise be averaged out, raising effective SNR and making denoising too easy; thus prior work scales noise (or reweights the loss) to keep difficulty/SNR consistent across resolutions (Hoogeboom et al., 2023; Chen, 2023; Kingma & Gao, 2023). This motivates architectural adap- tations tailored to high-resolution pixel modeling: SiD2 (Hoogeboom et al., 2024) trims U-Net skip connections and reduces high-resolution feature capacity; PixelFlow (Chen et al., 2025a) alternates denoising with progressive upsam- pling; and methods such as PixNerd (Wang et al., 2025), PixelDiT (Yu et al., 2025b), and DiP (Chen et al., 2025b) introduce specialized heads to better handle fine-grained inputs. JiT (Li & He, 2025) takes a complementary min- imalist stance, training a plain ViT generator directly on raw patches without tokenizers, pretraining, or auxiliary losses. Unlike these pixel-space approaches that focus on learning a generator, we study both latent-space inference (via tokenization) and generation. Concurrent works: Several concurrent papers have ex- plored closely related directions toward unifying tokeniza- tion and latent generative modeling. The closest to our set- ting is Googleâs Unified Latents (Heek et al., 2026), which studies end-to-end training of a tokenizer together with a latent generator and is closely aligned with our separate- weights ablation (i.e., an encoder and denoiser trained jointly without parameter sharing). In contrast to our single-stage results, their strongest numbers rely on an additional second- stage diffusion fine-tuning step (see Appendix B of Heek et al. 2026). Latent Forcing (Baade et al., 2026) extends pixel-space diffusion (JiT) by denoising pretrained DINO 3 End-to-End Training for Unified Tokenization and Latent Denoising latents alongside image patches through a shared bottleneck, but does not learn the latent space from scratch. Another concurrent effort (Chefer et al., 2026) adds an auxiliary self- supervised objective alongside the diffusion/flow objective, but similarly operates in a latent space defined by a pre- trained encoder rather than jointly learning the tokenizer and generator end-to-end. In contrast to these works, our primary emphasis is on understanding the capabilities of a single-stage, end-to-end trained latent diffusion model. To this end, we study a perspective in which encoding & denoising are performed by the same network parameters. 3. Unifying Tokenization & Latent Denoising Can we jointly train a tokenizer and a generator end-to-end in a single stage, such that gradients from one objective meaningfully shape the other? The answer is yes: we show that single-stage end-to-end training can learn a latent space that supports both high-fidelity reconstruction and iterative generation. While recent work has begun to explore end- to-end training, most approaches still rely on multi-stage pipelines (e.g., pretraining or freezing parts of the system) or introduce external supervision from pretrained represen- tation models. These design choices can be effective, but they make it harder to isolate and study the intrinsic inter- action between tokenization and generation. In this work, we take a step toward single-stage joint tokenization and generation without external supervision, using a single uni- fied network trained simultaneously with reconstruction and latent denoising objectives. In many ways, an early and elegant solution to this already exists: variational autoencoders (VAEs) jointly learn an encoderâdecoder for reconstruction while also imposing a simple latent prior, typicallyN(0,I), that enables sampling and generation. This classical design suggests that tokeniza- tion & generation need not be separated into distinct stages. 3.1. From VAE to UNITE In a VAE, the âtokenizerâ is the encoder,E θ : it maps a data pointxto a conditional latent distributionq(z | x) rather than a single code. The decoder,D Ď , reconstructs by samplingz âź q(z | x)and mapping back to data space via p(x | z). For any generative model, a central requirement is to map an easy-to-sample distribution into an expressive latent space that supports high-quality decoding. VAEs meet this requirement by regularizing the encoder so that its latent distribution remains close to a simple priorp(z) =N(0,I) (through the KL loss term), making generation as simple as sampling z âź p(z) and decoding. VAE: z = E θ (x); Ëx = D Ď (z); Notably, VAE-family encoderâdecoder tokenizers have be- come a standard building block in modern vision and video Generative Encoder (as tokenizer) Decoder K Registers ~ N(0,1) Generative Encoder (as denoiser) Tokenizer â Generator Image patches ! 0 "! 0 +noise # Recon. Pathway Gen. Pathway Gradients Figure 3. UNITE Training Pipeline uses two forward passes through the Generative Encoder: first, mapping (distilling) image patches into latent registers, and second, denoising a noised version of those latents, with weights shared across both passes. Training combines reconstruction losses with a denoising loss| Ě Ëz 0 âsg( Ěz 0 )|. foundation-model pipelines: high-dimensional visual in- puts are first compressed into latents via a VAE/VQ-style encoder, but generation is performed in latent space by a separate model. In this regime, these autoencoders function primarily as tokenizers rather than as the final generative model, since a simple Gaussian prior typically does not reach the sample fidelity of modern diffusion generators. Modern high-fidelity generative models therefore replace VAE-style Gaussian prior sampling with a learned iterative generative process, while retaining the VAEâs role as the tokenizer. In latent diffusion and flow models, a VAE-style encoder first maps data into a compact latent space, and a separate denoising model,G Ď , is trained to transform Gaussian noise into samples from the latent data distribution via iterative denoising. In practice, this is often implemented as a staged pipeline: the tokenizer is trained and frozen; the denoiser is trained on top of the fixed latent space. LDM: z = E θ (x); Ëx = D Ď (z); Ëz = G Ď (z t ,t); UNITE replaces the separate tokenizer and latent denoiser with a shared set of parameters, the Generative Encoder, as demonstrated in Fig. 2. This shared module retains the simplicity of the autoencoder interfaceâan encoder and a decoderâwhile enabling single-stage learning of both tokenization and generation. Paired with a decoderD Ď that maps latents back to image space, the Generative En- coderGE θ operates in two modes. In tokenization mode, GE θ maps an imagexto latent tokensz = GE θ (x)op- timized for reconstruction, without enforcing an explicit KL-to-Gaussian bottleneck. In generation mode, the same GE θ is used as a latent denoiser: given a noisy latentz t and noise levelt, it predicts the corresponding denoising target, enabling iterative sampling from Gaussian noise at inference time. Sharing parameters across these two modes lets gradients from both objectives jointly shape the same weights in a single training job. This yields a minimal end-to-end pipeline with performance approaching modern latent generative models, with the resulting formulation as: UNITE: z = GE θ (x); Ëx = D Ď (z); Ëz = GE θ (z t ,t); 4 End-to-End Training for Unified Tokenization and Latent Denoising 3.2. End-to-End Training for UNITE Training Pipeline: We adopt a Vision Transformer (ViT) (Dosovitskiy et al., 2021) backbone for both the generative encoder and the decoder, motivated by the strong empirical performance of Transformer architectures in diffusion/flow denoising. The generative encoderGE θ must support two operating modes with compatible input/output types: a tok- enization pathway, which ingests image patch tokens and produces a compact latent representation, and a generation (denoising) pathway, which ingests noisy latents along a flow or diffusion trajectory that connects the latent distribu- tion to a standard normal prior. To unify the input format across pathways, we represent the latentzas a fixed set ofKregister tokens. In the tok- enization pathway, we concatenate the image patch tokens withKregisters, initializing the registers as i.i.d. Gaussian noise,N(0,I), to match the input distribution at the maxi- mum noise level. The concatenated sequence is processed with self-attention in a first forward pass throughGE θ . We then discard the patch tokens and retain only the updated registers. These updated registers serve as the image latents z 0 , having absorbed the relevant information from patches through attention. The decoderD Ď consumesz 0 and re- constructs the image using a ViT-style stack followed by a lightweight unpatchification head to produce pixels. In the generation (denoising) pathway, we first corrupt the clean latentsz 0 to obtain a noisy latentz t at noise levelt (using our rectified-flow / flow-matching corruption process) and then usez t to initialize the sameKregisters. No image patches are concatenated in this pathway. A second forward pass throughGE θ (now in generation mode, conditioned ontand optional class information) predicts the denoising target; in our implementation we usex-start prediction, i.e., Ëz 0 = GE θ (z t ,t), so that the denoiser output lies in the same space as the tokenization output. To avoid degenerate solutions where the denoiser objective collapses the latent space, we stop gradients through the clean latents used to form z t (i.e., we detach z 0 before noising). The final layer ofGE θ is a normalization module. Em- pirically, we find that LayerNorm (Ba et al., 2016) with learnable scale and shift parameters performs best. As a result, the clean latentsz 0 (from the tokenization pathway), the denoised predictionsËz 0 (from the generation pathway), and the model outputs at each denoising step during infer- ence are all normalized. Overall, each training iteration performs two forward passes through the sharedGE θ : an image-conditioned pass to pro- duce clean latents for reconstruction, followed by a latent- only pass to denoise a corrupted version of those latents. The full system is trained end-to-end in a single stage by jointly optimizing a pixel-space reconstruction objective F low loss Recon loss Reconstruction Noising Scale 0.0 (max noise) 0.7 (default) 0.81.0 (no noise) Epochs Figure 4. UNITEâs Training dynamics: The conflicting nature of the reconstruction and denoising objectives leads to an adversarial training behavior when trained jointly. The dotted lines (zoomed in) represent different ablations (see Appendix) based on the scale of noise added in reconstruction pathway for decoder robustness. (via D Ď ) and a latent-space denoising objective (via GE θ ). Training Objectives: We optimize two losses computed from the two forward passes described above. For recon- struction, we encode the image into clean latentsz 0 = GE θ (x), inject small Gaussian noise Ěz 0 = z 0 + ĎÎľwith reconstruction noise scaleĎ = 0.7following Leng et al. (2025); Yu et al. (2025a), and decodeËx = D Ď ( Ěz 0 ). The re- construction loss combines pixel-level and perceptual terms: L recon = âĽËx â x⼠1 + LPIPS(Ëx,x). For generation, we apply rectified flow matching (Liu et al., 2023) on the la- tents. Given clean latentsz 0 , we construct noisy latents z t = tz 0 +(1ât)ÎľwithÎľâźN(0,I)andtâźU[0, 1](where t=1corresponds to clean data andt=0to pure noise), then train the generative encoder to predict clean latents viaËz 0 = GE θ (z t ,t) . We minimizeL flow =E t,Îľ [âĽËz 0 â sg(z 0 )⼠2 2 ], wheresg(¡)denotes stop-gradient to prevent degenerate so- lutions. The total objective is the sum of reconstruction and generation losses. Inference: At inference, the Generative Encoder can serve as the tokenizer by mapping an input image to its latent representation in a single forward pass. For generation, we start from a class label and noisy latent registers, and iteratively refine them through multiple passes of the GE into clean, decodable latents (shown as red loops in Fig. 1). 3.3. Understanding UNITEâs Training Dynamics The adversarial nature of joint training.Jointly training tokenization and generation under weight sharing induces non-trivial dynamics. In a standard LDM, the latent space is produced by a pretrained (and typically frozen) tokenizer, so the generative objective does not shape the latent interface. In UNITE, reconstruction and generative objectives are op- timized jointly over the same parameters, so each objective can influence the representations used by the other. This dynamic is best understood as the search for a latent space that satisfies two distinct pressures shaping its struc- ture. The reconstruction objective drives the encoder to 5 End-to-End Training for Unified Tokenization and Latent Denoising maximize information content, preventing the latent rep- resentation from becoming too coarse to capture instance- specific detail. Simultaneously, the generative objective constrains how this information is encoded: it penalizes learning fragile representations whose semantic content can be easily destroyed by noise, since such instability makes denoising harder. Consequently, joint optimization balances these pressures, finding a latent space that is rich enough for reconstruction yet robust enough against perturbations. By forcing the encoder to adopt this robust geometry, the gen- erative loss effectively molds the latent space into one that is intrinsically easier to denoiseâfacilitating high-fidelity generation. Empirically, this interaction can resemble an âadversarialâ game: the two losses do not necessarily decrease mono- tonically together. Improvements in generative fidelity can even coincide with an increase in denoising loss, as shown in Fig. 4 (see red curves with star markers). Crucially, a rising denoising loss does not imply worse generation. In- stead, it often signals that the latent space is becoming richer and more informative to satisfy the reconstruction objective, making the denoising task harder but the resulting samples more realistic. During training, we often observe genera- tion metrics (e.g., FID/IS) improving even as the denoising loss increases, until the system reaches a stable equilibrium. Similar to GAN-style training, the goal is therefore not to drive all losses to zero, but to reach stable training dynamics where the latent space balances information density with generative robustness. This perspective is also consistent with modern diffusion/flow models, where the denoising loss typically stabilizes at a non-zero value. 4. Analyzing UNITEâs Generative Encoder In UNITE, we pursue end-to-end training by sharing pa- rameters between the encoder and denoiser roles of a single network. This choice suggests a natural hypothesis: pa- rameter tying encourages the model to develop a common latent âlanguageââshared internal features and transforma- tions that simultaneously support reconstruction and itera- tive denoising-based sampling. To better understand this design choice, we study two alter- native routes to end-to-end latent diffusion training that each relax a component of our Generative Encoder mechanism. First, we remove parameter tying, maintaining separate en- coder and denoiser networks while still training both objec- tives jointly. Second, we remove the stop-gradient through clean latents, allowing denoising gradients to backpropagate into the tokenization pathway. Together, these alternatives help isolate the role of weight sharing and gradient flow in our end-to-end formulation. Finally, we also study these end-to-end training approaches through the lenses of repre- sentation alignment and compression. FrĂŠchet inception distance (FID) Solid Lines (gFID) Dashed Lines (rFID) # Denoising Iters (per Recon. Iter) Recon. FID (rFID) Gen. FID (gFID) Figure 5. Weight-shared vs. Separate Enc-Denoiser training. UNITE uses a single Generative Encoder, sharing weights between tokenization and generation. To isolate the effect of weight shar- ing, we keep the rest of the end-to-end training pipeline fixed, including the stop-gradient that prevents denoising gradients from flowing into the tokenized output. Both UNITE and the separate encoder-denoiser ablation attain competitive performance, with UNITE benefiting from more denoising-to-reconstruction steps ra- tio during training, achieving the best overall rFIDâgFID trade-off. Weight-Shared vs. Separate EncoderâDenoiser Training: Our Generative Encoder ties the encoding and denoising roles by sharing parameters. As an ablation, we keep the en- tire end-to-end pipeline fixedâincluding the stop-gradient that prevents denoising gradients from flowing through the tokenization output into the encoderâbut instantiate separate networks for the encoder and the denoiser. In this separate- networks ablation, the encoder & denoiser are optimized for their own objectives, with no gradient interaction involved. If the weight-shared Generative Encoder matches or im- proves upon this separate-weights variant, it already offers a practical advantage: fewer parameters to store and up- date, and a shorter description length (MDL) for the learned model. Fig. 5 shows that, while the separate-weights abla- tion is competitive, parameter tying yields the best overall reconstructionâgeneration trade-off. Specifically, we report rFID and gFID as a function of the number of denoising (flow) steps performed per reconstruction step during train- ing. Under weight sharing, increasing the number of flow steps consistently improves generation fidelity, reducing gFID from 3.33 to 2.12 as the number of flow iterations is increased by14Ă. This indicates that the latent space becomes more sampleableâwhile maintaining, or slightly improving, reconstruction fidelity, suggesting that the repre- sentation also remains information-preserving at the chosen compression dimension. Next, we study the role of the stop gradients operator between the denoiser and the tokenizer. 6 End-to-End Training for Unified Tokenization and Latent Denoising Centred Kernel Alignment (CKA) Layer ID Cosine Similarity (between final outputs) Denoising timestampLayer ID Weight Sharing Impact on To k e n i z a t i o n-Generation Alignment (Denoiser Ă latent Ă Encoder) Gradient Backprop Impact on To k e n i z a t i o n-Generation Alignment Centred Kernel Alignment (CKA) Figure 6. Representation alignment between tokenization and generation pathways. We measure alignment between tokenization and denoising activations using CKA and cosine similarity. Given an input image, we first record intermediate activations along the tokenization pathway, then corrupt the encoded latent and record the corresponding denoising-pathway activations. Left: both the weight-shared UNITE model and the separate encoderâdenoiser ablation exhibit strong alignment, especially in later layers, indicating that tokenization and denoising are intrinsically aligned tasks. Middle: removing the stop-gradient and backpropagating denoising gradients through the latent weakens late-layer alignment, even though the denoising objective still matches the final latent target. Right: cosine similarity on the final latents decreases at lower denoising timesteps in the no-stop-gradient setting, suggesting that direct gradient backpropagation from denoising into tokenization leads to a less cleanly shared representation (see Fig. 7 for visual interpretation). Backpropagating Denoising Gradients through the En- coder:Throughout this work, we stop denoising gradients from flowing through the clean latent into the tokeniza- tion pathway. Concretely, after the tokenization pass pro- ducesz 0 = GE θ (x), we applysg(¡)before constructing the noised latentz t used in the denoising pass. As a result, the flow-matching objective updates GE θ only through the second (denoising) forward pass, rather than also directly shaping tokenization through gradients flowing into z 0 . Importantly, this does not decouple tokenization and gen- eration: in the weight-shared Generative Encoder, recon- struction and denoising still act on the same set of network parameters, so both objectives jointly shape the learned representation. The stop-gradient only removes the more direct route in which denoising gradients also flow through the clean latent itself. In the separate encoderâdenoiser setting, removing this stop-gradient yields a two-network end-to-end regime closely analogous to concurrent work on Unified Latents (UL) (Heek et al., 2026), which jointly trains separate encoder and denoiser modules without pa- rameter sharing. We therefore study what happens when denoising gradients are allowed to backpropagate through the clean latent (termed the no-stop-grad setting in the fol- lowing paragraphs), both in our weight-shared GE setting and in the separate encoderâdenoiser ablation. Looking at rFID/gFID, removing the stop-gradient improves the separate encoderâdenoiser ablation from2.60/1.30to 2.24/0.85(gFID/rFID), indicating that end-to-end joint training of tokenization and generation is promising. As noted in the concurrent Unified Latents (Heek et al., 2026) (their Appendix B), obtaining the best performance in the no-stop-gradient setting requires tuning the denoising-to- reconstruction loss ratio. By contrast, for UNITE, we obtain the best performance (gFID= 2.12, rFID= 1.1) with stop- gradient in place. One possible hypothesis is that, under weight sharing, the two objectives already interact through a common parameter set, so allowing denoising gradients to additionally flow through the clean latent introduces extra (asymmetric) gradient interference. In this sense, weight sharing itself acts as a natural coupling mechanism between the two tasks: simply increasing the number of flow iter- ations improves performance, without requiring as much loss-weight tuning. We now showcase representation align- ment and compression-based analysis. Tokenization-Generation Representation Alignment Analysis:As shown in Fig. 2, tokenization can be viewed as a generative process under strong observability,p θ (z | x), whereas generation corresponds to unconditional sampling from the induced prior,z âź p θ (z). This viewpoint suggests that the two tasks may be aligned, and motivates measuring representational alignment between the two modes. We test this by measuring alignment between tokenization-pathway & denoising-pathway activations using Centered Kernel Alignment (CKA / CKNNA) and Cosine Similarity (Fig. 6). Several aspects of our design encourage alignment. First, both modes are trained to operate in the same latent space: the denoiser is supervised to predict the corresponding clean latent for a corrupted version of the encoded latent. Second, the GE receives the same latent register parameterization in both modes; during tokenization these registers are ini- tialized fromN(0, 1), reducing input-domain mismatch between tokenization and generation. Finally, we adopt ar- chitectural and optimization choices that limit drift between modes: (i) consistent normalization throughout the network 7 End-to-End Training for Unified Tokenization and Latent Denoising UNITE Sep. Enc Denoiser UNITE (no - stop - grad) Sep Enc - Den (no - stop - grad) T=0.2T=0.6T=0.8ReconstructionT=0.2T=0.6T=0.8Reconstruction Figure 7. Analyzing the denoising trajectory. Given an input image, we first encode it into latents, corrupt the latent with noise, and then decode the denoised prediction at different noise levels (first three columns). The final column shows direct decoding of the clean latent. Although all four models achieve competitive aggregate rFID/gFID, the stop-gradient variants (first two rows)âUNITE and the separate encoder-denoiser ablationâexhibit markedly cleaner intermediate denoising trajectories, with higher PSNR to the input image across all noise levels. This result is consistent with the representation-alignment in Fig. 6, which shows drop in alignment at final layers. (within blocks and at the encoder output), (i) matched condi- tioning interfaces across modes (e.g., time and class signals injected in analogous ways) to avoid mode-specific short- cuts, (i) conservative optimization (learning-rate warmup and schedules) to prevent one objective from dominating shared parameters early in training. With these choices in place, we find that both the weight- shared Generative Encoder and the separate encoder- denoiser variant exhibit high CKA/CKNNA alignment (Fig. 6, left), indicating that tokenization and denoising are intrinsically aligned tasks in our setting. This also clari- fies the role of weight sharing: when the two tasks already align, parameter tying becomes a principled way to remove redundancyâespecially in the reusable functional sublayers (attention and MLPs)âwhile retaining strong reconstruc- tion and generation fidelity. When analyzing the no-stop-gradient alternative (Fig. 6, middle and right), we observe that, for both the weight- shared and separate encoderâdenoiser settings, CKA and cosine-similarity alignment between the outputs of the tok- enization and denoising pathways is reduced, relative to the stop-gradient variants, despite the denoising objective en- couraging agreement at the final latent target. Further, Fig. 7 shows that the no-stop-gradient models produce noticeably noisier intermediate denoised reconstructions. Taken to- gether, these observations suggest that stopping denoising gradients through clean latent may help preserve a more cleanly shared representation between tok. & generation. Entropy / Compression Analysis. We next study the encoderâdenoiser relationship through the lens of com- pressibility, motivated by a Minimum Description Length (MDL) perspective: if tokenization and denoising imple- ment closely related computations, then a unified latent- generation program might admit a shorter description than two independently parameterized modules. Concretely, we estimate an empirical description-length proxy for model weights using per-tensor histogram entropy. We begin with the separate encoderâdenoiser setting. Com- pared to random weights, the total entropy of the en- coder drops from179.2MB at random initialization to 121.9MB after training, with both normalization parameters (60.0â 30.7MB) and functional attention/MLP parameters (119.2â 91.2MB) becoming substantially more structured as a result of training. In the weight-shared Generative Encoder setting, the en- tropy of the functional attention/MLP parameters remains nearly unchanged relative to the separate encoder (91.2MB â 90.8MB), while the main increase is concentrated in normalization-related parameters, whose entropy rises mod- estly from30.7MB to42.0MB and closely matches that of the separate denoiser (42.0MB). Thus, unifying tokenization and denoising does not require a more complex functional 8 End-to-End Training for Unified Tokenization and Latent Denoising Table 1. ImageNet 256Ă256 generation. Our approach out- performs both recent single-stage pixel baselines and standard two-stage latent diffusion frameworks by a large margin. Method Aux. Token ParamsFIDâISâ Single-stage Frameworks JiT-B/16-131M3.66275.1 UNITE-B (Ours)Joint217M2.12294.1 RIN-410M3.42182.0 JiT-L/16-459M2.36298.5 ADM-G-554M4.59186.7 UNITE-L (Ours)Joint589M1.73296.0 PixelFlow-XL/4-677M1.98282.1 PixNerd-XL/16-700M2.15297 UNITE-XL (Ours)Joint806M1.75309.9 JiT-H/16-953M1.86303.4 SiD-2B2.44256.3 VDM++-2B2.12267.7 JiT-G/16-2B1.82292.6 Two-stage Frameworks DiT-XL/2SD-VAE675M+49M2.27278.2 SiT-XL/2SD-VAE675M+49M2.06277.5 Two-stage Frameworks with Aux Supervision (DINOv2) REPA-BSD-VAE130M+49M2.15268.3 RAE-BRAE-tok130M+415M2.08275.1 REPA-SiT-XL/2SD-VAE675M+49M1.42305.7 LightningDiT-XL/2VA-VAE675M+49M1.35295.3 DDT-XL/2SD-VAE675M+49M1.26310.6 RAE-DiT DH -XL/2RAE839M+415M1.13262.6 Concurrent works LF-DiT-LDINOv2465M2.48â backbone; instead, the shared model reuses essentially the same attention/MLP computation and expresses the residual mode-specific adaptation primarily through normalization and scale parameters. This provides a complementary MDL- style interpretation of why sharing works in our setting: pa- rameter tying may yield a shorter description of the joint latent-generation program, not by substantially altering the main reusable computation, but by preserving a common functional backbone while allocating only a small additional entropy budget to normalization. This interpretation aligns with our CKA analysis, since pathway alignment remains high while CKA is largely insensitive to norm/scale changes, suggesting that tokenization & denoising differ more in fea- ture calibration than in core representational geometry. 5. Experimental Results Can a single training job produce both a strong tokenizer and a strong generator? In this section, we show that UNITE achieves nearâstate-of-the-art performance on both recon- struction and generation tasks across image and molecule modalities. See Appendix for more ablations. Table 2. ImageNet 256Ă256 reconstruction. Our tokenizer achieves competitive rFID without adversarial loss (Adv.) or pre- trained encoders. All UNITE rows use base backbone at 120 eps. TokenizerAdv.Pretrained EncoderrFIDâ With adversarial / external supervision SD-VAEâ-0.62 DC-AE-f32â-0.69 RAEâDINOv20.58 VA-VAEâDINOv20.28 Without adversarial or external supervision ViTok-B/16 â --1.63 UNITE-B (Ours)--1.01 + GAN decoder ft â â-0.51 w/ separate weights--1.38 â Stage-1 only (L2+LPIPS+KL). â Decoder-only ft, 16 epochs. 5.1. ImageNet-256 Results Generation. Tab. 1 summarizes our main generation re- sults on ImageNet-256. The results indicate that truly end-to- end training of tokenization and generation is not only feasi- ble, but also competitive. In particular, UNITE-B reaches an FID of 2.12, substantially improving over the single-stage baseline JiT-B/16 (Li & He, 2025) (FID 3.66). Increasing the model capacity further improves performance: UNITE- L (encoder size: L, default patch size: 16) reduces FID to 1.73, surpassing two-stage approaches such as DiT-XL/2 (FID 2.27) and SiT-XL/2 (FID 2.06), suggesting that the unified setup continues to benefit from scale. Fig. 8 shows representative samples from UNITE-XL (more uncurated class-conditional generations in Appendix). Unlike previous latent diffusion pipelines that train VAEs with GAN-based adversarial objectives, UNITE uses no adversarial loss. Unlike RAE (Zheng et al., 2025) and REPA (Yu et al., 2025a), which fundamentally rely on pretrained vision en- coders, our single-stage approach reaches comparable per- formance while training from scratch, without requiring an external pretrained representation model. 1 This simplic- ityâone encoder that serves both tokenization & generation via weight sharingâmakes the system easier to train and de- ploy, reducing reliance on external pretrained components. Compared with concurrent works including Latent Forcing (Baade et al., 2026) (LF-DiT-L as mentioned in Tab. 1), Uni- fied Latents (Heek et al., 2026), and Self-Flow (Chefer et al., 2026), our method is trained fully from scratch while achiev- ing stronger generation FID. Unified Latents reports results only on ImageNet-512, with its best performance further relying on a second-stage diffusion fine-tuning step. Self- Flow, in contrast, builds on a pretrained DINO tokenizer and reports only unconditional generation results. 1 Our ImageNet training uses LPIPS loss, which requires a pretrained VGG. However, (a) training VGG is inexpensive; (b) our molecule gen. results do not use LPIPS. 9 End-to-End Training for Unified Tokenization and Latent Denoising Figure 8. Selected samples from UNITE-XL. Generated using 50 steps with CFG. This model achieves FID 1.75. Reconstruction.Tab. 2 compares the reconstruction qual- ity of UNITE against existing tokenizers. Most prior meth- ods decouple reconstruction and generation, and low-rFID tokenizers such as VAEs and VQGANs typically rely on adversarial objectives in addition to reconstruction losses. More recent approaches further improve reconstruction fi- delity by leveraging externally pretrained self-supervised encoders (e.g., DINOv2). As a reference point, a vanilla ViT autoencoder (ViTok-B/16 Stage 1 (Hansen-Estruch et al., 2025)), trained from scratch with only L2+LPIPS+KL losses and no adversarial training, attains an rFID of 1.63. Despite being trained jointly with a generative objective, UNITE-B (217M parameters) achieves an rFID of 1.01 after 120 epochs, already outperforming the vanilla autoen- coder baseline. A lightweight adversarial fine-tuning stage â which freezes the Generative Encoder and updates only the decoder for 16 epochsâfurther reduces rFID to 0.51, surpassing all baselines, including RAE (0.58) and SD-VAE (0.62), without relying on any self-supervised pretraining. Finally, removing weight sharing (Tab. 2 last row) substan- tially degrades reconstruction quality (rFID 1.38), further supporting the claim that shared parameterization benefits both reconstruction and generation. 5.2. Beyond Vision: Application to Domains Without Pretrained Encoders Recent approaches such as REPA (Yu et al., 2025a) and RAE (Zheng et al., 2025) crucially depend on pretrained representation models, e.g., DINOv2 (Oquab et al., 2023), to strengthen latent diffusion. This reliance makes their trans- fer to domains where such encoders are unavailableâor Table 3. QM9 molecule generation. UNITE-S achieves the best reconstruction accuracy (99.37% match) and uniqueness (99.71%) under single-stage training. Crystal generation results on MP20 are provided in Appendix C.2. ReconstructionGeneration Method Match (%) RMSD ( Ě A) Valid (%) Unique (%) EDM (Hoogeboom et al., 2022)â91.990.7 GeoLDM (Xu et al., 2023)â93.892.9 ADiT Tokenizer (Joshi et al., 2025)97.200.075â ADiT-S QM9-only (Joshi et al., 2025)â96.0297.76 UNITE-S (Ours)99.370.03994.9099.71 expensive to obtainâless straightforward, especially in set- tings with limited data or weaker pretraining ecosystems. In contrast, our end-to-end formulation does not require pretrained encoders: tokenization and generation are learned jointly from scratch in a single training run âmaking latent generative modeling applicable to domains where strong pretrained representation models do not exist. We demonstrate this capability on QM9 molecule genera- tion, a setting with no DINO-equivalent pretrained encoder. As shown in Tab. 3, UNITE achieves state-of-the-art per- formance, matching or surpassing the All-atom Diffusion Transformer (ADiT) (Joshi et al., 2025)âthe current best method that relies on a separate VAE tokenizer. Notably, we obtain a 99.37% reconstruction match rate (vs. 97.20% for ADiT) and 99.71% uniqueness among generated molecules (vs. 97.76%), while training fully end-to-end and without any pretrained components. These results further motivate studying true end-to-end train- ing of tokenization and generationâwhere the two objec- tives are optimized jointly and gradients from each task shape the same representation spaceâas a means of enhanc- ing latent diffusion models, rather than leveraging pretrained encoders trained on additional data. 5.3. Training Efficiency We report total training FLOPs measured with gradient checkpointing enabled. For UNITE-B, each training sam- ple costs approximately 3.5 TFLOPs (forward + backward), including one tokenization pass throughGE θ (512 tokens), fourteen denoising mini-batch passes (256 tokens each), one decoder pass, and one forward pass through a frozen VGG network for the LPIPS loss. Over 120 ImageNet epochs, UNITE-B requires approximately6.7Ă 10 20 FLOPs and reaches an FID of 2.18. Reducing the number of denoising iterations per reconstruction iteration can further lower train- ing cost, at the expense of a modest increase in gFID. This is approx.15Ăcheaper than the end-to-end cost of meth- ods that rely on pretrained DINOv2 encoders. RAE (Zheng et al., 2025) and LF-DiT (Baade et al., 2026) both depend on DINOv2 features, whose ViT-g/14 pretraining and dis- 10 End-to-End Training for Unified Tokenization and Latent Denoising tillation together require approx. 27,000 A100-GPU-hours, corresponding toâź 1.0Ă 10 22 model FLOPs 2 . This con- stitutes a fixed upfront cost inherited by any downstream method built on top of these features. In contrast, UNITE eliminates this overhead entirely by training from scratch. Compared with standard two-stage latent diffusion models, our total compute is comparable: UNITE-B surpasses DiT- XL/2 (Peebles & Xie, 2023) (FID 2.27) at nearly matched total FLOPs (6.7Ă 10 20 vs.6.4Ă 10 20 ), while using3Ă fewer parameters (217M vs. 724M). In addition, UNITE jointly learns a tokenizer whose latent space is shaped by both reconstruction and generation objectives (Tab. 2). Among single-stage methods, UNITE-B (6.7Ă10 20 FLOPs, 217M parameters) achieves an FID of 2.18 at total compute comparable to JiT-G/16 (Li & He, 2025) (âź 8.8Ă 10 20 FLOPs, 2B parameters, FID 1.82), while using approxi- mately10Ăfewer parameters. Moreover, UNITE produces a reusable latent tokenizer alongside the generator, a capa- bility that pixel-space methods such as JiT do not offer. 6. Conclusion We present UNITE, a unified approach to joint tokeniza- tion and generation. Our encoder, termed the Generative Encoder, serves as both tokenizer and latent denoiser, with weights shared across the two objectives. This shared param- eterization allows reconstruction and generation gradients to jointly shape the representation space, encouraging a com- mon latent âlanguageâ that supports both tasks. UNITE is trained end-to-end in a single stage, with each iteration performing two forward passes through the same Generative Encoder: one for tokenization/reconstruction and one for latent denoising. Across ImageNet and molecule genera- tion, UNITE achieves near-state-of-the-art fidelity: the base model reaches 2.12 gFID on ImageNet 256Ă256, and scal- ing to XL improves this to 1.75 gFID. We further analyze the Generative Encoder through the lenses of representation alignment and compression. More broadly, our results suggest two practical implications. First, it removes the reliance on pretrained encoders such as DINO for generative modeling, opening the door to la- tent generative modeling in domains where such encoders are unavailable. Second, our unified architecture is simpler and more efficient than conventional two-stage pipelines, reducing both implementation complexity and overall com- putational requirements. 7. Discussions The core contribution of UNITE is to align tokenization and generation by training both over a shared latent space. The 2 27,316 A100-GPU-hoursĂ312 TFLOP/s (A100 BF16 peak) Ă0.4 model-FLOP utilizationâ 1.0Ă 10 22 (Oquab et al., 2023). two objectives we consider are denoising and reconstruc- tion. While reconstruction is a natural objective for learning compressed representations that preserve input information, exploring alternative objectives for tokenization beyond re- construction is an interesting direction for researchâfor ex- ample, jointly training the Generative Encoder with DINO- or JEPA-style objectives. This is especially appealing for robotics, where generative modeling can provide a useful world model of the environment. However, naively training such a world model on standard VAE latents may not yield actionable latents that matter most for decision-making. Another point worth discussing is the vision-language mod- eling capability of the Generative Encoder. The Generative Encoder idea is loosely reminiscent of the classical wake- sleep algorithm, whose broader goal was to bridge discrim- inative and generative modeling. While UNITE achieves strong reconstruction and generation fidelity, the linear prob- ing accuracy of the Generative Encoder remains comparable to that of other generative tokenizers, such as VAEs and VQGANs, at around30%. We believe that linear probing (LP) alone may not be fully predictive of the discriminative strengths of highly compressed latent representations. In particular, stronger compression may require greater down- stream decoding capacity before the representation becomes predictive for a given task. For this reason, evaluating the tokenizer in a VLM setting may provide a more informative picture of its discriminative capabilities than LP alone. Furthermore, the results in Fig. 7 suggest that weight shar- ing and end-to-end joint training of tokenization and gener- ation may support further progress toward faster generative models, potentially enabling high-quality one- to few-step generation. This also raises another interesting related ques- tion: can the process of mapping images to latents itself benefit from multiple iterative refinement loops? Prior work, such as ALIT (Duggal et al., 2025), has explored this di- rection and reported improvements in linear probing and token-level object binding with additional iterations. Acknowledgements We are grateful to Jyo Pari, Shamit Lal, Tianyuan Zhang, Suwan Kim, Peter Holderrieth & Qianwei Jia for fruitful dis- cussions and constructive suggestions. We also thank Prof. Kaiming He for inspiring discussions on earlier iterations of this project. This work is in part supported by MIT-IBM Watson AI Lab; ONR MURI grant #033697-00007; the National Science Foundation under Cooperative Agreement PHY-2019786 (The NSF AI Institute for Artificial Intelli- gence and Fundamental Interactions, http://iaifi.org/). S.D. is further supported by Amazon AI Research Innovation Fel- lowship; X.B. is supported by MongoDB PhD fellowship. 11 End-to-End Training for Unified Tokenization and Latent Denoising References Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016. Baade, A., Chan, E. R., Sargent, K., Chen, C., Johnson, J., Adeli, E., and Fei-Fei, L. Latent forcing: Reordering the diffusion trajectory for pixel-space image generation, 2026. URLhttps://arxiv.org/abs/2602.1 1401. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. Advances in Neural Infor- mation Processing Systems, 33:1877â1901, 2020. URL https://papers.nips.c/paper/2020/ha sh/1457c0d6bfcb4967418bfb8ac142f64a-A bstract.html. Caron, M., Touvron, H., Misra, I., J Ě egou, H., Mairal, J., Bojanowski, P., and Joulin, A. Emerging properties in self-supervised vision transformers. Proceedings of the IEEE/CVF International Conference on Computer Vision, p. 9650â9660, 2021. Chefer, H., Esser, P., Lorenz, D., Podell, D., Raja, V., Tong, V., Torralba, A., and Rombach, R. Self-supervised flow matching for scalable multi-modal synthesis. arXiv preprint arXiv:2603.06507, 2026. Chen, J., Yu, J., Ge, C., Yao, L., Xie, E., Wu, Y., Wang, Z., Kwok, J., Luo, P., Lu, H., and Li, Z. Pixart-Îą: Fast training of diffusion transformer for photorealistic text-to- image synthesis. arXiv preprint arXiv:2310.00426, 2023. URL https://arxiv.org/abs/2310.00426. Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018. URLhttps://papers.nips.c/paper /2018/hash/69386f6b1dfed68692a24c86 86939b9-Abstract.html. Chen, S., Ge, C., Zhang, S., Sun, P., and Luo, P. Pixelflow: Pixel-space generative models with flow. arXiv preprint arXiv:2504.07963, 2025a. Chen, T. On the importance of noise scheduling for diffusion models. arXiv preprint arXiv:2301.10972, 2023. Chen, Z., Zhu, J., Chen, X., Zhang, J., Hu, X., Zhao, H., Wang, C., Yang, J., and Tai, Y. Dip: Taming diffusion models in pixel space. arXiv preprint arXiv:2511.18822, 2025b. Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URLhttps://op enreview.net/forum?id=YicbFdNTTy. Duggal, S., Isola, P., Torralba, A., and Freeman, W. T. Adap- tive length image tokenization via recurrent allocation. In The Thirteenth International Conference on Learning Representations, 2025. URLhttps://openreview .net/forum?id=mb2ryuZ3wz. Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 12873â12883, 2021. Esser, P., Kulal, S., Blattmann, A., Entezari, R., M Ě uller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. In Proceedings of the 41st Interna- tional Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, p. 12606â 12633. PMLR, 2024. URLhttps://proceedings. mlr.press/v235/esser24a.html. Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial networks. Communications of the ACM, 63(11):139â144, 2020. Hansen-Estruch, P., Yan, D., Chung, C.-Y., Zohar, O., Wang, J., Vishwanath, S., Vajda, P., and Chen, X. Learnings from scaling visual tokenizers for reconstruction and gen- eration. arXiv preprint arXiv:2501.09755, 2025. He, K., Chen, X., Xie, S., Li, Y., Doll Ě ar, P., and Girshick, R. Masked autoencoders are scalable vision learners. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, p. 16000â16009, 2022. Heek, J., Hoogeboom, E., Mensink, T., and Salimans, T. Unified latents (ul): How to train your latents, 2026. URL https://arxiv.org/abs/2602.17270. Ho, J., Jain, A., and Abbeel, P. Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, volume 33, p. 6840â6851. Curran Associates, Inc., 2020. Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning, p. 8867â8887. PMLR, 2022. 12 End-to-End Training for Unified Tokenization and Latent Denoising Hoogeboom, E., Heek, J., and Salimans, T. simple diffu- sion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, p. 13213â13232. PMLR, 2023. Hoogeboom, E., Mensink, T., Heek, J., Lamerigts, K., Gao, R., and Salimans, T. Simpler diffusion (sid2): 1.5 fid on imagenet512 with pixel-space diffusion. arXiv preprint arXiv:2410.19324, 2024. Jain, A., Ong, S. P., Hautier, G., Chen, W., Richards, W. D., Dacek, S., Cholia, S., Gunter, D., Skinner, D., Ceder, G., and Persson, K. A. Commentary: The materials project: A materials genome approach to accelerating materials innovation. APL Materials, 1(1):011002, 2013. Joshi, C. K., Fu, X., Liao, Y.-L., Gharakhanyan, V., Miller, B. K., Sriram, A., and Ulissi, Z. W. All-atom diffusion transformers: Unified generative modelling of molecules and materials. arXiv preprint arXiv:2503.03965, 2025. Kingma, D. and Gao, R. Understanding diffusion objectives as the elbo with simple data augmentation. Advances in Neural Information Processing Systems, 36:65484â 65516, 2023. Kingma, D. P. and Welling, M. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2014. Kornblith, S., Norouzi, M., Lee, H., and Hinton, G. Sim- ilarity of neural network representations revisited. In Proceedings of the 36th International Conference on Ma- chine Learning, volume 97 of Proceedings of Machine Learning Research, p. 3519â3529. PMLR, 2019. URL https://proceedings.mlr.press/v97/ko rnblith19a.html. Koroteev, M. V. Bert: a review of applications in natural language processing and understanding. arXiv preprint arXiv:2103.11943, 2021. Leng, X., Singh, J., Hou, Y., Xing, Z., Xie, S., and Zheng, L. Repa-e: Unlocking vae for end-to-end tuning with latent diffusion transformers. arXiv preprint arXiv:2504.10483, 2025. Li, T. and He, K. Back to basics: Let denoising generative models denoise. arXiv preprint arXiv:2511.13720, 2025. Li, T., Li, H., and Deng, M. Autoregressive image gener- ation without vector quantization. Advances in Neural Information Processing Systems, 37, 2024. Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In International Conference on Learning Representations, 2023. URLhttps://arxiv.org/abs/2210.0 2747. Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint arXiv:2209.03003, 2023. Ma, N., Goldstein, M., Albergo, M. S., Boffi, N. M., Vanden- Eijnden, E., and Xie, S. Sit: Exploring flow and diffusion- based generative models with scalable interpolant trans- formers. arXiv preprint arXiv:2401.08740, 2024. URL https://arxiv.org/abs/2401.08740. ECCV 2024. Ong, S. P., Richards, W. D., Jain, A., Hautier, G., Kocher, M., Cholia, S., Gunter, D., Chevrier, V. L., Persson, K. A., and Ceder, G. Python materials genomics (pymatgen): A robust, open-source python library for materials analysis. Computational Materials Science, 68:314â319, 2013. Oquab, M., Darcet, T., Moutakanni, T., Vo, H., Szafraniec, M., Khalidov, V., Fernandez, P., Haziza, D., Massa, F., El-Nouby, A., Assran, M., Ballas, N., Galuba, W., Howes, R., Huang, P.-Y., Li, S.-W., Misra, I., Rabbat, M., Sharma, V., Synnaeve, G., Xu, H., J Ě egou, H., Mairal, J., Labatut, P., Joulin, A., and Bojanowski, P. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023. URLhttps://arxiv.or g/abs/2304.07193. Peebles, W. and Xie, S. Scalable diffusion models with trans- formers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), p. 4195â4205, 2023. URLhttps://openaccess.thecvf.co m/content/ICCV2023/html/Peebles_Scal able_Diffusion_Models_with_Transform ers_ICCV_2023_paper.html. Polyak, A., Zohar, A., Brown, A., Tjandra, A., Sinha, A., Lee, A., Vyas, A., Shi, B., Ma, C.-Y., Chuang, C.-Y., et al. Movie gen: A cast of media foundation models. arXiv preprint arXiv:2410.13720, 2024. Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. In Proceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, p. 8748â8763. PMLR, 2021. URL https://proceedings.mlr.press/v139/r adford21a.html. Ramakrishnan, R., Dral, P. O., Rupp, M., and von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022, 2014. Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International 13 End-to-End Training for Unified Tokenization and Latent Denoising Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=PxTI G12RRHS. Tian, K., Jiang, Y., Yuan, Z., Peng, B., and Wang, L. Visual autoregressive modeling: Scalable image generation via next-scale prediction. In Advances in Neural Information Processing Systems, volume 37, 2024. URLhttps: //arxiv.org/abs/2404.02905. Van Den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017. Wan, T., Wang, A., Ai, B., Wen, B., Mao, C., Xie, C.-W., Chen, D., Yu, F., Zhao, H., Yang, J., et al. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.20314, 2025. Wang, S., Gao, Z., Zhu, C., Huang, W., and Wang, L. Pixnerd: Pixel neural field diffusion. arXiv preprint arXiv:2507.23268, 2025. Xu, M., Powers, A. S., Dror, R. O., Ermon, S., and Leskovec, J. Geometric latent diffusion models for 3d molecule gen- eration. In International Conference on Machine Learn- ing, p. 38592â38610. PMLR, 2023. Yu, S., Kwak, S., Jang, H., Jeong, J., Huang, J., Shin, J., and Xie, S. Representation alignment for generation: Training diffusion transformers is easier than you think. In ICLR, 2025a. Yu, Y., Xiong, W., Nie, W., Sheng, Y., Liu, S., and Luo, J. Pixeldit: Pixel diffusion transformers for image genera- tion. arXiv preprint arXiv:2511.20645, 2025b. Zhai, S. et al. Normalizing flows are capable generative models. arXiv preprint arXiv:2412.06329, 2024. Zheng, B., Ma, N., Tong, S., and Xie, S. Diffusion trans- formers with representation autoencoders. arXiv preprint arXiv:2510.11690, 2025. 14 End-to-End Training for Unified Tokenization and Latent Denoising Appendix In this appendix, we first provide more details on the reconstruction fidelity results in Sec. A. Next, we share evaluation details (Sec. B), along with additional uncurated samples generated by our model UNITE-XL shown in Fig. 9. We also provide architectural and training details in Tab. 7 and Tab. 8. Finally, Sec. C presents additional results on the molecule generation task and ablations on ImageNet. A. Reconstruction Fidelity Details Table 2 in the main paper summarizes our reconstruction results. Here, we provide additional details on the adversarial fine-tuning procedure that reduces rFID from 1.01 to 0.51, without changing gFID. GAN Decoder Fine-Tuning. After UNITE joint training converges, we optionally apply a lightweight adversarial fine- tuning stage that targets only the decoder. Concretely, we freeze the Generative Encoder entirely and train the decoder with an additional GAN loss for 16 epochs. The discriminator is initialized from our Generative Encoder, which already encodes rich semantic features from joint training; this eliminates the need for an external pretrained network (e.g., DINOv2) as the discriminator backbone. Due to limited compute budget, we did not explore fine-tuning both the encoder and decoder jointly with adversarial training. Effect of Weight Sharing on Reconstruction. As shown in Table 2, removing weight sharing between the encoder and denoiser (in the stop-gradient setting) degrades rFID from 1.01 to 1.38. We attribute this to the fact that in the shared-weight setting, the generation objective acts as an implicit regularizer on the encoder, encouraging latent representations that are both reconstructive and generatively useful. Separate weights remove this coupling, leading to a less structured latent space. That said, the separate encoderâdenoiser variant without stop-gradient achieved a much lower rFID, indicating that joint training of tokenization and generation is beneficial. B. Evaluation Protocol For reproducibility, we detail the full evaluation protocol used for all ImageNet-256 generation results. See also Fig. 9 for additional uncurated samples generated by UNITE-XL. FID Computation.We compute Fr Ě echet Inception Distance (FID) using thetorch-fidelitylibrary with InceptionV3 features. Reference statistics are computed on the full ImageNet-1K training set (1281167 images). All reported FID scores use 50K generated samples. Sampling Protocol. We adopt class-balanced sampling: exactly 50 images are generated per class for the 1K ImageNet classes, totaling 50K images. This follows the protocol used by VAR (Tian et al., 2024), MAR (Li et al., 2024), and RAE (Zheng et al., 2025), among others. As shown in RAE (Table 14), class-balanced sampling yields about 0.1 lower FID than the uniform random class sampling used in some prior work (e.g., DiT (Peebles & Xie, 2023), SiT (Ma et al., 2024)). We note this systematic difference when comparing absolute FID values across methods. ODE Solver and Inference Details.At inference time, we solve the probability flow ODE over the interval[0.1, 1.0]with classifier-free guidance. We sweep the CFG scaleĎfrom 1.0 to 4.0 in increments of 0.2 and report best FID for each model. For all reported FID numbers in the main paper (Tab. 1), we use the adaptive fifth-order DormandâPrince solver (dopri5, fromtorchdiffeq(Chen et al., 2018)), following the default configuration of the SiT codebase (Ma et al., 2024). For our model, we observe thatdopri5usesâź108 NFEs on average per sample (estimated from wall-clock timing), compared to exactly 100 NFEs for the fixed-step Heun solver with 50 steps. Tab. 4 compares FID under different evaluation protocols for the same UNITE-B checkpoint. Switching to a fixed- step second-order Heun solver with 50 steps (100 NFEs) yields FID withinâź0.05ofdopri5, consistent with prior observations that flow-matching models produce near-linear trajectories that are well approximated by low-order fixed-step integrators (Lipman et al., 2023; Liu et al., 2023). This small gap is also consistent with the SiT authorsâ report that the FID difference between dopri5 and fixed-step solvers is <0.1. 3 3 See https://github.com/willisma/SiT/issues/21. 15 End-to-End Training for Unified Tokenization and Latent Denoising Table 4. Effect of evaluation protocol on reported FID. All rows use the same UNITE-B checkpoint (240 epochs). âBalancedâ denotes 50 images per class; âRandomâ denotes uniformly sampled class labels. NFE = number of function evaluations. ODE SolverClass SamplingNFEFIDâISâ Heun (50 steps)Balanced1002.789287.6 dopri5 (adaptive)Balancedâź1082.735268.1 dopri5 (adaptive)Randomâź1082.885274.7 C. Additional Results C.1. QM9 Molecular Generation The QM9 dataset (Ramakrishnan et al., 2014) contains approximately 130K stable small organic molecules with up to 9 heavy atoms from the setC, N, O, F. Following Joshi et al. (2025), we represent molecules with explicit hydrogen atoms and use 3D Cartesian coordinates for both training and generation. Each molecule is preprocessed to ensure correct bond valencies and stable conformations, with coordinates normalized to have zero center of mass. Our training configuration employs the UNITE-S architecture with a DiT-S backbone containing approximately 33M parameters. The model is trained end-to-end for 8000 epochs with a batch size of 512 using the AdamW optimizer with a learning rate of1Ă 10 â4 . This single-stage approach contrasts with ADiTâs two-stage training, which requires 5000 epochs for the tokenizer followed by another 5000 epochs for the diffusion model, totaling 10000 epochs of training across two separate optimization phases. For evaluation, we compute four primary metrics on 10000 generated samples. The match rate measures the percentage of reconstructed molecules that exactly match the input structure after discretization. The RMSD (Root Mean Square Deviation) in Angstroms quantifies reconstruction error in atomic positions. Validity percentage indicates the proportion of generated molecules satisfying chemical constraints including proper valencies, reasonable bond lengths, and absence of steric clashes. Uniqueness measures the percentage of distinct molecules among valid generations, computed using canonical SMILES representations to identify duplicates. The UNITE-S architecture employs a weight-shared encoder-denoiser operating in a 16-dimensional latent space, signifi- cantly compressed from the original 3D coordinate space. This compression factor of approximately 20:1 (from 29 atoms Ă 3 coordinates to 16 dimensions) requires the model to learn highly efficient representations while maintaining recon. fidelity. C.2. MP20 Crystal Generation The MP20 dataset from the Materials Project (Jain et al., 2013) contains 45,231 inorganic crystal structures with up to 20 atoms per unit cell. We train UNITE-S for 10000 epochs with batch size 512, following the same single-stage approach as QM9. Evaluation follows Joshi et al. (2025), computing structural validity (pairwise distances> 0.5 Ě A, unit cell volume > 0.1 Ě A 3 ), compositional validity (charge neutrality and electronegativity balance), and match rate using pymatgenâs (Ong et al., 2013) Structure Matcher. Table 5. MP20 crystal generation results. Evaluation on 10K generated samples. MethodSizeTrainingStruct.Comp.OverallMatch ADiT Tokenizer-----84.50 ADiT MP20-onlyDiT-BTwo-stage99.690.590.1- ADiT JointDiT-BTwo-stage99.792.191.9- UNITE-S (Ours)DiT-SSingle-stage99.089.987.975.7 Table 5 shows that UNITE-S achieves 87.9% overall validity on MP20, approaching ADiTâs 90.1% despite using single- stage training. Our structural validity of 99.0% nearly matches ADiTâs 99.6%, demonstrating effective learning of crystal geometry constraints. The match rate of 75.7% is reasonable considering ADiTâs dedicated tokenizer achieves 84.50% after separate optimization. These results validate that our unified approach generalizes well from molecules to crystalsâthe same architecture that achieves a 99.37% match rate on QM9 also performs competitively on the more complex MP20 dataset without modification. 16 End-to-End Training for Unified Tokenization and Latent Denoising class 12: house finch, linnet, Carpodacus mexicanusclass 39: common iguana, iguana, Iguana iguana class 99: gooseclass 108: sea anemone, anemone class 144: pelicanclass 207: golden retriever class 309: beeclass 470: candle, taper, wax light class 725: pitcher, ewerclass 930: French loaf Figure 9. Uncurated, class-conditional samples on ImageNet 256Ă256 using UNITE-XL. We show images using CFGĎ = 4.0. Each grid contains 21 randomly sampled images, demonstrating consistent quality across diverse categories including animals, objects, and scenes. 17 End-to-End Training for Unified Tokenization and Latent Denoising Table 6. Ablation study on ImageNet-256 generation. We systematically evaluate key design choices across architecture, normalization, training dynamics, and augmentation strategies. All ablations are done using the base backbone and trained for 120 epochs. StepsFIDâr-FIDâISâ Reconstruction noiseĎ (c) 0.0 (full noise)2.871.09281.1 0.62.801.21267.1 0.72.711.01282.2 0.82.871.42292.1 1.0 (no augmentation)6.581.60275.7 Noise shiftÎą (e) 0.0 (No Shift)3.141.41267.1 0.52.711.01282.2 0.752.931.26278.0 C.3. Ablation Studies on ImageNet 256Ă256 In addition to the weight-sharing and stop-gradient ablations shown in the main paper, we provide a few additional ablations that further improve UNITEâs reconstruction and generation fidelity. Reconstruction Noise Level. Table 6 (top half) investigates the impact of noise augmentation during reconstruction training, where Gaussian noise is injected into latent representations prior to decoding. Consistent with recent findings in RAE (Zheng et al., 2025) and TARflow (Zhai et al., 2024), this acts as a useful regularizer by preventing the decoder from overfitting to noise-free latent codes, thereby improving generative capability. Notably, due to our modelâs learnable affine normalization, the system can autonomously calibrate its internal signal-to-noise ratio (SNR) to accommodate varying noise scales. As a result, the model exhibits strong robustness to the exact noise level, maintaining a nearly constant FID of around 2.7 across a range of noise levels. Noise Schedule Shifting. Following RAE (Zheng et al., 2025), we find that noise-schedule shifting is important for our 32-dimensional latent space. Table 6 (bottom half) shows that, without shifting, FID degrades to 3.14. Our best setting, a shift of 0.5, adapts the noise schedule to the compressed latent dimensionality and improves both FID and IS. This shift is equivalent to using an anchor dimension ofd anchor = 4096, matching the uncompressed token dimension before projection into the 32-dimensional latent space. Overall, this adaptation is important when working with highly compressed latents. D. Architectural Details Tab. 7 and Tab. 8 provide additional architectural and training details. For more details, refer to the codebase. Table 7. Detailed architecture configurations for UNITE models. Component Encoder/DenoiserDecoder DiT-BDiT-LDiT-XLViT-BViT-L Hidden Dimension768102411527681024 Layers1224281224 Attention Heads1216161216 MLP Ratio44444 Patch Size161616-- Latent Dimension3232323232 Latent Resolution16Ă1616Ă1616Ă1616Ă1616Ă16 Parameters (M)86.2458.2675.3130.6303.9 18 End-to-End Training for Unified Tokenization and Latent Denoising Table 8. Training configuration for ImageNet-256 experiments. All models trained with mixed precision BF16 & gradient clipping at 3.0. HyperparameterValueHyperparameterValue Base Learning Rate1Ă 10 â4 Warmup Epochs20 Global Batch Size1024Total Epochs240 OptimizerMuonLR ScheduleCosine AdamW Betas(0.9, 0.999)Min LR1Ă 10 â6 Weight Decay0Gradient Clip3.0 Reconstruction Noise (Ď)0.7EMA Decay0.9978 Flow Steps (Training)1000ODE Solver (Inference)dopri5 (adaptive,âź108 NFE) Flow Mini-batches14Noise Schedule Shift (Îą)0.5 CFG Scale (Ď)Sweep [1.0, 4.0], step 0.2Integration Interval[0.1, 1.0] 19