💻 computer science

One Pass Is Not Enough: Recursive Latent Refinement for Generative Models

The paper introduces RTM, a generative model that replaces single-pass latent mapping with recursive refinement to explicitly prioritize mode coverage, thereby achieving superior precision and recall alongside competitive FID scores across multiple datasets.

Original authors: Mehdi Esmaeilzadeh, Alexia Jolicoeur-Martineau, Chirag Vashist, Ke Li

Published 2026-05-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mehdi Esmaeilzadeh, Alexia Jolicoeur-Martineau, Chirag Vashist, Ke Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot artist how to paint portraits. You show it thousands of photos of people. The goal is for the robot to eventually paint a new picture of a person that looks real but has never existed before.

For a long time, the main way we judged these robots was by asking: "Does this painting look sharp and realistic?" If the robot painted 1,000 copies of the exact same handsome man, it would get a perfect score because every single painting looked great. But this is a problem: the robot failed to learn that people have different hair colors, ages, and expressions. It got stuck in a "mode collapse," only painting one type of person.

This paper introduces a new method called RTM (Recursive Token Mapper) to fix this. Here is how it works, using simple analogies:

1. The Problem: The "One-Shot" Artist

Most current AI models (like the popular StyleGAN) work like a student taking a final exam in one single, frantic rush.

The Old Way: The AI takes a random noise signal (like a blank canvas) and runs it through a long chain of 8 steps (an 8-layer "MLP") all at once to decide what the picture should look like.
The Flaw: Because it has to decide the face shape, the skin tone, the hair texture, and the lighting all in one go, it often gets overwhelmed. It tends to play it safe, picking the "average" or most common features to ensure the image looks sharp. This leads to high quality but low diversity (the "mode collapse" problem).

2. The Solution: The "Iterative Sketch" Artist

The authors propose a new approach called RTM. Instead of doing everything in one rush, the AI acts like a master artist who sketches and refines.

The Analogy: Imagine an artist who doesn't try to paint a masterpiece in one stroke.
1. First Pass (Coarse): They quickly sketch the basic outline: "Okay, this is a face, it's facing left, the hair is short."
2. Second Pass (Refinement): They look at that sketch and say, "The jawline needs to be sharper, and the eyes need more detail."
3. Third Pass (Polishing): They add the final touches: "Let's make the skin texture look real and adjust the lighting."

RTM does exactly this. It takes the random noise and runs it through a small, reusable "block" of logic multiple times.

Early cycles establish the big picture (identity, pose, composition).
Later cycles refine the details (texture, color, sharpness).

Because the AI gets to "look back" at its own work and fix mistakes, it doesn't get stuck on just one type of image. It can explore a wider variety of faces while still keeping them looking realistic.

3. The "Recursive" Secret Sauce

You might ask, "Why not just make the chain of 8 steps longer (e.g., 32 steps)?"
The paper argues that just making the chain longer is like giving the student a longer list of instructions to memorize without letting them think. It's inefficient and can actually make things worse.

RTM is "Recursive." This means it uses the same small set of instructions over and over again, but each time it applies them to the improved result of the previous step.

Analogy: It's like having a single, very smart editor who reviews a draft, edits it, reviews the new draft, edits it again, and so on. This is much more effective than hiring 32 different editors who each do a tiny part of the job without talking to each other.

4. The Results: Better Quality AND More Variety

The authors tested this on several datasets (like CIFAR-10, which has small images of cats, dogs, and cars, and CelebA-HQ, which has faces).

The Old Metric Trap: They noted that the standard score (FID) is getting "saturated." It's hard to make it lower, and a low score doesn't guarantee the AI is learning the full variety of the data.
The New Goal: They focused on Precision (how realistic the images look) and Recall (how many different types of images the AI can generate).
The Outcome: RTM beat the previous best models. It didn't just make sharper images; it made more diverse images.
- On the "Face" dataset, it generated faces with a much wider variety of ages, skin tones, and expressions compared to the old models, while still looking very realistic.
- It worked not just for their specific training method (IMLE) but also improved the standard StyleGAN models.

Summary

Think of the old AI as a photocopier that only knows how to copy one specific photo perfectly. The new RTM AI is like a creative director who can look at a rough idea, refine it step-by-step, and produce a gallery of unique, high-quality portraits that cover the entire spectrum of human diversity.

The paper claims this is achieved by replacing the "one-shot" mapping network with a recursive loop that allows the AI to iteratively refine its latent "blueprint" before generating the final image.

Technical Summary: One Pass Is Not Enough: Recursive Latent Refinement for Generative Models

Problem Statement

Despite significant progress in image generation, the field faces a critical evaluation bottleneck. The dominant metric, Fréchet Inception Distance (FID), conflates sample fidelity with mode coverage, often saturating at low single-digit values. This saturation masks a persistent issue: models can achieve low FID scores by generating a small set of sharp, near-duplicate images (mode collapse) while failing to cover the full data distribution.

Current generative paradigms exhibit a trade-off:

Diffusion and Flow-Matching models: Achieve high fidelity but often require iterative sampling (multiple neural function evaluations). Distilled one-step variants accelerate inference but consistently sacrifice diversity and coverage. Furthermore, diffusion solvers systematically under-cover low-density regions of the data distribution.
Generative Adversarial Networks (GANs): Offer single-step inference but are prone to mode collapse, concentrating on high-fidelity modes while dropping others.
Implicit Maximum Likelihood Estimation (IMLE): Guarantees mode coverage by construction (ensuring every training image has a nearby generated sample) and supports single-step inference. However, IMLE-based methods have historically lagged in sample quality. This limitation stems partly from the generator's architecture: the mapping network in StyleGAN-based IMLE models is a standard Multi-Layer Perceptron (MLP) processed in a single forward pass. This forces the network to determine all aspects of the style code (identity, structure, texture, fine detail) simultaneously, making the output highly sensitive to inaccuracies in the initial latent mapping.

