Technical Summary: InstructMoLE

Problem Statement

Parameter-Efficient Fine-Tuning (PEFT) of Diffusion Transformers (DiTs) for multi-conditional image generation (e.g., multi-subject composition, image editing) faces significant challenges when using monolithic adapters like LoRA, primarily due to task interference and catastrophic forgetting. While Mixture-of-Experts (MoE) architectures offer a modular solution by diversifying experts for different tasks, their application to diffusion models has been limited by routing policies inherited from language models.

Current approaches typically employ token-level routing, where each image patch independently selects an expert based on local hidden states. The paper argues that this local decision-making process is fundamentally misaligned with the global nature of user instructions. In complex generation tasks (e.g., "a baby crawling on grass, a white horse nearby, and a football helmet"), instructions establish global semantic relationships. Token-level routing fails to maintain these relationships, leading to:

1. Spatial Fragmentation: Region-wise discontinuities in style and hue.
2. Routing Jitter: Amplification of high-frequency noise.
3. Semantic Drift: Failure to adhere to the holistic intent of the instruction, resulting in incoherent subjects or ignored attributes.

Methodology

The authors propose InstructMoLE, a framework that replaces token-level routing with a globally consistent, instruction-guided policy. The architecture is built upon the Flux.1 Kontext backbone and utilizes Low-Rank Adaptation (LoRA) experts.

1. Instruction-Guided Routing (IGR)

Instead of selecting experts per token, InstructMoLE employs an instance-level routing policy.

• Global Signal Distillation ($Z_{global}$): A global routing signal is distilled from the user's textual instruction ($I_c$). This signal fuses two representations:
  • CLIP Embedding: Provides a robust, holistic semantic anchor.
  • Perceiver-style Distillation: A learnable latent query iteratively attends to token-level features from a T5 encoder. This captures fine-grained compositional details (e.g., specific color-subject associations) that a simple pooled CLIP embedding might average out.
• Unified Expert Council: At each layer of the model, the router computes a single set of gating weights based on $Z_{global}$. This "expert council" (a Top-$k$ set of experts) is then broadcast identically to all spatial tokens within that layer. This ensures that the entire image generation process at a specific abstraction level is governed by a coherent set of experts aligned with the user's intent.

2. Output-Space Orthogonality Loss

To prevent expert collapse (where distinct experts learn redundant functions), the authors introduce a novel regularizer.

• Standard load-balancing losses only ensure uniform utilization frequency but do not guarantee functional diversity.
• The Orthogonality Loss ($L_{ortho}$) operates on the pre-gating outputs of all experts. It calculates the squared cosine similarity between the flattened output vectors of every pair of experts.
• By minimizing this loss, the model is forced to learn mutually orthogonal functional roles for each expert, ensuring that the selected council offers a diverse set of capabilities rather than redundant solutions.

3. Training Objectives

The total loss function combines three components:
$$L_{total} = L_{flow} + \lambda_{aux}L_{aux} + \lambda_{ortho}L_{ortho}$$

• $L_{flow}$: Standard flow-matching/diffusion loss for generation accuracy.
• $L_{aux}$: Auxiliary load-balancing loss to prevent over-reliance on a few experts.
• $L_{ortho}$: The proposed orthogonality loss for functional diversity.

Key Contributions

1. Identification of Routing Misalignment: The paper identifies the critical mismatch between local token-level routing and global instruction semantics in multi-conditional image generation.
2. Instruction-Guided Routing (IGR): A novel framework that conditions expert selection entirely on a global signal derived from the user's instruction, ensuring structural and semantic consistency across the image.
3. Orthogonality Regularization: A new loss function that explicitly promotes functional diversity among experts, mitigating representational collapse more effectively than load balancing alone.
4. Empirical Validation: Extensive experiments demonstrating that InstructMoLE outperforms existing LoRA adapters and MoE variants (including token-level routing baselines) on challenging multi-conditional benchmarks.

Experimental Results

The model was fine-tuned on the Flux.1 Kontext backbone using a diverse dataset of synthesized and public data (including face swapping, style transfer, multi-subject composition, and spatial control).

• Multi-Subject Generation (XVerseBench): InstructMoLE achieved the highest average score (71.07), outperforming the strong Flux.1 Kontext baseline (68.04) and DreamO-v1.1 (69.57). It showed particular strength in identity preservation (ID-Sim) and object fidelity (IP-Sim).
• Single-Image Editing (GEdit-EN-full): The method achieved the highest average score (6.17) among compared models, surpassing OmniGen2 (5.45) by over 13% and the Flux.1 Kontext baseline (6.06). It excelled in tasks requiring both global and local manipulations.
• In-Context Generation (OmniContext): InstructMoLE achieved the highest average ID-Sim score (38.85) among released models, balancing instruction adherence and identity preservation better than OmniGen2 (which had high prompt following but poor identity preservation).
• Ablation Studies:
  • Routing Policy: IGR significantly outperformed token-level routing (Standard, Expert Choice, Expert Race) in spatial alignment tasks (Depth RMSE, Canny F1, Pose F1), confirming that global consistency is crucial for structural integrity.
  • Components: Removing the orthogonality loss or the Perceiver-style signal distillation resulted in measurable performance drops, validating the necessity of both functional diversity and fine-grained instruction encoding.

Significance and Claims

The paper claims that Instruction-Guided Routing (IGR) represents a more effective paradigm than traditional token-level routing for instruction-driven fine-tuning of generative models. By aligning expert selection with the holistic intent of the user, InstructMoLE enables:

• Superior Compositional Control: The ability to render complex scenes with multiple subjects and specific spatial relationships without fragmentation.
• Fidelity to User Intent: Improved adherence to complex, chain-of-thought instructions and preservation of identity in editing tasks.
• Robustness: The framework maintains high fidelity even when processing complex instructions, though the authors note limitations under severe textual corruption or extreme reference loads (binding >3 entities).

The work establishes a robust, generalizable framework for parameter-efficient adaptation of diffusion transformers, addressing the specific challenges of multi-conditional generation that monolithic adapters and local-routing MoE models fail to resolve.