PaperlayerView exact revision on arXiv

Source-linked AI summary

LoopVAE: Recurrent Depth Across Scales for Visual Tokenization

Zhiying Lu

arXiv:2609.11516v1cs.CV

TL;DR

Hierarchical visual tokenizers must decide how to allocate processing parameters across depth and spatial scales. LoopVAE shares a conditioned core within and across scales while keeping resolution transitions independent, and reports a compact CNN operating point alongside matched sharing, recurrence, and runtime analyses. The paper presents recurrent depth across scales as a parameter-sharing organization, while showing that it does not automatically provide efficient execution or reliable early exits.

  • Problem

    The paper asks how hierarchical visual tokenizers should allocate processing parameters across depth and spatial scales.

  • Method

    LoopVAE reuses a scale- and loop-conditioned core within and across scales while retaining independent resolution-changing transitions.

  • Results

    29M parameters, 0.28 rFID, and 32.54 dB PSNR are reported for the CNN on ImageNet-256, while global sharing remains competitive in matched Transformer fidelity metrics and unshared blocks improve LPIPS.

  • Takeaways & Limitations

    Recurrent depth across scales is a viable parameter-sharing organization for visual tokenization, but parameter savings leave arithmetic, runtime, and intermediate-output calibration costs.

  • Takeaways & Limitations

    The sharing comparison is a single-run, non-adversarial Transformer experiment, and the loop study is a targeted TokBench case analysis rather than a full benchmark evaluation.

Abstract

from arXiv · show

Hierarchical visual tokenizers typically allocate different processing blocks to different spatial scales. We ask how much of this computation can use the same parameters. LoopVAE reuses a scale- and loop-conditioned core within and across scales, while keeping resolution-changing transitions independent. A four-block core executes 28 block applications per encoder or decoder. On ImageNet-256, the 29M-parameter convolutional model reaches 0.28 rFID and 32.54 dB PSNR under an approximately 30-epoch two-stage training budget, using approximately 65% fewer parameters than the 84M reference VAEs. A non-adversarial Transformer ablation with the same execution graph finds competitive PSNR and SSIM under global sharing, although unshared blocks improve LPIPS. Targeted loop interventions show that completing the trained recurrence improves reconstruction and that even small feature updates can have substantial downstream effects. Truncation also exposes output-range errors, distinguishing useful recurrent computation from reliable early exit. Runtime profiling reveals the execution tradeoff: fewer stored weights require more arithmetic and longer runtime in the tested configurations. With convolutional and Transformer operators and single- or multi-resolution latent interfaces, LoopVAE establishes recurrent depth across scales as a parameter-sharing design axis for visual tokenization.

1 Introduction

LoopVAE asks whether hierarchical visual tokenizers can share processing parameters across depth and spatial scales while retaining independent resolution transitions. It evaluates reconstruction quality, recurrent-depth contributions, and the tradeoff between parameter storage and execution cost.

  • Method: LoopVAE reuses a scale- and loop-conditioned core across spatial scales while retaining independent parameters for resolution-changing transitions.The encoder and decoder each use a shared fixed-resolution processing core, with transitions and heads remaining scale specific.
  • Method: A four-block core executes 28 block applications per encoder or decoder, separating executed depth from the number of stored core blocks.The design uses recurrent passes within scales and weight reuse across scales.
  • Empirical results: 29M parameters, 0.28 rFID, and 32.54 dB PSNR are reported for the final CNN on ImageNet-256 under an approximately 30-epoch two-stage budget.The result accompanies evidence from sharing and intervention studies rather than establishing a universal quality advantage.
  • Empirical results: Global sharing gives the highest PSNR and SSIM in the matched Transformer comparison, whereas fully unshared processing gives the lowest LPIPS.The comparison varies stored core depth while holding the execution graph fixed.
  • Empirical results: Completing the trained recurrence improves reconstruction, but truncation can produce raw-output errors, so recurrent computation does not guarantee reliable early exits.Small relative feature changes can still substantially affect final outputs.
  • Execution cost: 65.3% fewer parameters coincide with 1.57× the estimated MACs and 2.85× the batch latency for the 8x looped CNN versus a local Flux-style reference.Weight reuse reduces stored parameters without removing repeated high-resolution computation.

2 Related Work

