PaperlayerView exact revision on arXiv

Source-linked AI summary

Embedding Prediction Helps Image Generation

Sihan Xu, Ji Xie, Zilin Wang, Hui Shen, Stella X. Yu

arXiv:2610.02203v1cs.CVcs.LG

TL;DR

Diffusion-transformer conditions are embedded once and reused through denoising, motivating whether predictions from the current noisy image can adapt the condition. The paper trains NEPA to predict multiple clean-image embeddings and uses them to condition a DiT at every step; combined with REPA, NEPA-DiT-XL reaches FID 1.32 using about a third of REPA’s training compute.

  • Problem

    Can a condition predicted from the current noisy image at every denoising step replace the fixed condition reused by diffusion transformers?

    Source: §1 · p. 1

  • Method

    ECG conditions a DiT at every denoising step on clean-image embeddings that NEPA predicts from the condition and noisy image using Multi-Embedding Prediction.

    Source: §1 · p. 1, §3.2 · p. 4

  • Results

    FID 1.32 with the 250-step SDE sampler; combined with REPA, NEPA-DiT-XL uses about a third of REPA’s training compute.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Predictive embeddings can serve as an image generator’s condition in class-conditional ImageNet 256 × 256 experiments.

    Source: §1 · p. 2, §5 · p. 9

    The approach trains and runs two networks.

    Source: §5 · p. 9

Abstract

from arXiv · show

In diffusion transformers, a class label or a text prompt is embedded once, and the same condition is reused at every denoising step. We ask whether predicted embeddings can serve as this condition instead. Next-Embedding Predictive Autoregression (NEPA) trains a Transformer to predict the next continuous embedding in a sequence. In generation, the clean image follows the noisy image, so its embeddings are the next embeddings after the condition and the noisy image. We train a NEPA model to predict them all at once with Multi-Embedding Prediction, and in Embedding Conditioned Generation, a DiT generator is conditioned on these predictions, recomputed at every denoising step, so the conditioning signal adapts to the current noisy state. Experiments on class-conditional ImageNet $256\times256$ study the condition of the generator, the design of Multi-Embedding Prediction, and the scaling of both models. The NEPA model adds a second network to every sampling step; with it, and combined with REPA, our final model, NEPA-DiT-XL, reaches an FID of 1.32 using about a third of the training compute of REPA.

1 INTRODUCTION

The introduction reframes conditioning as a state-dependent prediction problem: ECG supplies predicted clean-image embeddings to a DiT at every denoising step rather than reusing a fixed label or prompt. On class-conditional ImageNet 256 × 256, NEPA-DiT-XL combined with REPA reaches FID 1.32.

1 INTRODUCTION

  • Diffusion-transformer conditions are usually embedded once and reused at every denoising step, prompting the question of whether the condition itself can be predicted from the current noisy image..Source: §1 · p. 1
  • NEPA predicts the next continuous embedding rather than reconstructing pixels or discrete tokens; in generation, the clean image follows the condition and noisy image in the sequence.This makes the clean image's predicted embeddings a candidate generator condition.Source: §1 · p. 1
  • ECG conditions a DiT on clean-image embeddings predicted by a frozen NEPA model and recomputed from the current noisy state at every denoising step.Multi-Embedding Prediction trains NEPA to predict the clean-image embeddings from the condition and noisy image at once.Source: §1 · pp. 1–2, Figure 1 · p. 1
  • FID 1.32: NEPA-DiT-XL combined with REPA reaches this result on class-conditional ImageNet 256 × 256.The introduction also describes studies of the generator condition, Multi-Embedding Prediction design, and scaling of both models.Source: §1 · p. 2

2 RELATED WORK

The related work places ECG alongside diffusion conditioning and predictive representation learning. It contrasts fixed conditions and denoiser-carried estimates with a separate predictor, and positions Multi-Embedding Prediction as an extension of next-embedding prediction.

2.1 DIFFUSION MODELS AND REPRESENTATIONS FOR GENERATION

  • Diffusion conditions stay fixed across denoising; prior work also uses visual representations as generation targets, alignment signals, or generative latents.Self-conditioning and RIN reuse estimates or latent states carried by the denoiser, whereas this paper uses a separate NEPA predictor recomputed from the noisy input.Source: §2.1 · p. 2

2.2 PREDICTIVE REPRESENTATION LEARNING

  • Predictive representation learning predicts learned embedding targets rather than reconstructing pixels or discrete tokens; NEPA predicts the next embedding, and MEP extends this to several embeddings at once.MEP's embedding-space predictions are model outputs, unlike predicted tokens that require decoding from hidden states.Source: §2.2 · p. 3

3 METHOD

The method trains NEPA to predict continuous image embeddings, extends it with MEP to predict the clean image embeddings from the condition and noisy image, then uses those predictions to condition a DiT generator at each denoising step.

