---
title: Token Merge with Attention (ToMA)
url: https://www.emergentmind.com/topics/token-merge-with-attention-toma
type: topic
---

# Token Merge with Attention (ToMA)

Search arXiv for 2411.07132 and 2509.10918 and related ToMeSD/ToFu semantic binding/token merging diffusion papers.
Token Merge with Attention (ToMA) denotes a token-aggregation paradigm for diffusion-based image generation that appears in two distinct but related formulations in the recent literature. In one formulation, ToMA is a reinterpretation of Token Merging (ToMe) for training-free semantic binding in text-to-image synthesis: subject, attribute, and sub-object tokens are merged into composite text tokens so that they share a single cross-attention map, with auxiliary unsupervised losses refining those composites during early denoising [2411.07132]. In a second formulation, ToMA is an efficiency method for diffusion transformers: token reduction is recast as submodular selection plus attention-like merge/unmerge linear maps, with GPU-aligned matrix operations and pattern reuse across layers and timesteps to reduce generation latency while preserving image quality [2509.10918]. The shared premise is that carefully structured token aggregation can alter the behavior of diffusion models without retraining, but the optimization target differs fundamentally between semantic correctness and computational efficiency.

## 1. Terminological scope and conceptual split

The term ToMA is not used for a single canonical algorithm across the supplied literature. One usage arises from a reinterpretation of ToMe as a “Token Merge with Attention” framework for semantic binding in text-to-image synthesis, while the later paper titled “ToMA: Token Merge with Attention for Image Generation with Diffusion Models” uses the same acronym for a latency-oriented token reduction method [2411.07132; 2509.10918]. This distinction is substantive rather than merely terminological.

| Usage | Primary objective | Core mechanism |
|---|---|---|
| ToMe reinterpreted as ToMA | Semantic binding in T2I | Merge subject+attribute text tokens into composite tokens; apply end-token substitution and two unsupervised losses |
| ToMA (2025) | Efficiency in diffusion models | Select diverse tokens by submodular optimization; merge/unmerge via dense matrix operations; reuse patterns across layers and timesteps |

In the semantic-binding setting, token merging acts on CLIP-tokenized text embeddings and modifies cross-attention conditioning. In the efficiency setting, token merging acts on transformer token features inside the diffusion model and reduces effective attention workload. A plausible implication is that both methods exploit redundancy in tokenized representations, but they do so for different failure modes: semantic misbinding in one case and quadratic attention complexity in the other.

## 2. Semantic-binding ToMA: prompt structure and composite token formation

In the semantic-binding formulation, the prompt after CLIP tokenization is written as
$$
P = [\mathrm{SOT}, w_1, w_2, \dots, w_L, \mathrm{EOT}, \dots, \mathrm{EOT}],
$$
with embeddings
$$
C = [c_0^{\mathrm{SOT}}, c_1, \dots, c_{M-1}].
$$
The prompt is assumed to describe $K$ main entities $n^1, n^2, \dots, n^K$, each associated with a set of attribute or sub-object tokens
$$
a^k = \{t_{k,1}, \dots, t_{k,r_k}\}.
$$
For each entity, the index set $S_k \subset \{1,\dots,M-1\}$ denotes the tokens corresponding to the $k$-th noun and its attributes. The supplied example is “a cat wearing sunglasses and a dog with a hat,” with $n^1=\text{“cat”}$, $a^1=\{\text{“wearing”},\text{“sunglasses”}\}$, $n^2=\text{“dog”}$, and $a^2=\{\text{“with”},\text{“a”},\text{“hat”}\}$ [2411.07132].

To bind each subject $n^k$ with its attributes $a^k$, ToMA introduces a composite token $t_c^{(k)}$ by merging the embeddings in $S_k$:
$$
\hat c_k = \sum_{i\in S_k} w_i c_i, \qquad w_i \ge 0,\quad \sum_{i\in S_k} w_i = 1.
$$
In the simplest case,
$$
\hat c_k = \frac{1}{|S_k|}\sum_{i\in S_k} c_i.
$$
All original tokens in $S_k$ are then removed and replaced by the composite embedding sequence
$$
c' = [\hat c_1, \dots, \hat c_K].
$$

