PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusable Latents from Structure-Agnostic Distillation

Adrien Ramanana Rahary, Nicolas Dufour, Patrick Pérez, David Picard

arXiv:2609.39657v1cs.CV

TL;DR

Position-wise distillation ties autoencoder latents to a teacher’s token grid, which is unavailable for some latent shapes and cross-modal teachers. This paper tests pooled and relational objectives that do not require token-level correspondence. Pooled matching performs as well as or slightly better than position-wise alignment with image teachers and extends improvements to 1D latents and text teachers.

  • Problem

    Position-wise distillation requires a shared patch grid, which compact 1D latents and teachers from another modality may not have.

    Source: §1 · p. 1

  • Method

    The method distills pooled image-level descriptors or between-image similarity relations, leaving latent token layout unconstrained.

    Source: §1 · p. 1, §2 · p. 3

  • Results

    Pooled matching performs as well as or slightly better than position-wise alignment with image teachers and extends to 1D latents and text teachers, where Pool-Align and Soft-KL improve generation quality.

    Source: §5 · p. 4

  • Takeaways & Limitations

    Structure-agnostic distillation is a flexible alternative to position-wise alignment that leaves the latent’s layout free.

    Source: §5 · p. 4

Abstract

from arXiv · show

Distilling pretrained foundation models into an autoencoder bottleneck improves latent diffusability, enabling diffusion models to converge faster and reach higher sample quality. Standard distillation aligns the latent at each position to a co-located teacher feature, tying the latent layout to the teacher's. We show this constraint is unnecessary: aligning a single pooled image-level descriptor to the teacher's performs as well as or slightly better than dense position-wise distillation. We compare first-order and relational pooled objectives across latent shapes and teacher modalities. First-order matching extends naturally to 1D token-sequence latents and across modalities, where distilling a text encoder into an image autoencoder still improves diffusability; a relational objective based only on each image's nearest neighbours improves it as well. Code and blog post are available at https://github.com/AdrienRR/structure-agnostic-distillation and https://kyutai.org/blog/2026-09-28-structure-agnostic-distillation/.

1 Introduction

Prior distillation improves latent diffusability but ties the latent to a shared patch grid, which compact 1D latents and cross-modal teachers may lack. The paper tests pooled objectives that remove token-level correspondence and finds pooled matching can replace position-wise alignment while improving diffusability.

Motivation

  • Position-wise distillation ties latent positions to co-located teacher features, but compact 1D latents and cross-modal teachers may lack a shared patch grid.Prior work also uses frozen representation models as the diffusion space, but both approaches require a shared patch grid.Source: §1 · p. 1

Approach

  • Structure-agnostic distillation pools latent tokens into one image-level descriptor, leaving the latent’s internal layout free.The paper compares direct descriptor matching with relational objectives based on between-image similarities.Source: §1 · pp. 1–2

Finding

  • Pooled matching performs as well as or slightly better than position-wise alignment where applicable, and extends to 1D latents and text teachers.The paper reports that a text encoder can improve the diffusability of an image latent.Source: §1 · p. 2

2 Method

The method adds a distillation term to the autoencoder encoder’s objective, using pooled descriptors or between-image similarities rather than requiring token-level alignment.

Preliminaries

  • The encoder minimizes the VAE objective plus a weighted distillation term that aligns its latent to a frozen foundation encoder.A diffusion prior is trained afterwards in the autoencoder’s latent space.Source: §2 · p. 2

Pooled objectives

  • Mean-pooling produces one normalized descriptor per image and discards token layout, relaxing position-wise, descriptor-level, or relational correspondence.The pooled objectives can be used regardless of latent shape.Source: §2 · p. 2

Pool-Align

  • Pool-Align matches each image’s pooled student descriptor to its teacher descriptor using a learned linear map and cosine distance.Source: §2 · p. 3

Relational objectives

  • Relational objectives match student and teacher between-image similarity matrices rather than descriptors, so they need no projector.CKA matches the global geometry; Soft-KL matches each image’s soft distribution over nearest neighbours.Source: §2 · p. 3

3 Experiments

