Local Representative Token Guided Merging (ReToM)

• Local Representative Token Guided Merging (ReToM) is an adaptive token merging technique that partitions tokens into local windows to reduce the quadratic cost of self-attention.
• It employs cosine similarity to select a representative token per window, ensuring that salient visual features are preserved during the merge.
• Empirical results on Stable Diffusion demonstrate improved FID and CLIP scores with efficient near-linear scaling in attention computations.

Local Representative Token Guided Merging (ReToM) is a token merging strategy designed to accelerate attention-based text-to-image generation models, specifically targeting the computational bottleneck in the self-attention layers of architectures such as Stable Diffusion’s U-Net. ReToM introduces local adaptive windowing, representative token selection grounded in window-level similarity dynamics, and similarity caching, aiming to reduce the quadratic cost of attention while preserving visual fidelity and salient feature detail. Experimentally, ReToM improves both FID and CLIP scores over prior merging methodologies with comparable inference times and without requiring additional training (Lee et al., 17 Jul 2025).

1. Motivation and Limitations of Prior Work

The computational inefficiency of self-attention in U-Net architectures—a core component in models like Stable Diffusion—is due to the $O(N^2)$ scaling with respect to the number of input tokens $N$. For high-resolution image synthesis, this cost results in slow generation and heavy memory consumption. Token merging strategies, such as ToMeSD [Bolya & Hoffman ’23], attempt to reduce this cost by grouping similar tokens for attention, but these methods exhibit several limitations:

• Fixed windowing: Prior approaches commonly employ a fixed region size throughout all attention layers, neglecting the changing receptive field across the U-Net’s depth.
• Suboptimal token destination selection: Merging source tokens into destination tokens is typically performed via random selection or bipartite matching, risking the loss of salient local structures.
• Computational redundancy: Similarity computations are performed in every diffusion timestep, maintaining a significant quadratic complexity burden.

ReToM addresses these constraints by employing adaptive local windows, selection of a single representative token per window based on similarity, and cached similarity updates at reduced temporal frequency (Lee et al., 17 Jul 2025).

2. Local Window Partitioning Formalism

At any given attention layer $\ell$ in the U-Net Transformer, denote the sequence of tokens as $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$. ReToM partitions the index set $\{1,\ldots,N\}$ into $K$ non-overlapping windows:

$$W^1, W^2, ..., W^K \qquad \bigcup_{k=1}^K W^k = \{1, ..., N\}, \qquad W^k \cap W^{k'} = \emptyset ~\text{for}~ k \neq k'$$

Each window $W^k$ contains $N_{W^k}$ tokens, with window size $m_\ell$ adaptively chosen per layer.

When tokens map to an $N$0 grid (common in vision Transformer applications), the window $N$1 is constructed as:

$N$2

where $N$3 and indices traverse to exhaustively and non-overlappingly cover all positions.

3. Representative Token Selection within Windows

Within each local window $N$4, ReToM performs all-pairs cosine similarity analysis and, for each token $N$5 in $N$6, computes its average similarity to all other tokens:

$N$7

$N$8

The representative token $N$9 is the one maximizing the above similarity criterion:

$\ell$0

All tokens in $\ell$1 are deemed possible sources for merging into $\ell$2, concentrating window-level context into locally salient features.

4. Merging and Processing Algorithm

Merging is executed only every $\ell$3 diffusion timesteps, leveraging cached window-level similarities for computational amortization.

For each window $\ell$4 at a merge timestep $\ell$5:

• Compute or load (if $\ell$6) $\ell$7 for all $\ell$8.
• Select $\ell$9, the representative token as per the maximal average similarity.
• Determine the number of source tokens, $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$0, using a merge ratio $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$1.
• Select the top $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$2 non-representative tokens in $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$3 by ranking $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$4.
• The merged token is computed as:

$$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$5

where $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$6 is a balancing hyperparameter and $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$7 are ranked source tokens.

• $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$8 is replaced by $$X = [x_{1}, ..., x_{N}]^\top \in \mathbb{R}^{N \times D}$$9 for window $\{1,\ldots,N\}$0.
• The modified token set is then used in conventional multi-head self-attention, drastically reducing per-window compute.

5. Complexity and Efficiency Analysis

The baseline U-Net self-attention at each layer requires $\{1,\ldots,N\}$1 operations. By constraining self-attention to merged tokens per window (size $\{1,\ldots,N\}$2 post-merge), per-window cost drops from $\{1,\ldots,N\}$3 to $\{1,\ldots,N\}$4. The major overhead becomes the computation of similarities for determining representative tokens, yet this is amortized across $\{1,\ldots,N\}$5 timesteps.

Total per-layer amortized cost becomes:

$\{1,\ldots,N\}$6

where $\{1,\ldots,N\}$7 is average window size at layer $\{1,\ldots,N\}$8. Setting $\{1,\ldots,N\}$9 sufficiently large and keeping windows small in early/late layers brings near-linear scaling for attention on most timesteps while accommodating context in deeper layers with larger windows.

6. Empirical Results and Ablations

All experiments were performed on Stable Diffusion with 50 sampling steps on 512×512 ImageNet images (guidance=7.5). The following summarizes the key quantitative results:

Method	FID	CLIP	Speed (s/im)
Baseline SD	37.02	38.10	2.62
ToMeSD (fixed 2, rand)	37.20	36.00	2.10
ReToM (fixed 2, rep)	35.00	37.90	2.25
ReToM (adaptive + rep)	34.89	39.40	2.17

• ReToM with adaptive windows and representative selection achieves FID = 34.89 (–5.8% vs. baseline), CLIP = 39.40 (+1.3), and 2.17 s/im inference speed.
• Excessively large fixed windows (size 16) degrade FID to 36.53, demonstrating the need for adaptive window sizing.
• Substituting the selection of the “most representative” with the “least representative” token raises FID and lowers SSIM (0.72→0.74), confirming the necessity of the optimal average similarity criterion.
• Visualizations show cosine similarities within windows remain above 0.9 over $K$0 steps, supporting the temporal caching approach.

7. Summary and Significance

ReToM institutes a scalable, adaptive local merging paradigm for U-Net Transformers, distinguished by local window partitioning per receptive field, maximally representative token selection via average similarity, top-r token convex merging, and efficient similarity caching. This methodology yields improved generation metrics and substantial reduction in quadratic attention cost over prior approaches—accomplishing efficiency and fidelity gains without supplementary training phases or fine-tuning (Lee et al., 17 Jul 2025).