The explicit motivation is semantic binding, defined as associating a given object with its attribute, termed attribute binding, or linking it to related sub-objects, termed object binding. The method is training-free and avoids full-model fine-tuning or externally specified layouts. Because the merged token stands in for an entire subject-attribute group, the method is designed to prevent the separate tokens from competing for inconsistent spatial assignments during generation.

## 3. Cross-attention integration and end-token substitution

In a standard diffusion U-Net at timestep $t$, the flattened feature queries $Q \in \mathbb{R}^{N\times d}$ attend to text keys
$$
K = C W_K^\top \in \mathbb{R}^{M\times d}
$$
through
$$
A = \operatorname{softmax}(QK^\top/\sqrt{d}),
$$
where $A \in \mathbb{R}^{N\times M}$. After merging, the textual key matrix becomes
$$
K' = [W_K \hat c_1;\dots;W_K \hat c_K].
$$
Since
$$
W_K \hat c_k = W_K \sum_{i\in S_k} w_i c_i = \sum_{i\in S_k} w_i (W_K c_i),
$$
the new attention scores satisfy
$$
Q(W_K \hat c_k)^\top
= \sum_{i\in S_k} w_i \bigl(Q(W_K c_i)^\top\bigr).
$$
After the softmax, the composite attention map $A'_{:,k}$ is therefore a weighted mixture of the original maps $A_{:,i}$ for $i\in S_k$ [2411.07132].

The stated consequence is that the newly introduced token $\hat c_k$ “inherits” and consolidates the spatial focus of all constituent tokens, guaranteeing that the bound subject, attributes, and sub-objects now share a single cross-attention map. Within the paper’s framing, this is the mechanism by which semantic binding is enforced without retraining.

A complementary mechanism is End-Token Substitution (ETS). The motivation is that CLIP’s $[\mathrm{EOT}]$ tokens often accumulate global semantics and can re-introduce stray attributes. To remove this spurious coupling, the method replaces the original end-token embeddings with embeddings drawn from a “plain” prompt containing only the subjects $\{n^1,\dots,n^K\}$:
$$
C^*_{\mathrm{EOT}}[j] = c^{\mathrm{EOT}}_{\text{subjects-only}}
\qquad \text{for all } j \ge L+1.
$$
The final textual condition is then
$$
C' = [c_0^{\mathrm{SOT}}, c_1,\dots,c_L^{\dots}, \hat c_1,\dots,\hat c_K, c^{\mathrm{EOT}*},\dots,c^{\mathrm{EOT}*}].
$$
This isolates the subject set from residual prompt-wide semantics carried by repeated end tokens.

## 4. Training-free refinement, algorithmic workflow, and reported performance