Across shared-grid, 1D-sequence, and cross-modal settings, the experiments compare distillation objectives under matched architecture and training budgets. Pooled and relational objectives improve generation quality across the tested settings, with image-teacher pooling matching or slightly outperforming position-wise alignment.

Experimental settings

  • The study holds architecture and training recipe fixed while varying the distillation term across a 2D grid with an image teacher, a 1D sequence with an image teacher, and a 2D grid with a text teacher.The position-wise baseline is unavailable in the 1D and cross-modal settings because they lack a shared grid.Source: §3 · p. 3

1D sequence

  • In the reported comparison, Pool-Align achieves final gFID 15.62 versus 26.49 for No-Distillation, and the objectives still beat No-Distillation in the same order.Source: §3 · p. 3

Cross-modal

  • With a text teacher, Pool-Align (6.97) and Soft-KL (8.51) improve over No-Distillation (9.52), while CKA (10.41) does not.Source: §3 · p. 3

Shared grid

  • Pool-Align reaches gFID 5.77 vs. VF 6.04 (and 1.99 vs. 2.12 with guidance), while every structure-agnostic objective beats No-Distillation at 9.52.Source: §3 · p. 3, Table 1 · p. 4

Convergence

  • In the 2D-grid and 1D-sequence settings, every distilled latent reaches lower final gFID than No-Distillation and reaches low gFID in fewer training steps.Figure 2 also marks how many fewer steps Pool-Align takes to reach No-Distillation’s final gFID.Source: Figure 2 · p. 4

4 Discussion

Pooling may average out positional and high-norm artifacts in teacher features, and the latent PCA is consistent with position-wise alignment introducing a spatial gradient artifact. With a text teacher, Pool-Align and Soft-KL improve generation quality, though the text-teacher evaluation uses one encoder and caption source and does not test text-to-image generation.

Why pooling can match position-wise alignment

  • Pooling may average out positional and high-norm artifact tokens that distract from content the decoder must reconstruct.A latent PCA is consistent with position-wise VA-VAE latents inheriting a smooth spatial gradient artifact absent from pooled latents.Source: §4 · p. 4

What transfers across modalities

  • With a text teacher, Pool-Align brings gFID from 9.52 to 6.97, close to its 5.77 with DINOv2, and Soft-KL also improves it.The authors note that captions largely restate each image’s ImageNet class, so part of the gain may be class-level semantics.Source: §4 · p. 4
  • Only CKA fails across modalities; the authors hypothesize its full centered between-image geometry is the strictest constraint.They test one text encoder and caption source, not text-to-image generation.Source: §4 · p. 4

5 Conclusion

The paper shows latent distillation can improve diffusability without requiring spatial correspondence. Pooled alignment performs as well as or slightly better than position-wise alignment for image teachers and extends to 1D latents and text teachers.

5 Conclusion

  • Pooled representations perform as well as or slightly better than position-wise alignment for image teachers.The conclusion frames this as improving diffusability without requiring spatial correspondence.Source: §5 · p. 4
  • Structure-agnostic distillation extends to 1D latents and to a text teacher, where Pool-Align and Soft-KL improve generation quality.Source: §5 · p. 4
  • Structure-agnostic distillation is a simple, flexible alternative to position-wise alignment that leaves the latent’s layout free.Source: §5 · p. 4

A Architecture and training details

The architecture builds on VA-VAE’s convolutional tokenizer and LightningDiT prior, with a Perceiver-style bottleneck variant for 1D latents. Frozen vision or text encoders provide teacher features, while the diffusion pipeline and distillation term are specified for the experiments.

A Architecture and training details

  • VA-VAE aligns latent cells to co-located vision-foundation features, and its VF baseline combines position-wise cosine matching with a within-image relational term.The retrained VF baseline reproduces gFID 2.11 under the shared 64-epoch LightningDiT budget.Source: A Architecture and training details · p. 7
  • The 2D tokenizer uses a convolutional VAE with a 32 × 16 × 16 posterior; the 1D variant replaces only its bottleneck with 32 learned Perceiver-style latents.The 1D resampler cross-attends to convolutional tokens and uses learned position queries for decoding.Source: A Architecture and training details · p. 7
  • LightningDiT-XL uses 28 transformer blocks, hidden size 1152, and 16 heads; the 1D setting adds learned positional embeddings and disables RoPE.Under guidance, VF gives FID 2.12 versus 2.11 reported for VA-VAE, while Pool-Align gives 1.99; guided FID is reported only for the 2D setting.Source: A Architecture and training details · p. 7
  • The vision teacher is frozen DINOv2 ViT-L/14, while the cross-modal teacher encodes each training image’s caption with BGE-large.Pool-Align uses a learned bias-free linear map from the latent into the teacher dimension.Source: A Architecture and training details · p. 7