LoopVAE builds on recurrent-depth, visual-compression, and autoencoder architectures while applying weight sharing across a spatial hierarchy. Its related work distinguishes the tokenizer setting and multi-resolution latent interfaces from earlier recurrent or looped systems.

  • Continuous visual tokenizers: Perceptual and adversarial objectives, stronger compression, semantic supervision, and Transformer scaling represent complementary directions for improving visual tokenizers.LoopVAE instead focuses on how processing parameters are allocated across depth and spatial scales.
  • CNN and Transformer autoencoders: Convolutional and Transformer autoencoders provide alternative spatial mixers inside LoopVAE’s shared recurrent hierarchy, rather than defining its central research question.The related work contrasts locality and translation equivariance with content-dependent global interaction.
  • Weight tying and recurrent depth: LoopVAE extends recurrent-depth ideas from Universal Transformers and looped Transformers to a visual autoencoder whose core is tied across recurrent passes and spatial scales.Prior work separates executed depth from independently parameterized depth, while LoopVAE studies that distinction in visual tokenization.
  • Recurrence in visual compression and generation: Unlike recurrent image compression with progressive quantized coding, LoopVAE uses a continuous latent interface and reuses a core across a spatial hierarchy.The comparison places LoopVAE in visual compression while distinguishing its objective and architecture.
  • Multi-resolution latent representations: LoopVAE-Multi attaches several posterior heads to one shared hierarchy, extending common single-bottleneck tokenizer designs toward multi-resolution latent representations.The related-work context links compression ratio and latent spatial size to the representation exposed to generative models.

3 Method

LoopVAE replaces scale-specific fixed-resolution processing with a scale- and loop-conditioned core reused across encoder or decoder scales and recurrent steps, while retaining independent resolution-changing modules. Its single- and multi-resolution designs support convolutional and Transformer operators with shared latent interfaces and joint training objectives.

  • Scale-Conditioned Depth Recurrence: LoopVAE reuses one conditioned fixed-resolution operator across scales and loop steps, while downsampling, upsampling, projections, and heads remain independently parameterized.The encoder and decoder use separate operators, but each branch shares its core wherever tensor geometry is unchanged.
  • Scale-Conditioned Depth Recurrence: Input injection adds a scale-specific projected feature before the shared core on later iterations, using a zero-initialized learned scalar gate.The injection occurs once before the entire core rather than before every block; the first iteration omits it.
  • Operator Instantiations: The core supports conditioned ConvNeXt-style convolutional blocks with AdaGN and DiT-style Transformer blocks with AdaLN, attention, and GELU MLPs.Both alternatives use two gated residual branches, while scale and loop embeddings condition the blocks.
  • Single- and Multi-Resolution Tokenizers: LoopVAE-Single uses an f8d16 latent, while LoopVAE-Multi attaches independent posterior heads at 8x, 16x, and 32x scales and decodes each through a shared hierarchy suffix.The multi-resolution model gives each latent shape its own input projection and trains reconstruction, KL, and adversarial losses jointly across interfaces.
  • Training Objective: Training combines reconstruction, perceptual, and KL terms in Stage 1, then continues from those weights with an adversarial objective in Stage 2.The adversarial stage is intended to sharpen high-frequency detail.

4 Experiments

The experiments evaluate LoopVAE’s reconstruction quality, sharing scope, recurrent-depth sensitivity, and execution behavior across CNN and Transformer configurations. They show competitive quality with substantial parameter savings, while repeated computation remains useful and parameter reuse increases measured execution cost.

  • 4.2 Reconstruction Results: 0.28 rFID and 32.54 PSNR are achieved by the 29M-parameter LoopVAE-CNN on ImageNet-256 under an approximately 30-epoch two-stage budget.It uses approximately 65% fewer parameters than the 84M reference VAEs, while baseline training recipes are not matched.
  • 4.3 Non-Adversarial Sharing Ablation: Global sharing has the highest PSNR and SSIM in the Transformer ablation, exceeding fully unshared processing by 0.120 dB and 0.0023 with one seventh as many stored core blocks.Fully unshared processing achieves the lowest LPIPS, improving over global sharing by 0.0055.
  • 4.4 Recurrent Depth and Loop-Position Sensitivity: 32.28 dB mean PSNR is reached as the encoder f8 prefix grows from one to four passes, with non-decreasing PSNR and SSIM across all evaluated multi-pass stages and images.The intervention keeps the remaining network on its trained schedule, so it tests recurrent depth without retraining.
  • 4.4 Recurrent Depth and Loop-Position Sensitivity: The second encoder f8 pass and third decoder f2 pass are the most bypass-sensitive positions on every tested image, showing that shared weights do not make loop positions equivalent.Bypass losses cannot be treated as independent additive contributions, and the experiments do not isolate loop conditioning from the evolving hidden state.
  • 4.4 Recurrent Depth and Loop-Position Sensitivity: A relative feature update below 0.1 can still matter: the fourth encoder f8 pass changes features by 7.47–9.36% but bypassing it loses 4.00–11.94 dB PSNR.The relative norm describes current feature coordinates rather than a calibrated output-error estimate or validated stopping criterion.