3 METHOD

  • The method has two stages: train NEPA with MEP to predict clean-image embeddings, then train a DiT generator conditioned on NEPA’s predictions at each denoising step.NEPA is frozen during generator training.Source: §3 · p. 3, §3.4 · p. 5

3.1 PRELIMINARIES

  • NEPA patchifies images into continuous embeddings and uses an autoregressive Transformer to predict each following embedding, trained with an embedding-space distance.The diffusion generator operates in a pretrained autoencoder’s latent space and predicts the flow target; its condition is replaced by NEPA predictions from the condition and current noisy state.Source: §3.1 · p. 3

3.2 MULTI-EMBEDDING PREDICTION

  • MEP extends NEPA to predict the next K embeddings from the same context in one pass; NEPA is the special case K = 1.Unlike multi-token prediction, MEP predicts continuous embeddings directly rather than decoding hidden states into tokens.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 NEPA FOR GENERATION

  • For image generation, MEP predicts all clean-image patch embeddings in one pass from the condition and noisy image, pairing each noisy patch with the clean patch that follows it.The noise level is sampled uniformly so one model covers every denoising step; InfoNCE uses predicted embeddings as queries and clean embeddings as keys, with other patches serving as negatives.Source: §3.3 · pp. 4–5

3.4 EMBEDDING CONDITIONED GENERATION

  • At every denoising step, a frozen NEPA model predicts clean-image embeddings from the condition and current noisy image, and the DiT generator uses them as its condition.The generator is trained with weighted flow matching; gradients update only the generator.Source: §3.4 · p. 5

4 EXPERIMENTS

Experiments evaluate the approach on class-conditional ImageNet-1K at 256 × 256, from small-scale ablations through scaling and comparison with prior methods. The final NEPA-DiT-XL model combines Embedding Conditioned Generation with REPA and reaches an FID of 1.32.

4.1 ABLATION STUDIES

  • InfoNCE gives FID 31.36, versus 35.23 for MSE and 37.82 for cosine similarity; its patch-discriminative predictions retain detail useful for generation.Fine-tuned NEPA classification accuracy follows the opposite loss ordering.Source: §4.1 · p. 6, Table 1 · p. 6
  • Stop-gradient on targets prevents convergence; target normalization raises FID from 31.36 to 34.76, while adding the timestep changes it to 31.40.The chosen setup keeps gradients through targets, uses raw targets, and leaves the NEPA model timestep-free.Source: §4.1 · p. 6, Table 2 · p. 6, Table 3 · p. 6
  • The instance, batch, and full negative banks give FID 31.36, 32.13, and 31.21; patch size 4 gives 31.00 versus 31.36 for size 2 and 33.20 for size 8.Patch size 4 uses a quarter of the tokens of size 2, making each NEPA query correspondingly cheaper; the instance bank needs no cross-device communication.Source: §4.1 · p. 6, Table 5 · p. 7, §4.2 · p. 7

4.2 CONDITION OF THE GENERATOR

  • With the generator and recipe fixed, MEP-trained frozen NEPA predictions give FID 25.04, beating class conditioning at 36.39 and matched-budget alternatives at 29.52 and 30.86.The matched alternatives use an NEPA-XL-architecture network trained end-to-end or pretrained with flow matching and then frozen.Source: §4.2 · p. 7, Table 7 · p. 7

4.3 SCALING BEHAVIOR

  • FID decreases monotonically as either model scales: generator scaling improves FID by about 14 points and NEPA scaling by about 5 points across counterparts.All nine NEPA-B/L/XL and DiT-B/L/XL combinations are trained for 400K iterations and evaluated with 64 ODE steps without guidance.Source: §4.3 · p. 7, Table 8 · p. 7
  • At 400K iterations, NEPA-DiT-B/L/XL with the same frozen NEPA-XL reach FID 25.04, 12.86, and 11.27; larger generators lead at every checkpoint.At 200K iterations, NEPA-DiT-L and XL already reach 17.55 and 15.23, ahead of NEPA-DiT-B at 400K.Source: §4.3 · p. 7, Figure 6 · p. 7

4.4 COMPARISON WITH PREVIOUS METHODS

  • Scale 3.6 on t ∈ [0.4, 1] gives FID 1.57 and 1.32; querying every 8 steps gives 1.82 versus 1.57 every step, at 92 GFLOPs per step.With 96 ODE steps, total sampling cost is 62 TFLOPs per image, below SiT-XL/2 with REPA at 91 TFLOPs using 250 SDE steps.Source: §4.4 · pp. 8–9
  • FID 1.32 with the 250-step SDE sampler and 1.57 with the 96-step ODE sampler; NEPA-DiT-XL + REPA uses about a third of SiT-XL/2 + REPA’s training compute.The comparison is in the SD-VAE latent space; NEPA-DiT-XL uses 3.1 × 10^20 FLOPs versus 1.42 FID for SiT-XL/2 with REPA.Source: Table 9 · p. 8, §4.4 · p. 9

5 CONCLUSION