A.1 Training and evaluation hyperparameters

The training and evaluation configuration is shared across arms, with the autoencoder trained before the diffusion prior on frozen latents. The table specifies ImageNet-1k data settings and identifies the gradient-norm-balancing target.

A.1 Training and evaluation hyperparameters

  • Every arm shares the training and evaluation configuration; only per-method distillation settings and latent shape differ.Source: §A.1 · p. 7
  • The autoencoder trains first, then the diffusion prior trains on frozen latents using ImageNet-1k images at 256 × 256.Images are normalized to [−1, 1] and randomly cropped without horizontal flips.Source: Table 2 · p. 8
  • The table’s weight target r is the gradient-norm-balancing factor defined in Appendix A.2.Source: Table 2 · p. 8

A.2 VAE training loss

The VAE tokenizer objective combines reconstruction, perceptual, variational, and adversarial terms, while gradient-norm balancing sets the distillation weight. The CKA objective uses an unbiased estimator based on between-image structure.

A.2 VAE training loss

  • Lvae combines L1 pixel reconstruction, VGG-LPIPS perceptual loss, learned-variance negative log-likelihood, a KL regularizer, and hinge adversarial loss.The KL term is weighted by 10^-6; the adversarial loss is enabled after 5000 steps and weighted by 0.5 times an adaptive factor.Source: §A.2 · p. 8
  • Gradient-norm balancing sets the distillation weight λ at the encoder’s last layer, and distillation acts on the posterior mean rather than the reparameterized sample.The posterior-mean choice avoids letting the encoder inflate posterior variance.Source: §A.2 · pp. 8–9

A.3 Unbiased CKA

  • The CKA objective uses a minibatch unbiased linear-CKA estimator that normalizes cross-space HSIC by within-space HSIC terms.The B × B kernels have zeroed diagonals, so only between-image structure enters.Source: §A.3 · p. 9

B Additional results and analyses

The Pool-Align margin over position-wise VF is close to the reported scale of gFID variation, so the authors read the methods as on par or slightly different. Latent PCA shows position-wise alignment retaining a spatial gradient while pooled latents track image content.

B.1 Significance of gFID gaps

  • Pool-Align’s gFID margin over VF is 0.27, or 0.13 with guidance, just above the ∼0.1–0.2 variation scale; with one run per objective, the authors read them as on par or slightly different.Improvements over No-Distillation exceed the variation scale, including the smallest cross-modal Soft-KL improvement of ∼1 FID.Source: §B.1 · p. 9

B.2 Latent PCA

  • Top-3 PCA shows the position-wise VF latent carrying a smooth horizontal gradient, while pooled latents pick out the foreground object without that gradient.The undistilled latent is nearly featureless; pooled latents track content rather than residual positional-embedding information in the teacher representations.Source: Figure 3 · p. 9

B.3 Samples

Figure 4 presents samples from Pool-Align on a 2D grid using a DINOv2 teacher in setting A, with the reported training and sampling recipe corresponding to gFID 1.99.

B.3 Samples

  • The displayed samples come from Pool-Align on the 2D grid with a DINOv2 teacher in setting A.The autoencoder trains for 50 epochs and the diffusion prior for 64 epochs (80k steps).Source: Figure 4 · p. 9
  • Sampling uses 250 Euler steps with classifier-free guidance and corresponds to gFID 1.99.The caption refers to the recipe in Appendix A.Source: Figure 4 · p. 9
  • The sample figure reports a 50-epoch autoencoder and a 64-epoch diffusion-prior training setup.The prior training is reported as 80k steps.Source: Figure 4 · p. 9

Related papers

Loading 2609.39657v1…