(b) Cost of bypassing the same pass

The interventions show that bypassing trained passes can substantially damage reconstruction, while runtime profiling separates parameter savings from execution efficiency.

  • Reconstruction sensitivity: 7.78 dB and 12.92 dB PSNR gains follow the final encoder f8 and decoder f2 passes, respectively, in the tested checkpoint.Decoder f2 prefixes also produce severe out-of-range RGB values before the complete schedule.
  • Downstream application: The tested downstream diffusion runs reach FID-50k values of 39.906 for B/2 and 18.198 for XL/2, but do not isolate tokenizer quality.No matched alternative tokenizer is included, and equal update counts do not imply equal training budgets.
  • Runtime trade-off: 65.3% fewer parameters still require 1.57× the estimated MACs and 2.85× the batch latency of the Flux-style reference.The benchmark reports 281.05 versus 98.73 ms batch latency and higher peak allocated memory for the looped CNN.
  • Runtime trade-off: The finest decoder stage accounts for 63.4% of spatially linear CNN core work, concentrating arithmetic where truncation is also quality-sensitive.This 64/101 share is an analytical decomposition of the specified graph, excluding transitions and overhead.

5 Discussion and Limitations

LoopVAE’s evidence supports recurrent parameter sharing as a capacity choice, but leaves important questions about convergence, execution trade-offs, and generality unresolved.

  • Sharing trade-offs: Global sharing remains competitive in paired fidelity, while unshared processing improves LPIPS; unmatched baseline recipes prevent a causal advantage claim.The sharing comparison is a single-run, non-adversarial Transformer experiment.
  • Recurrence and truncation: Trained passes can matter despite small feature changes, but truncation produces severe raw-output errors and does not establish calibrated early exit or convergence.Bypass tests dependence within the trained network rather than the performance of a shallower model trained from scratch.
  • Execution trade-offs: Because the finest decoder stage dominates work and remains truncation-sensitive, reducing deployment cost likely requires retraining or schedule redesign rather than simply omitting passes.The paper proposes matched-budget schedule comparisons and intermediate-output supervision or distillation as next tests.
  • Scope and evaluation: The evidence is bounded by a single sharing run, targeted loop cases, unavailable matched tokenizer comparisons, and untested higher-resolution, video, and text-conditioned generation.Repeated sharing or schedule runs and broader intervention evaluations are identified as needed extensions.

6 Conclusion

LoopVAE demonstrates recurrent depth across scales as a parameter-sharing organization for visual tokenization. Its benefits are paired with sensitivity to the trained schedule and substantial arithmetic and runtime costs.

  • Conclusion: A four-block shared core executes 28 block applications per branch, while the 29M-parameter CNN reaches 0.28 rFID and 32.54 dB PSNR on ImageNet-256.The model uses an approximately 30-epoch two-stage training budget and supports CNN and Transformer implementations with single- or multi-resolution interfaces.
  • Conclusion: Trained passes contribute to reconstruction even when relative feature changes are small, but truncated outputs are not reliably calibrated for early exit.The conclusion distinguishes useful recurrent computation from convergence-like or anytime behavior.
  • Conclusion: High-resolution recurrence costs substantial arithmetic and runtime despite the small stored core, making schedule design and intermediate-output training concrete improvement directions.The paper frames these directions as ways to improve the execution trade-off.

A.1 LoopVAE-Single Configuration

