Trait-Routing Attention in Controllable Diffusion

• Trait-Routing Attention (TA) is a token-level, layer-adaptive MoE mechanism that routes heterogeneous attributes like texture, shape, and color into diffusion models.
• It employs cross-attention and top-2 expert routing within each U-Net self-attention block to fuse condition features and achieve fine-grained controllability.
• Empirical studies show that TA reduces FID and improves CLIP-I, SSIM, and LPIPS metrics by mitigating attribute entanglement compared to static feature fusion.

Trait-Routing Attention (TA) is a token-level, layer-adaptive mixture-of-experts (MoE) mechanism designed to enable dynamic, disentangled injection of heterogeneous condition features—such as texture, shape, and color—into diffusion models for controllable image synthesis. Developed as the core architectural innovation in VersaVogue’s U-Net, TA addresses the pervasive problems of attribute entanglement and semantic interference inherent in prior approaches that relied on feature concatenation or static layer-wise injections. TA enables each condition feature to be routed to semantically specialized experts at distinct generative layers, resulting in significant improvements in visual fidelity, semantic consistency, and fine-grained controllability in applications such as unified garment generation and virtual dressing (Yu et al., 8 Apr 2026).

1. Integration within Diffusion U-Nets

TA is integrated into the SDXL-based denoising U-Net architecture by situating a dedicated TA module within each self-attention block, after the standard self-attention on latent features $Z^{(\ell)}\in\mathbb{R}^{m\times d}$ at layer $\ell$. Multiple heterogeneous condition streams, denoted $\{F_i^{(\ell)}\}_{i=1}^N$, are extracted by a parallel fitting U-Net, where each $F_i$ encodes a distinct attribute (e.g., shape mask, texture exemplar, color hint) using isolated self-attention branches. These feature maps are fed layer-by-layer into corresponding TA modules. Each TA module fuses the condition features with the current latent representation and injects the output back into the U-Net’s residual processing stream. This architectural placement allows TA to mediate the integration of multi-source features at every generative layer, offering precise control over the spatial and semantic flow of information (Yu et al., 8 Apr 2026).

2. Mathematical Formulation

At every layer $\ell$, the input latent $Z\in\mathbb{R}^{m\times d}$ and a set of $N$ condition features $F_i\in\mathbb{R}^{m\times d}$ are processed as follows:

1. Cross-Attention per Condition: For each stream $i$:
   • $Q = Z W_q$,
   • $\ell$0,
   • $\ell$1, where $\ell$2 are frozen SDXL projections, and $\ell$3 are LoRA residuals. Attribute-aware interactions are computed:

$\ell$4

1. Token-Level Mixture-of-Experts Routing: $\ell$5 experts (empirically $\ell$6) are instantiated per layer. For each token $\ell$7:

$\ell$8

with $\ell$9. The top-$\{F_i^{(\ell)}\}_{i=1}^N$0 ($\{F_i^{(\ell)}\}_{i=1}^N$1) indices $\{F_i^{(\ell)}\}_{i=1}^N$2 are selected, and gating weights are normalized:

$\{F_i^{(\ell)}\}_{i=1}^N$3

2. Expert Output Aggregation: Each selected expert $\{F_i^{(\ell)}\}_{i=1}^N$4 computes $\{F_i^{(\ell)}\}_{i=1}^N$5, and the outputs are aggregated:

$\{F_i^{(\ell)}\}_{i=1}^N$6

yielding $\{F_i^{(\ell)}\}_{i=1}^N$7.

3. Final Residual Fusion: All $\{F_i^{(\ell)}\}_{i=1}^N$8 refined features are summed with the output of self-attention on $\{F_i^{(\ell)}\}_{i=1}^N$9:

$F_i$0

($F_i$1, $F_i$2). This comprehensive fusion enables expert-specialized updates per token and condition, facilitating semantically aligned synthesis.

3. Mixture-of-Experts Routing Scheme

