Source-linked AI summary
Pooling Representation Autoencoders for Efficient Diffusion
Ramón Calvo-González, Youssef Saied, François Fleuret
TL;DR
Dense representation token grids make generative modeling computationally expensive, despite strong correlations between neighboring features. PoolDINO learns affine pooling of neighboring frozen encoder features jointly with an RGB decoder, then trains generation on the compressed representations within the two-stage RAE pipeline. At 4× compression, generation quality remains comparable to the uncompressed reference while throughput increases; 16× trades visual fidelity for higher throughput, and classification accuracy decreases while dense-prediction gains over average pooling are inconsistent.
Problem
Dense representation token grids make generative modeling computationally expensive, despite strong correlations between neighboring features.
Source: §1 · p. 2Method
PoolDINO learns affine pooling of neighboring frozen encoder features jointly with an RGB decoder, then trains generation on the compressed representations within the two-stage RAE pipeline.
Source: §1 · p. 2Results
Takeaways & Limitations
Jointly learned local pooling compresses RAE latents without requiring a separate feature autoencoder or changing the standard two-stage training pipeline.
Source: §6 · p. 9Generation experiments use ImageNet-256 and a single autoencoder and generator architecture, and competitive quality relies heavily on guidance.
Source: §6 · p. 9
Abstract
from arXiv · showhide
Representation Autoencoders (RAEs) generate images from pre-trained visual fea- tures, but their dense token grids make generative modeling expensive. Motivated by local feature correlations, we introduce PoolDINO, a learned affine pooling operator that merges neighboring tokens. Training the pooling operator jointly with the RGB decoder preserves the standard two-stage RAE procedure without a separate feature auto-encoder. On ImageNet-256, 4x token compression retains comparable generation quality under internal guidance, while 16x compression trades some quality for greater efficiency. At a fixed budget of 100 sampling steps, latent-sampling throughput increases by 3.7x and 9.0x, respectively, relative to the unpooled baseline. Classification and dense prediction evaluations show that comparable guided generation quality can coexist with weaker performance on other tasks.
1 INTRODUCTION
Dense, high-dimensional representation tokens are costly for Transformer generation, although nearby features are strongly correlated. PoolDINO learns to compress neighboring tokens jointly with RGB decoding, preserving the two-stage RAE pipeline; evaluations report a compression–throughput–quality trade-off and weaker downstream performance on some tasks.
1 INTRODUCTION
- The introduction situates RAEs as a two-stage approach that decodes frozen vision-encoder features and then trains a flow-matching model to generate them.It also notes prior optimization challenges in high-dimensional latents and later RAEv2 gains.Source: §1 · p. 2
- Nearby feature tokens are strongly correlated, motivating the hypothesis that a high-dimensional pooled token can summarize its local neighborhood and reduce generative-model memory and compute.The introduction frames this as a hypothesis about local feature capacity.Source: §1 · p. 2
- PoolDINO learns an affine operator that merges neighboring encoder patches and trains it jointly with an RGB decoder, preserving the standard two-stage RAEv2 pipeline.The encoder stays frozen; the pooling operator and decoder are later frozen while the generator is trained on compressed representations.Source: §1 · p. 2
- At 4× compression, guided generation quality remains comparable to the uncompressed reference, while more aggressive compression trades fidelity for throughput and downstream task performance can weaken.Classification accuracy decreases despite comparable guided quality at 4×, and learned pooling does not consistently beat average pooling on dense prediction.Source: §1 · p. 2
- The evaluation examines compression, guided generation quality, and sampling throughput on ImageNet-256, alongside operator analysis, classification, segmentation, and depth estimation.These evaluations assess both generation and the semantic content of pooled representations.Source: §1 · p. 2
2 RELATED WORK
Prior work uses pretrained visual features for generation and explores compressing them, but some approaches add a feature-compression training stage or retain the original spatial token count. PoolDINO instead explicitly pools local windows while training compression for image reconstruction.
2 RELATED WORK
- RAEs use frozen vision-encoder features directly as generative latents, while REPA aligns diffusion hidden states to pretrained features and other methods align or tune the tokenizer and diffusion model.RAEv2 combines multiple encoder layers to improve reconstruction and refines generator training.Source: §2 · p. 2
- FAE reduces RAE latent channel dimensionality but adds a feature-autoencoder training stage and retains the spatial token count, leaving the generation-stage bottleneck unaddressed.FlatDINO also adds a compression stage; most of its tokens compress fixed spatial chunks independent of image content.Source: §2 · p. 3
- FlatDINO compresses dense DINOv2 features into a shorter sequence with a Transformer autoencoder, but most tokens encode fixed spatial chunks independent of image content.Under the evaluated sampling configurations, PoolDINO has a stronger quality–throughput trade-off than FlatDINO.Source: §2 · p. 3
- PoolDINO uses an affine map to pool local windows jointly with RGB reconstruction, avoiding the separate feature autoencoders used by FAE and FlatDINO while preserving the two-stage RAE pipeline.Its compression is directly optimized for image reconstruction.Source: §2 · p. 3
3 PRELIMINARIES
Flow matching learns a time-dependent velocity field from interpolated noise–data pairs and generates samples by integrating that field. RAE uses pretrained encoder features as spatial latents for image reconstruction, with RAEv2 averaging normalized activations across selected layers.
3.1 FLOW MATCHING
- Flow matching trains a time-dependent velocity predictor on interpolations between data and Gaussian noise, whose constant target velocity is the data–noise difference.The optimal predictor averages velocities compatible with the observed interpolated sample.Source: §3.1 · p. 3
- Generation starts from Gaussian noise and numerically integrates the learned velocity field to data; the same formulation applies when the modeled samples are latent representations.A model predicting the clean sample can be converted analytically to a velocity for t < 1.Source: §3.1 · p. 3
3.2 REPRESENTATION AUTOENCODERS
- An RAE maps an image to a spatial latent with a pretrained vision encoder and trains an image decoder to reconstruct the image from that latent.The latent grid has spatial dimensions h × w and token dimensionality D.Source: §3.2 · p. 3
- RAEv2 averages layer-normalized activations from selected intermediate encoder blocks, using DINOv3-L/16 blocks 11, 13, 15, 17, 19, 21, and 23 in this work.The layer selection balances earlier fine-grained spatial structure with deeper semantic context.Source: §3.2 · p. 3
- The pooled-decoder schematic shows a frozen DINOv3-L/16 patch field, one projected token per pooling window, and repeated token copies passed to ViT-XL for reconstruction.The displayed 4 × 4 field with 2 × 2 windows is illustrative.Source: Figure 2 · p. 4
4 METHOD
PoolDINO compresses local windows of RAE latents with learned affine pooling and trains the operator jointly with an RGB decoder. Generation then uses a model trained on the compressed grid, with internal guidance and auxiliary objectives supporting sampling.
4.1 IMAGE DECODING
- Local feature correlation and retained channel dimensionality motivate pooling each neighborhood into one token while preserving semantic structure for image decoding.The authors hypothesize this can heavily compress the token grid while retaining structure needed for high-quality decoding.Source: §4.1 · p. 4
- Each non-overlapping latent window is mapped to one token by a shared learned affine map applied to its concatenated features.The map is equivalent to a convolution whose kernel size and stride both equal the pooling-window dimensions.Source: §4.1 · p. 4
- Before RGB decoding, each pooled token is repeated over its original window so a ViT-XL decoder reconstructs the image on the original grid.This repetition keeps decoder architecture and input sequence length fixed across pooling geometries to control decoder compute.Source: §4.1 · p. 4
4.2 IMAGE GENERATION
- For each pooling geometry, a class-conditional DiTDH-XL is trained on the compressed latent grid with RAEv2’s transport objective and internal guidance.The generator uses channel-wise latent statistics and class-condition dropout; sampling can extrapolate between intermediate and full-depth predictions without a separate unconditional evaluation.Source: §4.2 · pp. 4–5
- Training combines full-generator flow matching with auxiliary intermediate-prediction and dense encoder-field reconstruction losses; sampled latents are denormalized, repeated, and decoded.The auxiliary heads are applied after the eighth DiT block, and generation uses the frozen RGB decoder.Source: §4.2 · p. 5
5 RESULTS
The results show that learned pooling preserves reconstruction and, at 4× compression, comparable guided generation quality, while higher compression and fewer sampling steps offer efficiency trade-offs. Classification and dense-prediction results show that this generation performance does not imply stronger downstream performance.
5.1 RECONSTRUCTION
- At 4× compression, learned pooling reaches 0.36 rFID, close to the unpooled reference’s 0.32 and below average pooling’s 0.65.At 16×, learned pooling reaches 0.41 rFID versus 1.81 for average pooling and 1.72 for max pooling; PSNR and LPIPS also favor learned pooling.Source: §5.1 · p. 5, Table 1 · p. 5
5.2 IMAGE GENERATION
- At 4× compression, IG-only generation reaches 1.09 FID versus 1.08 for the uncompressed baseline; 16× reaches 1.44 FID, with CFG further improving the lowest observed FID.CFG yields the largest improvements at 16×, but requires an additional unconditional forward pass per sampling step.Source: §5.2 · p. 6, Table 2 · p. 7
- The 2×2 average-pooled generator achieves 1.70 FID unguided versus 3.00 for its learned counterpart, but responds less effectively to IG at 1.29 FID and does not match the uncompressed baseline.Source: §5.2 · p. 6
- Extended training improves generation quality at both 2×2 and 4×4 pooling, but further training to approximately match the unpooled compute budget does not improve the best observed FID.The extended runs train the 2×2 and 4×4 generators for 180 and 300 epochs, respectively.Source: §5.2 · p. 7, Table 3 · p. 7
- Reducing sampling from 100 to 50 steps maintains comparable generation quality while providing a further twofold increase in latent-sampling throughput.Source: §5.2 · p. 7
5.3 CLASSIFICATION AND DENSE PREDICTION
- Learned-pooled features yield lower ImageNet classification accuracy than all three simpler baselines at every compression rate on both linear probing and k-NN evaluations.The learned operator retains a different subspace from average pooling and local PCA; its alignment with both decreases as compression increases, and it preserves less total variance than PCA at the same output dimension.Source: §5.3 · p. 7
- At 4× compression, average pooling performs better on segmentation and depth; at 16×, learned pooling has higher segmentation mIoU while average pooling retains lower depth error.These results do not establish a consistent transfer advantage for pooling learned through RGB reconstruction.Source: §5.3 · p. 9, Table 4 · p. 9
6 CONCLUSION
PoolDINO compresses RAE latents with a learned spatial pooling operator trained jointly with an RGB decoder, preserving the standard two-stage training pipeline. The conclusion reports a quality–throughput trade-off, while transfer results and experimental scope qualify the method’s performance.
6 CONCLUSION
- At 4× compression, PoolDINO substantially increases sampling throughput while retaining a guided FID highly competitive with the uncompressed baseline; 16× trades visual fidelity for higher throughput.The method trains a local affine mapping jointly with an RGB decoder, avoiding complex feature autoencoders and preserving the two-stage RAE pipeline.Source: §6 · p. 9
- Learned pooling yields lower k-NN and linear-probe accuracy than average pooling, local PCA, and random projection, and does not consistently outperform average pooling on dense prediction.The conclusion notes that strong guided generation quality does not guarantee equally strong downstream performance.Source: §6 · p. 9
- Generation experiments are limited to ImageNet-256 and one autoencoder and generator architecture; the pooling-and-decoder training budget is fixed, and longer first-stage training at higher compression remains untested.Competitive generation quality relies heavily on guidance, and unguided performance noticeably degrades under compression.Source: §6 · p. 9
A TRAINING AND MODEL DETAILS
- The first stage learns the local affine pooling operator and RGB decoder with the vision encoder frozen; the second freezes these components and trains the class-conditional generator.The pooled generator predicts the clean pooled latent from noisy latents, while auxiliary heads predict the same clean target and dense encoder features.Source: A TRAINING AND MODEL DETAILS · p. 12, Figure 5 · p. 12
- The RGB reconstruction objective combines L1 pixel error, LPIPS perceptual loss, and an adaptively weighted adversarial term, with discriminator and adversarial training introduced after six and eight epochs.Pooled-token noise is added during training, while evaluation uses clean tokens and EMA parameters.Source: A TRAINING AND MODEL DETAILS · p. 12
A.1 EXTENDED-TRAINING COMPUTE BUDGET
- The 2 × 2 and 4 × 4 runs use 180 and 300 epochs at approximately 81% and 77% of the unpooled baseline’s training budget, respectively.Compute-matched endpoints use 222 and 390 epochs at approximately 100.13% and 99.89% of the baseline budget.Source: §A.1 · p. 14
A.2 EFFECT OF TRAINING DURATION
Under the tested recipe, extending generator training brought little or no improvement in best observed FID, though Inception Score generally rose at fixed guidance scales. The throughput comparisons use specified hardware and measurement protocols, with cross-paper comparisons also reflecting software-stack differences.
A.2 EFFECT OF TRAINING DURATION
- For 2 × 2 pooling, best observed FID was 1.0462 at 180 epochs versus 1.0476 at 222; for 4 × 4, it was 1.2880 at 300 versus 1.3160 at 390.These results suggest diminishing returns under the tested recipe; repeated evaluations would be needed to assess statistical significance.Source: §A.2 · p. 15, Table 9 · p. 15, Table 10 · p. 15
- IS generally increases with training duration at fixed guidance scales, so FID and IS trends differ.The main generation tables and throughput figures retain the 180- and 300-epoch operating points for 2 × 2 and 4 × 4 pooling.Source: §A.2 · p. 15
B THROUGHPUT MEASUREMENT
- Throughput reporting distinguishes individual guided denoiser-step speed from complete latent-sampling trajectories; measurements use BF16 model computation and return full-precision model outputs.Means and sample standard deviations are computed across 30 runs of 20 denoiser steps or one complete sampling trajectory. Latent samples/s counts completed latent samples, not end-to-end RGB outputs.Source: B THROUGHPUT MEASUREMENT · p. 16
- The literature throughput comparison includes software-stack differences: baselines use PyTorch/Inductor, whereas PoolDINO uses JAX/XLA.Published literature FIDs, training budgets, sampling recipes, and class-balance protocols also differ from the authors’ evaluations.Source: B THROUGHPUT MEASUREMENT · p. 17
B.1 SAMPLING-STEP ABLATION
Halving sampling from 100 to 50 Euler steps doubled latent-sampling throughput with little change in generation quality across evaluated pooling geometries and training durations. Additional appendix figures and protocols describe step-count, hardware, and classification evaluations.
B.1 SAMPLING-STEP ABLATION
- Halving sampling steps from 100 to 50 doubles latent-sampling throughput, with an absolute FID difference below 0.04 in every evaluated configuration.Each checkpoint’s selected IG scale remains fixed during the comparison.Source: §B.1 · p. 17, Table 11 · p. 18
- The 300-epoch 4 × 4 generator is also evaluated across 10 to 100 sampling steps, with FID and IS plotted at a fixed IG scale of 2.75.Figure 8 separately compares five measured 50-step PoolDINO configurations with literature baselines.Source: §B.1 · p. 17, Figure 9 · p. 18, Figure 8 · p. 18
B.2 BATCH-SIZE AND HARDWARE SWEEPS
- Generator forward–backward throughput benchmarks compare pooling grids and hardware, measuring the base generator and main prediction loss while excluding optimizer updates and auxiliary objectives.The uncompressed model runs out of memory at batch size 256 on both GPUs.Source: §B.2 · p. 19, Figure 11 · p. 19, Figure 10 · p. 19
C IMAGE CLASSIFICATION DETAILS
- Classification probes compare learned pooling with average pooling, training-fitted local PCA, and random projection using spatially averaged features; dense-task evaluations retain the token grid.The k-NN and linear-probe protocols use frozen representations and report validation top-1 accuracy.Source: C IMAGE CLASSIFICATION DETAILS · p. 20
D DENSE PREDICTION DETAILS
The dense-prediction appendix reports segmentation and monocular-depth evaluations under scratch and RGB-decoder initialization settings. It specifies shared training choices and task-specific data, loss, and validation protocols.
D DENSE PREDICTION DETAILS
- Dense prediction is evaluated both with task models trained from scratch and with models initialized from the corresponding RGB decoder; scratch controls separate representation information from decoder-initialization transfer.The task model uses a ViT-XL decoder with a newly initialized output head.Source: D DENSE PREDICTION DETAILS · pp. 21–22
- ADE20K segmentation uses the official 150-class splits, random-scale 512 × 512 crops for training, pixelwise cross-entropy, and checkpoint selection by validation mIoU.Validation preserves aspect ratio, resizes the longest side to 512 pixels, and pads to a square.Source: D DENSE PREDICTION DETAILS · p. 21, Table 13 · p. 21, Table 15 · p. 22
- NYUv2 depth evaluation retains 480 × 640 targets in meters, predicts log depth, and selects the checkpoint with the lowest validation AbsRel.Evaluation restricts valid pixels to specified image rows and columns, clips predictions to 0.1–10 meters, and uses no scale alignment.Source: D DENSE PREDICTION DETAILS · p. 21, Table 14 · p. 22, Table 16 · p. 22
E GUIDANCE ABLATIONS
The guidance ablations compare classifier-free guidance (CFG), internal guidance (IG), and encoder-reconstruction guidance separately and in combination. Their FID rankings depend on compression, guidance pairing, and training duration.
E.1 FULL GENERATION COMPARISON
- For 80-epoch models, combining IG with CFG yields the lowest overall FID, with the largest additional reduction at 16× compression; this advantage disappears for the 180-epoch 2 × 2 model.The comparison uses ImageNet-256 and 100 Euler steps.Source: §E.1 · p. 23, Table 17 · p. 23
E GUIDANCE ABLATIONS
- The isolated sweeps vary one guidance scale at a time with the other two methods disabled; joint searches use measured configurations without interpolating unevaluated settings.The evaluations use 50,000 ImageNet-256 samples and 100 Euler steps; isolated sweeps use 80-epoch models.Source: E GUIDANCE ABLATIONS · p. 22, §E.2 · p. 23, §E.3 · p. 24
E.2 ISOLATED GUIDANCE SWEEPS
- Internal guidance gives the lowest observed isolated FID at every compression rate, while encoder-reconstruction guidance has a narrower optimum and degrades rapidly at large scales.Inception Score often continues to improve after FID reaches its minimum.Source: §E.2 · p. 24
E.3 JOINT GUIDANCE SEARCHES
- With CFG jointly tuned, encoder-reconstruction guidance has the lowest observed FID at 2 × 2, internal guidance is stronger at 2 × 4 and 4 × 2, and their 4 × 4 results are nearly identical.Both auxiliary heads were trained in every model; the comparison concerns sampling-time guidance.Source: §E.3 · p. 24
F LEARNED POOLING OPERATOR ANALYSIS
The learned pooling operator selects a row space distinct from spatial averaging and local PCA. Its measured alignment and captured feature energy characterize that geometry but do not show which directions support reconstruction or generation.
F LEARNED POOLING OPERATOR ANALYSIS
- Learned pooling selects a row space distinct from spatial averaging and local PCA; its adjusted alignment with both decreases as compression increases, and it captures less held-out feature variation than the training-fitted PCA baseline.These measurements characterize the operator but do not establish which retained directions support reconstruction or generation.Source: F LEARNED POOLING OPERATOR ANALYSIS · p. 29, Table 22 · p. 30
- The pooled token is y = Ax + b, and the input directions it retains are represented by an orthonormal basis of A’s row space; the bias does not affect that space.The analysis uses the right singular vectors of A, and the analyzed operators have numerical rank C.Source: F LEARNED POOLING OPERATOR ANALYSIS · p. 29
- The adjusted alignment score maps random-subspace overlap to zero and identical subspaces to one, but it is a geometric comparison rather than a significance test.The analysis uses one checkpoint per geometry and a finite block sample, so small differences warrant caution.Source: F LEARNED POOLING OPERATOR ANALYSIS · pp. 29–30
G GENERATED SAMPLES
The generated-sample figures show randomly selected examples from unpooled and pooled models, including standard and extended-training checkpoints. Corresponding grid positions share the same class condition.
G GENERATED SAMPLES
- Each model is shown with 48 randomly selected samples in an eight-column, six-row grid; corresponding positions share a class condition, and generation uses 100 Euler steps with IG alone.Each model uses its selected IG scale.Source: G GENERATED SAMPLES · p. 31
- The 80-epoch examples cover the unpooled reference and PoolDINO at 2 × 2, 2 × 4, and 4 × 4 pooling, with token compression of 4×, 8×, and 16×, respectively.The figure captions specify sIG values of 1.75, 2.00, and 2.75 for the three pooled settings.Source: Figure 17 · p. 31, Figure 18 · p. 32, Figure 19 · p. 33
- Extended-training examples show 2 × 2 pooling at 180 epochs and 4 × 4 pooling at 300 epochs, both sampled with IG alone.Their captions specify sIG = 2.00 and sIG = 2.75, respectively.Source: Figure 21 · p. 35, Figure 22 · p. 36