The main LoopVAE-Single configuration uses separate encoder and decoder parameters, a reverse-traversal decoder schedule, and a two-stage training recipe totaling approximately 30 nominal epochs.

  • Architecture: The decoder traverses scales from 8x to 2x, with stored loop list [1, 2, 4] interpreted in reverse-traversal order.The list denotes one, two, and four repeated passes at successive decoder scales.
  • Training: Stage 1 trains without GAN for 600k steps, after which Stage 2 continues from those weights with GAN training.The main CNN and ViT results use a nominal budget of approximately 30 epochs split roughly equally between stages.
  • Training: Under the reported data-iteration convention, 600k and 610k steps correspond to approximately 15.0 and 15.2 nominal epochs, while 220k corresponds to approximately 5.5.These are data-exposure estimates rather than exact reconstructions of checkpoint counters.
  • Training: Stage 2 uses a frozen DINOv2-small feature extractor, a trainable discriminator head, hinge loss, and differentiable augmentation with probability 1.0.The adversarial term activates after discriminator warm-up at the configured global-step threshold.

A.2 LoopVAE-Multi Configuration

The multi-resolution LoopVAE uses shared Transformer blocks across selectable latent interfaces and traversal scales, while changing entry points alters the executed hierarchy.

  • A.2 LoopVAE-Multi Configuration: The five encoder scales are 2x, 4x, 8x, 16x, and 32x, with loop schedule [2] [3] and six shared Transformer blocks at width 768.Attention uses local 8 × 8 windows at the two finest scales and global connectivity elsewhere.
  • A.2 LoopVAE-Multi Configuration: Posterior heads attach to scales 2, 3, and 4, while per-interface losses use weights 0.5, 1.0, and 1.5 with KL weight 10^-6.Stage 2 uses a convolutional discriminator with depth setting three.
  • A.2 LoopVAE-Multi Configuration: The decoder supports 8x, 16x, or 32x latent entry, skipping progressively more synthesis stages while reusing the same six Transformer blocks.These entry points execute 24, 42, and 78 core block applications, respectively.

A.3 Evaluation Protocol

Evaluation uses posterior-mode reconstruction metrics on matched image pairs, while reported results come from individual checkpoints with distinct evaluation scopes and no seed-level uncertainty.

  • A.3 Evaluation Protocol: PSNR, LPIPS, and SSIM are averaged over paired images after mapping tensors to the same image range, while rFID uses the same 50k ImageNet validation images as references and reconstructions.All validation images are encoded with the posterior mode and decoded once.
  • A.3 Evaluation Protocol: Reported scores are individual checkpoint evaluations rather than averages across training seeds, and the Transformer sharing ablation uses paired reconstruction metrics at Stage 1, step 400k.The ablation is non-adversarial and does not establish identical data ordering or seed-level uncertainty.
  • A.3 Evaluation Protocol: Final reconstruction, CNN diagnostics, Multi evaluation, and synthetic profiling use separate checkpoint or configuration scopes, so their results are not combined.Latent interpolation is not evaluated in this paper.

A.4 Loop Intervention Protocol and Output Range

The intervention protocol traces prefix and bypass behavior on a targeted image collection using raw outputs and explicit range-aware metrics. It tests recurrent sensitivity and output reliability while separating measured clipping effects from reconstruction error.

  • A.4 Loop Intervention Protocol and Output Range: The diagnostic uses raw step-610k CNN weights without EMA substitution and evaluates 16 targeted TokBench images after standardized 256 × 256 preprocessing.Each image yields 14 prefix records, 14 bypass records, and a normal 14-pass trace.
  • A.4 Loop Intervention Protocol and Output Range: Raw PSNR evaluates unclipped mapped RGB values, whereas standard PSNR and SSIM clip outputs to [0, 1]; out-of-range pixels count when any channel exceeds that interval.Clipping displacement measures mean absolute RGB change caused by clipping, not reconstruction error.
  • A.4 Loop Intervention Protocol and Output Range: 32.279 dB clipped PSNR and 32.251 dB raw PSNR are the full model’s subset averages, showing that out-of-range fraction alone does not determine distortion severity.One image has 55.73% out-of-range pixels but only 0.00169 mean clipping displacement.
  • A.4 Loop Intervention Protocol and Output Range: The trace analysis joins 224 trace–bypass pairs and measures feature-update variability across images, but these measurements characterize one trained checkpoint rather than semantic specialization or population-level behavior.RMS is computed over the entire feature tensor using the pre-injection state as reference.
  • A.4 Loop Intervention Protocol and Output Range: Generation runs use the Stage-2 610k CNN tokenizer and 50,000 generated images, while profiling times 100 batched passes after ten warm-ups and excludes diffusion sampling and discriminator execution.The XL/2 generator run uses 80 epochs independently of the tokenizer’s two-stage budget.
Loading 2609.11516v1…