During the first $T_{\mathrm{opt}}$ timesteps of denoising, given as “e.g. 20% of $T$ total,” the semantic-binding ToMA updates only the composite token embeddings via gradients on two unsupervised losses. The entropy regularization term encourages each composite token’s attention map to concentrate on a small region:
$$
\mathcal{L}_{\mathrm{ent}}
= \sum_{k=1}^K H(A'_{:,k})
= \sum_{k=1}^K \sum_{p\in A'_{:,k}} -\,p\log p.
$$
The semantic binding loss enforces that each $\hat c_k$ leads to the same denoising vector as the full prompt $C'$:
$$
\mathcal{L}_{\mathrm{sem}}
= \sum_{k=1}^K \left\| \epsilon_\theta(z_t,\hat c_k,t) - \epsilon_\theta(z_t,C',t) \right\|_2^2.
$$
The combined objective is
$$
\mathcal{L} = \mathcal{L}_{\mathrm{ent}} + \lambda \,\mathcal{L}_{\mathrm{sem}}.
$$
A small number of gradient steps is taken on $\{\hat c_k\}$ before diffusion continues. No network weights are changed; only composite tokens are refined [2411.07132].

The pseudocode given in the summary is sequential. It tokenizes the prompt, parses it into $K$ subject-plus-attribute groups $S_1,\dots,S_K$, initializes each composite token by uniform averaging over its group, substitutes the original tokens with $\hat c_k$, replaces post-terminal $[\mathrm{EOT}]$ embeddings with $\mathrm{EOT}^*$ from a subjects-only prompt, forms the updated text matrix $C'$, initializes $z_T \sim \mathcal{N}(0,I)$, and then denoises from $t=T$ to $1$. For $t > T_{\mathrm{opt}}$, denoising proceeds normally; otherwise the method extracts $Q$ and $K'$, builds $A'$, computes $\mathcal{L}_{\mathrm{ent}}$ and $\mathcal{L}_{\mathrm{sem}}$, updates $\{\hat c_k\}$, re-assembles $C'$, and denoises with the updated condition.

The reported evaluation uses T2I-CompBench, specified as “300 prompts each for color, texture, shape attributes, measured by BLIP-VQA”; a GPT-4o object binding benchmark with “50 hard prompts of the form ‘an X with itemA and a Y with itemB,’ scored by GPT-4o”; and human-preference via ImageReward plus DetScore. The metrics listed are BLIP-VQA, GPT-4o binding accuracy, and human-preference. For base SDXL, the reported BLIP-VQA values are Baseline SDXL $\{0.6369, 0.5637, 0.5408\}$, Best prior (CoMat) $\{0.7774, 0.6591, 0.5262\}$, and ToMA (Ours) $\{0.7656, 0.6894, 0.6051\}$. The GPT-4o score is reported as $0.9549$ for ToMA versus $0.6458$ for SynGen, and human preference/ImageReward logits are reported to rise by $+0.2$ overall. Qualitative examples are described as showing perfect attribute placement, such as “cat wears glasses ○, dog wears hat ○,” with no cross-object leakage, including in multi-object, multi-attribute scenes. The summary concludes that the method achieves state-of-the-art semantic binding on both compositional benchmarks and open-ended GPT-4o object-binding tests.

## 5. Efficiency-oriented ToMA: submodular token selection and attention-like merge/unmerge

The 2025 ToMA paper addresses a different problem: scalability limits in diffusion models caused by transformers’ quadratic attention complexity. Its starting point is a set of token features $x_i \in \mathbb{R}^d$ over a ground set
$$
U = \{1,2,\dots,N\}.
$$
The goal is to select a smaller subset
$$
S \subseteq U,\qquad |S| = M,
$$
or equivalently impose a budget $M=\lfloor (1-r)N\rfloor$ for reduction ratio $r\in(0,1)$. Selection is cast as maximizing a facility-location style submodular objective
$$
f(S)=\sum_{i\in U}\max_{j\in S}\kappa(x_i,x_j),
\qquad
\kappa(x_i,x_j)=\frac{x_i^\top x_j}{\|x_i\|\,\|x_j\|},
$$
with cosine similarity reported as giving the best empirical diversity, though an RBF or inner-product can also be used [2509.10918].

The paper adopts the classical greedy algorithm for monotone submodular maximization under a cardinality constraint:
initialize $S\leftarrow \emptyset$, then repeatedly choose
$$
j^*=\arg\max_{j\in U\setminus S}\Delta f(j\mid S),
$$
where
$$
\Delta f(j\mid S)=f(S\cup\{j\})-f(S).
$$
The summary states the standard guarantee from Nemhauser et al. (1978): a $(1-1/e)\approx 0.632$ approximation to the optimal $f(S)$. Marginal-gain updates are accelerated by caching $\max_{j\in S}\kappa(x_i,x_j)$ for every $i\in U$, reducing each update to $O(Nd)$ time rather than recomputing from scratch.

After selection, each original token is assigned to its closest center,
$$
c(i)=\arg\max_{j\in S}\kappa(x_i,x_j),
$$
which induces clusters
$$
C_j=\{\,i\in U \mid c(i)=j\}.
$$
Merge is implemented by a dense matrix $W_{\mathrm{merge}}\in\mathbb{R}^{M\times N}$ with entries
$$
(W_{\mathrm{merge}})_{p,i}
=
\begin{cases}
\tfrac{1}{|C_{j_p}|}, & \text{if } c(i)=j_p,\\
0, & \text{otherwise.}
\end{cases}
$$
For stacked token features $X\in\mathbb{R}^{N\times d}$, the merged representation is
$$
X' = W_{\mathrm{merge}}X \in \mathbb{R}^{M\times d}.
$$
Standard multihead attention is then applied on $X'$; for one head,
$$
Q' = X'W^Q,\qquad K' = X'W^K,\qquad V' = X'W^V,
$$
followed by
$$
\mathrm{Attn}(Q',K',V')
=
\mathrm{softmax}\!\left(\frac{Q'K'^\top}{\sqrt{d_h}}\right)V'.
$$
Unmerge is implemented via
$$
W_{\mathrm{unmerge}} = W_{\mathrm{merge}}^\top \in \mathbb{R}^{N\times M},
$$
with
$$
Z = W_{\mathrm{unmerge}} Z' \in \mathbb{R}^{N\times d}.
$$
Operationally, each original token copies the merged result of its cluster center.

The GPU-efficiency claim rests on representation rather than on token reduction alone. Merge and unmerge are both single dense GEMMs, $W_{\mathrm{merge}}X$ and $W_{\mathrm{unmerge}}Z'$, which the summary states can be fused into highly optimized kernels by cuBLAS and cuBLASLt. The paper explicitly contrasts this with ToMeSD and ToFu, which rely on sorting and scatter/gather-like operations.

## 6. Pattern reuse, complexity, empirical outcomes, and interpretive cautions

The efficiency-oriented ToMA further exploits latent locality and sequential redundancy. Within each transformer block, the assignment $c(i)$ is said to change little from layer to layer, so $W_{\mathrm{merge}}$ and $W_{\mathrm{unmerge}}$ can be reused across heads or entire blocks and recomputed only when the marginal gain of adding a new center exceeds a threshold $\tau$. Across diffusion timesteps, the last timestep’s clustering is used as the initial seed for the next, providing a warm start and often avoiding any change. The summary describes this “pattern reuse” as making the expensive greedy selection, with cost roughly $O(NMd)$ per update, occur “very infrequently—typically once every 4–8 layers” [2509.10918].

The complexity accounting is explicit. Baseline full attention per layer is given as approximately $4N^2d$ FLOPs. For ToMA, attention on merged tokens costs $4M^2d$, while the two GEMMs contribute approximately $2MNd$ FLOPs, yielding
$$
4M^2d + 2MNd
=
4(1-r)^2N^2d + 2(1-r)N^2d.
$$
For the typical value $r=0.2$, the summary computes
$$
4\cdot 0.64\,N^2d + 2\cdot 0.8\,N^2d
= 2.56N^2d + 1.6N^2d
= 4.16N^2d,
$$
described as about a $4.16/4.0 = 1.04\times$ baseline FLOPs. The paper’s central argument is that theoretical FLOPs alone are not decisive: the merged attention is on $M\times M$ rather than $N\times N$, the GEMMs are hardware-friendly, and submodular oracle calls are rare. The reported end-to-end GPU measurements with FlashAttention are that ToFu achieves only $5$–$8\%$ latency reduction, whereas ToMA “consistently achieves 20–25% end-to-end speedup vs full FlashAttention.”

The headline empirical results are more specific. On Stable Diffusion XL (SDXL) 1.0, ToMA reduces single-image generation latency by $24\%$; on Flux, it yields a $23\%$ speedup. Across $1\,000$ ImageNet prompts, the reported DINO-CLIP similarity shift satisfies $\Delta_{\mathrm{DINO}}<0.07$ relative to the full-token baseline, characterized as visually indistinguishable. Ablations over reduction ratio $r\in\{0.1,0.2,0.3\}$ report that quality is essentially flat up to $r=0.25$, then degrades gracefully. Reclustering every layer versus every $4$ layers versus every $8$ layers is reported to show that $4$–$8$ gives $20$–$23\%$ speedup with no measurable quality loss. Visual examples are said to show that even at $30\%$ token reduction, fine textures such as hair strands and foliage remain crisp, with only imperceptible blur in uniform backgrounds.

A recurrent source of confusion is to treat the two ToMA formulations as interchangeable. That would be inaccurate. The semantic-binding variant merges text tokens to enforce subject-attribute coherence and uses entropy plus semantic binding losses during early denoising. The efficiency variant merges internal transformer tokens to reduce latency and uses submodular selection, dense merge/unmerge maps, and reuse heuristics. This suggests a broader design space around token aggregation in diffusion models, but the supplied papers do not claim that the semantic-binding objective and the efficiency objective are the same problem, nor that one method subsumes the other.

Source: https://www.emergentmind.com/topics/token-merge-with-attention-toma