Recursive Token Mapper (RTM)

• Recursive Token Mapper (RTM) is a recursive neural module that maps tokens to representations by iterative refinement using parameter-shared blocks.
• RTMs employ dynamic recursion depth and router-guided early exits to balance global feature extraction with fine-grained details.
• RTMs improve generative sample quality and diversity while reducing computational overhead, as evidenced in style-based image synthesis and language models.

A Recursive Token Mapper (RTM) is a neural module that maps input tokens or latent vectors to output representations through a sequence of recursive refinement steps. Unlike traditional mapping strategies relying on single-pass feedforward transformations, RTMs iteratively refine hidden states using parameter-shared blocks, with recursion depth that may be fixed or dynamically determined per token. Originally introduced to enhance latent mapping in generative models, notably in style-based image synthesis and later generalized in LLMs, RTMs improve both sample quality and diversity, increase parameter and compute efficiency, and enable dynamic, adaptive computation tailored to the complexity of each token or sample (Esmaeilzadeh et al., 14 May 2026, Bae et al., 14 Jul 2025).

1. Core Architecture and Algorithmic Structure

In the context of generative models, such as those based on StyleGAN, RTMs replace the standard mapping network — typically a multi-layer perceptron (MLP) — with a recursively applied module. Instead of processing the input noise vector $z$ in a single forward pass, the RTM transforms $z$ into a grid of latent tokens $Z_0$, which are then refined through $H$ outer "refinement steps", each with $L$ inner "cycles". Each update is performed by a parameter-shared block, $f_\theta$, often instantiated as an MLP-Mixer layer with RMSNorm/LayerNorm and channel as well as token mixing via two SwiGLU MLPs.

For each recursion step $h=1,\dots,H$, $Z_L$ (inner state) is updated for $L$ cycles, receiving $Z_H$ (outer state) and input injection of $z$0. The outer state $z$1 is then updated based on the refined $z$2. The shared parameters $z$3 are re-used across all iterations, enabling an effective depth of $z$4 with constant parameter budget. Once the refinement steps conclude, $z$5 is flattened and projected to obtain the final style code $z$6 for the synthesis network (Esmaeilzadeh et al., 14 May 2026).

Recent generalizations in language modeling (e.g., Mixture-of-Recursions, MoR) instantiate the RTM as a single Transformer block $z$7 repeatedly applied across all tokens. Each token $z$8's hidden state $z$9 is recursively computed as $Z_0$0, with the recursion depth either fixed or token-specific and learned by a lightweight router network (Bae et al., 14 Jul 2025).

2. Mathematical Recursion Formalism

The RTM recursion for generative latent mapping is given as:

• Initial projection: $Z_0$1
• Inner loop (for $Z_0$2): $Z_0$3
• Outer update: $Z_0$4
• Final output: $Z_0$5

Token-level recursion in the MoR RTM setting is formalized as: $Z_0$6 with the recursion depth $Z_0$7 assigned per token (either deterministically via $Z_0$8 or stochastically), typically informed by router outputs $Z_0$9 calculated as softmax-normalized scores from the hidden state.

3. Routing, Adaptive Depth, and Memory Efficiency

A distinguishing feature of advanced RTMs is dynamic, token-level recursion depth assignment. In the MoR architecture, router networks output probability vectors $H$0 for each possible depth $H$1, enabling per-token exit at custom depths:

• Hard routing: $H$2
• Soft routing: weighted sum of final states across depths

Load-balancing, entropy, and auxiliary $H$3-loss regularizers ensure uniform utilization of recursion depths and stable routing. To address quadratic attention bottlenecks, only tokens still "active" at each recursion receive further compute, and key–value (KV) caching and sharing are used to minimize memory and redundant computation. Recursive KV sharing reuses the computed KV pairs from the first recursion across later depths, reducing memory and I/O by up to $H$4 (Bae et al., 14 Jul 2025).

4. Training Objectives and Losses