Methodology: Recursive Token Mapper (RTM)

The authors propose replacing the single-pass MLP mapping network in style-based generators with a Recursive Token Mapper (RTM). RTM adapts the "Tiny Recursive Model" (TRM) architecture to the generative setting, enabling iterative refinement of the latent space without increasing the parameter count.

Core Architecture

Recursive Structure: Instead of a deep stack of independent layers, RTM uses a shared block $f$ $f$ applied iteratively. The process involves two levels of recursion:
- Inner Cycles ( $L$ ): An inner state $Z_L$ is updated $L$ times per step, conditioned on an outer state and the original noise vector.
- Refinement Steps ( $H$ ): The outer state $Z_H$ is updated $H$ times.
Mechanism: The mapper produces a coarse estimate of the style vector $w$ first and progressively corrects it over multiple cycles. Early cycles establish coarse structures (identity, pose), while later cycles refine texture and color.
Parameter Efficiency: Because the same block $f$ is reused, increasing $H$ or $L$ increases the effective computational depth without adding parameters.
Short-Gradient Optimization: To manage memory during training, only the final refinement step is differentiated through; intermediate steps are detached. This preserves the representational benefits of deep recursion while keeping training memory costs comparable to a single feed-forward block.
Integration: RTM is designed as a "drop-in" replacement for the mapping network in both IMLE-based generators (specifically RS-IMLE) and adversarially trained StyleGANs (StyleGAN2 and StyleGAN2-ADA). The decoder architecture remains unchanged.

Key Contributions

Recursive Token Mapper (RTM): A novel mapping network architecture that replaces the standard single-pass MLP in style-based generators with an iterative refinement process.
Integration with RS-IMLE: When integrated with Rejection-Sampling IMLE (RS-IMLE), RTM improves FID, Precision, and Recall across nine few-shot benchmarks, unconditional CIFAR-10, and CelebA-HQ (256×256), while retaining IMLE's one-step inference and direct latent-to-image mapping.
Generalizability to Adversarial Training: RTM improves StyleGAN2 and StyleGAN2-ADA on CIFAR-10 and AFHQ-v1 (512×512), demonstrating that the benefits of recursive refinement are not specific to the IMLE training objective.
Inference Flexibility: The number of refinement steps ( $H$ ) acts as a computational hyperparameter. Models can be evaluated with different $H$ values at inference time without retraining, allowing a trade-off between inference cost and fidelity.

Experimental Results

The paper evaluates RTM using Precision, Recall, FID, and Inception Score (IS) across multiple datasets.

CIFAR-10 (32×32): RTM integrated with RS-IMLE achieves the highest Precision (0.896) and Recall (0.773) among all methods listed, including diffusion and flow-matching baselines. It lowers FID by 30% compared to the vanilla RS-IMLE baseline (3.97 vs. 5.69) while maintaining one-step inference.
CelebA-HQ (256×256): RTM reduces FID by 27% over the RS-IMLE baseline (10.67 vs. 15.43) and simultaneously improves Precision (0.952) and Recall (0.592).
Few-Shot Benchmarks: On nine standard few-shot datasets (e.g., Obama, Dog, Anime), RTM roughly halves the average FID of the matched RS-IMLE reproduction (13.94 vs. 25.2).
StyleGAN2/ADA: On CIFAR-10, RTM lowers FID from 3.88 to 3.55 for StyleGAN2 and maintains FID at 2.31 for StyleGAN2-ADA while improving Precision and Coverage. On AFHQ-v1, RTM improves FID (4.79 vs. 4.99) and Recall (0.565 vs. 0.507).
Depth vs. Recursion: An ablation study on CIFAR-10 shows that simply increasing the depth of a non-recursive MLP (to 32 layers) improves Precision but severely degrades Recall (0.550 vs. 0.664 baseline). RTM achieves better FID with 13× fewer parameters than the depth-matched MLP, suggesting recursion provides a superior inductive bias for distributional fidelity without sacrificing coverage.

Significance and Claims

The paper argues that the goal of image generation should shift from merely lowering FID (which is near saturation) to explicitly improving diversity and coverage (Recall) alongside fidelity (Precision).

Beyond FID: The authors claim that FID alone is insufficient because it rewards sharp, concentrated samples that may ignore large portions of the data distribution. Precision and Recall are essential complements.
Simultaneous Improvement: Unlike flow-matching baselines that often trade coverage for FID, or distilled diffusion models that lose diversity in one-step inference, RTM improves both quality and diversity simultaneously.
Architectural Insight: The paper posits that refinement is a qualitatively different computation than a single forward pass. By allowing the mapper to revise early decisions (coarse structure) in light of later computations (fine details), RTM overcomes the sensitivity of the decoder to latent code inaccuracies.
Modesty: The authors acknowledge limitations, noting that experiments are restricted to datasets where IMLE's nearest-neighbor search is computationally feasible (excluding ImageNet). They do not claim to have solved image generation but rather to have provided a method that better balances the fidelity-diversity trade-off within current evaluation frameworks.