TA employs a per-layer, per-token dynamic routing process, allowing each token in every attribute stream to select its own subset of experts. Each expert $F_i$3 is a two-layer MLP with hidden width approximately $F_i$4. The gating mechanism comprises a linear weight $F_i$5 and an auxiliary “noise branch” $F_i$6; random Gaussian noise is injected to encourage routing diversity and mitigate expert collapse. No explicit load-balancing regularization is used. Top-2 routing was empirically determined as optimal for balancing attribute diversity handling with parameter efficiency; more densely activated routing (top-3 or full MoE) introduced irrelevant noise and degraded performance on metrics such as FID and CLIP-I (Yu et al., 8 Apr 2026).

MoE Hyperparameters	Value	Empirical Role
Number of experts	4	Captures attribute diversity
Top-$F_i$7 routing	2	Sparsifies selection
Expert architecture	2-layer MLP, width~$F_i$8	Attribute specialization

4. Disentangled Layerwise Attribute Injection

Trait-Routing Attention adaptively decomposes attribute injection across the U-Net hierarchy:

• Early layers: Coarse, low-frequency generative layers route tokens representing structural silhouettes to shape-specialized experts.
• Middle layers: Attributes coding for fine-grained texture and pattern are preferentially routed to texture experts.
• Late layers: Features relevant for color harmonization and high-frequency details are routed to color and lighting experts.

Empirical analysis of routing entropy across layers confirms that TA autonomously learns a layerwise attribute-specific injection schedule, without any handcrafted heuristics or losses for explicit disentanglement. This suggests that layer-adaptive expert selection emerges as a consequence of end-to-end learning under TA’s design (Yu et al., 8 Apr 2026).

5. Training, Optimization, and Preference Alignment

The TA module components (LoRA residuals, gating networks, and expert MLPs) are trained in two stages:

• Stage 1 (standard): Only the aforementioned components are optimized under the canonical diffusion mean-squared error loss,

$F_i$9

No disentanglement or attribute supervision losses are introduced.

• Stage 2 (preference alignment): After convergence, the base model is frozen and “winner/loser” preference pairs are autonomously mined via a multi-perspective evaluator (evaluating content fidelity, textual alignment, perceptual quality). Gating and expert parameters are further refined by direct preference optimization (DPO),

$\ell$0

with

$\ell$1

and similarly for $\ell$2.

A plausible implication is that this two-stage objective sharpens both visual realism and attribute controllability without requiring additional manual disentanglement signals.

6. Empirical Analysis and Ablations

Extensive ablation studies show that TA confers substantial quantitative and qualitative improvements over concatenation and dense MLP fusion approaches:

Model Variant	FID	CLIP-I	SSIM	LPIPS
Baseline (no TA, no MPO)	51.22	0.777	0.754	0.310
+ TA only	45.32	0.805	0.802	0.244
+ TA + MPO	44.74	0.809	0.811	0.225

Removal of TA produces severe attribute entanglement, such as blurred logos, color bleeding, and mismatched silhouettes, while TA enables crisp pattern placement, faithful logos, and coherent shading. A key observation is that increasing the number of experts or the routing density beyond $\ell$3, top-2 selection, degrades performance due to routing noise. This supports the design choice for a lightweight, sparsely routed MoE (Yu et al., 8 Apr 2026).

7. Impact and Scope for Controllable Image Synthesis

Trait-Routing Attention represents an enabling mechanism for multi-condition, high-fidelity, and semantically consistent image synthesis in the fashion domain. By delivering dynamic, token-level, expert-specialized feature routing at every generative layer, TA supports a unified workflow for garment generation and virtual dressing, resolving limitations of previous approaches stemming from static fusion methods. Empirical results demonstrate that this routing-based disentanglement is critical for resolving conflicting condition cues and producing visually coherent outputs under heterogeneous multi-source control (Yu et al., 8 Apr 2026).

This methodology can be plausibly extended to other settings where attribute disentanglement and combinatorial control over image generation are necessary. The integration of preference-aligned optimization further highlights TA’s capacity for scalable, annotation-free improvement in synthesis realism and alignment.