For generative latent mapping, RTM is trained within the Implicit Maximum Likelihood Estimation (IMLE) framework. Each real image $H$5 in the dataset is paired to its nearest generated image in an embedding metric, ensuring strong mode coverage by construction. The IMLE loss remains: $H$6 where $H$7 indexes the nearest $H$8 in a large pool under the chosen feature extractor $H$9. Rejection-sampling is employed to exclude generated samples that are too close to any training image, closing the gap between train and test priors (Esmaeilzadeh et al., 14 May 2026).

For recursive mappers in language modeling, standard autoregressive or masked language modeling losses are coupled with router-specific regularizers (e.g., load balancing, entropy penalties).

5. Implementation and Quantitative Results

The RTM shared block is typically implemented as an MLP-Mixer layer: including RMSNorm, a SwiGLU MLP across tokens (sequence axis), and a SwiGLU MLP across channel axis (hidden size $L$0). Default hyperparameters include $L$1 and token count $L$2 for CIFAR-10, with $L$3 refinement schedules such as $L$4 for CIFAR-10 and $L$5 for CelebA-HQ. Gradient detachment ("short-gradient" trick) after each outer loop except the last manages memory (Esmaeilzadeh et al., 14 May 2026).

Quantitative Performance

Model/Setting	Precision ↑	Recall ↑	FID ↓	Few-Shot Acc ↑
RS-IMLE, CIFAR-10 (Baseline)	0.853	0.738	5.69	-
RS-IMLE + RTM (H=16,L=1)	0.896	0.773	3.97	-
RS-IMLE, CelebA-HQ (Baseline)	0.924	0.491	15.43	-
RS-IMLE + RTM (H=16,L=2)	0.952	0.592	10.67	-
MoR, LLM, 167M params	-	-	-	43.1%
Vanilla LM, 315M params	-	-	-	42.3%

On image datasets, RTM consistently raises both precision and recall while lowering FID (Frechet Inception Distance) compared to StyleGAN2, StyleGAN2-ADA, and IMLE baselines—across tasks including CIFAR-10, CelebA-HQ, and few-shot benchmarks. In LLMs, MoR-based RTMs deliver improved perplexity and few-shot accuracy with fewer parameters and higher throughput (Esmaeilzadeh et al., 14 May 2026, Bae et al., 14 Jul 2025).

6. Theoretical Insights and Benefits of Recursion

Recursive parameter sharing in RTM introduces a structural inductive bias toward multi-stage, coarse-to-fine refinement. Early recursion steps focus on global features (pose, composition), with later steps dedicated to fine-grained details (texture, localized variation). Parameter sharing across cycles regularizes the mapping function, counteracting propensity for mode collapse and preventing memorization of a limited set of mappings. Theoretical arguments demonstrate that RTM remains a continuous $L$6 map, inheriting the mode-coverage guarantees of IMLE (Esmaeilzadeh et al., 14 May 2026).

Empirical analysis shows 5–20% improvements in recall over single-pass MLP mappers, better nearest-neighbour preservation of distinctive image attributes, and simultaneous boosts in both diversity (recall) and fidelity (precision), unlike flow-matching or diffusion that often trade-off these metrics (Esmaeilzadeh et al., 14 May 2026).

A plausible implication is that RTM-like recursion mechanisms, when combined with adaptive token-level depth, could generalize further to diverse architectures and modalities, leveraging compute/parameter efficiency and structured inductive bias beyond image synthesis.

7. Extensions, Variants, and Cross-Domain Applications

The Mixture-of-Recursions framework consolidates token-level adaptive depth, shared parameter recursion, and router-based early exiting for efficiency and flexibility in LLMs. Potential extensions include multi-headed routers predicting both recursion depth and "width" (routing to different blocks), continuous (soft) recursion where final states aggregate information from all depths, joint cross-modal recursion, and hybrid schemes combining vertical (depth) and horizontal (sequence length) routing (Bae et al., 14 Jul 2025).

This suggests that RTM concepts serve as a unifying abstraction for recursive, parameter-efficient, and adaptive mapping modules across generative and sequential neural architectures, supporting state-of-the-art trade-offs in fidelity, diversity, efficiency, and throughput.

---

References:

• "One Pass Is Not Enough: Recursive Latent Refinement for Generative Models" (Esmaeilzadeh et al., 14 May 2026)
• "Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation" (Bae et al., 14 Jul 2025)