The conclusion presents predictive embeddings as a viable image-generator condition and reports NEPA-DiT-XL results on ImageNet 256 × 256. The model combines NEPA with REPA while using about a third of REPA’s training compute.

  • FID 1.32 with the 250-step SDE sampler and 1.57 with the 96-step ODE sampler are reported for NEPA-DiT-XL on ImageNet 256 × 256.Source: §5 · p. 9
  • Embedding Conditioned Generation conditions a DiT at every denoising step on clean-image embeddings predicted by a NEPA model trained with Multi-Embedding Prediction.Source: §5 · p. 9
  • NEPA-DiT-XL is trained with about a third of REPA’s compute, although both NEPA and the generator run as separate networks.Source: §5 · p. 9

A RELATION TO OTHER METHODS

The paper distinguishes NEPA from methods that condition diffusion transformers on noisy-input representations. NEPA is separately trained to predict clean-image embeddings and then frozen, rather than serving as a jointly trained generator component.

Relation to RepFusion

  • Unlike RepFusion, NEPA is trained from scratch on standard VAE latents with Multi-Embedding Prediction, without a representation autoencoder latent space or pretrained multimodal LLM.RepFusion adopts those existing components to make noisy latents compatible with pretrained language-model priors (Pan et al., 2026a).Source: A RELATION TO OTHER METHODS · p. 18

Relation to DDT

  • Unlike DDT, NEPA is separately trained to predict clean-image embeddings and then frozen, rather than jointly trained through the generator’s velocity loss and REPA alignment.DDT shares its encoder output between adjacent denoising steps to accelerate sampling; Table 14 studies the corresponding choice for NEPA.Source: A RELATION TO OTHER METHODS · p. 18

Relation to x0-prediction and self-conditioning

  • Unlike x0-prediction and self-conditioning, NEPA uses a separate frozen network trained with its own embedding-prediction objective, recomputed from the current noisy input at every step.The section also compares this approach with same-architecture networks trained end to end or pretrained using flow matching.Source: A RELATION TO OTHER METHODS · p. 18

B FUTURE WORK

The paper identifies text conditioning, other latent spaces, and higher resolutions as future directions. It also proposes that one NEPA model could serve as both a language model and a DiT conditioning model.

Text conditions

  • A NEPA model could read the prompt before the noisy image and predict clean-image embeddings from both, extending the class-label setup to text conditions.The section contrasts this possibility with text encoders that cannot adapt to the evolving noisy image and unified models that process prompt and image together.Source: B FUTURE WORK · p. 18

Unified models

  • Future work includes studying text conditions, other latent spaces, and higher resolutions.Source: B FUTURE WORK · p. 18
  • A single autoregressive NEPA model could in principle serve both as a language model with an output head and as a DiT conditioning model.Source: B FUTURE WORK · p. 18

C HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS

The implementation uses separate NEPA and generator stages, with a frozen NEPA model during generator training. The section also specifies sampling, baseline, and resource-accounting choices.

C HYPERPARAMETER AND MORE IMPLEMENTATION DETAILS

Further implementation details

ImageNet-1K accuracy

Sampling

  • System-level comparisons use a 250-step SDE Euler-Maruyama sampler or a 96-step Euler ODE sampler, with classifier-free guidance scale 3.6 only for t ∈ [0.4, 1].Source: C HYPERPARAMETER AND MORE… · p. 19

DiT baseline

  • The DiT-B baseline uses the same generator and training settings as NEPA-DiT-B but replaces predicted embeddings with a learned class-embedding token.Source: C HYPERPARAMETER AND MORE… · p. 19

Computing resources

  • NEPA and the generator with REPA together require 3.1 × 10^20 training FLOPs, versus about 1.0 × 10^21 for SiT-XL/2 with REPA.The resource accounting includes on-the-fly VAE encoding; Table 13 covers parameters and sampling FLOPs.Source: C HYPERPARAMETER AND MORE… · p. 19, Table 13 · p. 20

D ADDITIONAL ABLATIONS

The refresh-rate ablation reuses NEPA’s last predicted embeddings between queries while the generator continues denoising. Less frequent queries sharply worsen FID.

Refresh rate

  • FID rises from 1.57 to 1.82 when NEPA is queried every 8 steps, and reaches 242.63 when queried only at the first step.The comparison uses 96-step ODE sampling; the generator runs at every step and reuses the last predicted embeddings between queries.Source: D ADDITIONAL ABLATIONS · p. 20, Table 14 · p. 20
  • The ablation varies how often NEPA is queried while holding the generator’s denoising schedule at every step.Between NEPA queries, the generator reuses the last predicted embeddings.Source: Table 14 · p. 20

E ADDITIONAL QUALITATIVE RESULTS

Figures 8–10 provide additional NEPA-DiT-XL samples on ImageNet 256×256, selected from 16 samples per class.

E ADDITIONAL QUALITATIVE RESULTS

Loading 2610.02203v1…