PaperlayerView exact revision on arXiv

Source-linked AI summary

SAGE: Subspace Alignment for Classifier-Free Guidance in Mixture-of-Experts Diffusion Models

Boyu Zhang, Yangming Cheng, Ning Zhang, Pengfei Liu, Weijie Li, Yifan Gao, Hangyu Li, Litong Gong

arXiv:2609.34525v1cs.CV

TL;DR

In MoE diffusion models, independent routing can place CFG branches in different subspaces, creating a residual that guidance amplifies. SAGE aligns unconditional activations to the conditional subspace during training while preserving routing diversity. The authors report consistent gains, including a 9.3% improvement in peak DPG-Bench score for a 1B-parameter text-to-image model.

  • Problem

    Independent MoE routing can place CFG’s conditional and unconditional branches in different subspaces, causing guidance to amplify an unconditional residual.

    Source: §1 · p. 1

  • Method

    SAGE is a training-time regularizer that penalizes unconditional MoE activations outside the conditional subspace while preserving routing diversity and adding no inference cost.

    Source: §1 · p. 2

  • Results

    SAGE consistently outperforms baselines across CFG scales; on a 1B-parameter text-to-image model, it improves peak DPG-Bench score by 9.3%.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SAGE addresses the MoE-specific CFG mismatch without architectural changes, while preserving routing diversity and adding no inference cost.

    Source: §1 · p. 2

    The large-scale evaluation covers models only up to 1B parameters; scaling to 100B+ models remains future work.

    Source: §5 · p. 9

Abstract

from arXiv · show

Diffusion Transformers with Mixture-of-Experts (MoE) routing are a leading recipe for scaling generative models. Classifier-Free Guidance (CFG) is essential for generation quality, yet excessively high guidance scales trigger collapse. We identify a previously unreported failure mode in their combination: the two CFG branches route independently, so their realized activations occupy different subspaces. The unconditional write then leaves the conditional subspace, and CFG amplifies that residual linearly in the guidance scale. We propose SAGE, a training-time regularizer that aligns unconditional MoE activations to the conditional subspace without restricting routing diversity, at zero inference cost. Toy experiments show that SAGE dramatically suppresses extreme drift by 9.2x. When scaled to a 1B-parameter text-to-image model, SAGE significantly improves generation quality, delivering a 9.3% boost in peak DPG-Bench performance. Extensive experiments demonstrate that SAGE consistently outperforms the baseline.

1 INTRODUCTION

The introduction identifies a CFG failure specific to sparse MoE models: independently routed branches occupy different subspaces, so guidance amplifies an unconditional residual. SAGE aligns unconditional activations to the conditional subspace without restricting routing diversity.

1 INTRODUCTION

  • Forcing the two branches onto the same routes removes mismatch but eliminates expert specialization, and experiments find that restricting routing freedom degrades performance.Source: §1 · pp. 1–2
  • Routing misalignment places conditional and unconditional branches in different subspaces, and CFG amplifies the resulting leak linearly with guidance scale.The leak is specific to sparse MoE models and vanishes in dense architectures.Source: §1 · pp. 1–2
  • SAGE penalizes the unconditional output component outside the conditional subspace during training, preserving routing diversity without architectural changes or inference cost.The authors report that forcing hard routing agreement is markedly sub-optimal.Source: §1 · p. 2
  • SAGE consistently outperforms baseline models across CFG scales, while controlled toy experiments and ablations validate its mechanism and design choices.The proposed regularizer preserves routing diversity and adds no inference cost.Source: §1 · p. 2

2 RELATED WORK

Prior work scales diffusion transformers with sparse experts and develops several approaches to CFG, but the cited work leaves the interaction between MoE routing and CFG’s dual branches unstudied. SAGE addresses that mismatch through a training-time subspace constraint.

2.1 MIXTURE-OF-EXPERTS IN DIFFUSION MODELS

  • Prior diffusion MoE work explores different routing and scaling designs, but the cited studies do not examine how routing interacts with CFG’s conditional and unconditional passes.ProMoE assigns dedicated experts to the unconditional branch, but does not align remaining routed experts’ activations and sacrifices an expert’s conditional capacity.Source: §2.1 · p. 2

2.2 CLASSIFIER-FREE GUIDANCE

  • CFG research includes rescaling, schedules, guidance-term decompositions, manifold constraints, and alternative unconditional branches; SAGE instead addresses sparse-model subspace mismatch during training.The cited corrective methods target different guidance issues, while SAGE incurs zero inference cost.Source: §2.2 · pp. 2–3

3 SAGE: SUBSPACE ALIGNMENT FOR CLASSIFIER-FREE GUIDANCE

The section sets out flow matching, classifier-free guidance, and sparse MoE layers, then explains how independent routing creates an unconditional write outside the conditional subspace. SAGE penalizes that component while leaving routes free.

3.1 PRELIMINARIES

  • Classifier-free guidance combines conditional and null-condition velocity predictions as (1+s) vc −s vu at inference.Flow matching learns a conditional velocity field transporting Gaussian noise to the data distribution.Source: §3.1 · p. 3

3.2 ROUTING MISALIGNMENT AND SUBSPACE ALIGNMENT

  • Because conditioning changes hidden states, conditional and unconditional passes generally select different expert sets and can occupy different subspaces.The figure depicts this independent routing and the resulting mismatch in realized activations.Source: §3.2 · p. 3, Figure 1 · p. 4
  • The leak δ(s) is the unconditional write component outside the conditional subspace, and its norm grows exactly linearly with guidance scale s.It vanishes if and only if vu lies in VSc; exclusive experts are its MoE-specific source.Source: §3.2 · p. 4
  • SAGE penalizes only the unconditional write outside the conditional subspace; CFG remains in a common subspace without requiring shared routes or a global basis.The conditional subspace is estimated from stacked conditional activations using a sampled column subset and thin QR.Source: §3.2 · p. 4, Figure 1 · p. 4
  • Training adds a paired unconditional forward pass to form conditional and unconditional activations; standard CFG dropout supervision remains, and inference is unchanged.Source: §3.2 · p. 5

4 EXPERIMENTS

Experiments test SAGE first in a measurable 2-D toy setting, then at text-to-image scale, followed by ablations of its training choices. Across the large-scale benchmarks, SAGE improves quality and degrades more gracefully as CFG guidance rises.

4.1 TOY EXPERIMENT

  • At s=15, MoE’s maximum per-class excess drift reaches 12.9, while MoE + SAGE remains at 1.4, a 9.2× reduction.At s=7, MoE + SAGE also maintains mode purity of 0.781 versus 0.121 for MoE.Source: §4.1 · p. 6, Table 1 · p. 6

4.2.1 SETUP

  • The 1B-parameter text-to-image MoE-DiT evaluation compares Baseline, SAGE, and a KL routing constraint across DPG-Bench and GenEval and CFG scales.The model uses eight routed experts and one shared expert per layer, and all configurations train for 200,000 steps.Source: §4.2.1 · pp. 6–7

4.2.2 MAIN RESULTS

  • SAGE beats the Baseline at all s ≥3 on both benchmarks, with a 9.3% peak DPG-Bench gain and a +5.4 pp GenEval advantage at s=10.SAGE at s=7 also exceeds the Baseline’s DPG-Bench peak.Source: Table 2 · p. 7, §4.2.2 · pp. 7–8
  • The KL routing constraint underperforms on every metric, including a −21 pp GenEval drop versus Baseline at s=3; SAGE retains routing diversity.The authors attribute the constraint’s poor performance to loss of expert specialization.Source: §4.2.2 · p. 8
  • At s=10 and 15, Baseline images show over-saturation, white patches, and structural repetition, while SAGE preserves coherent composition and realistic textures.At moderate guidance (s ≤7), both methods produce plausible images.Source: §4.2.2 · p. 8, Figure 4 · p. 8

4.2.3 ABLATION STUDY

  • From-scratch SAGE performs better than post-hoc fine-tuning, and λ=1 is best among tested weights; intermittent activation retains most gains, while ProMoE remains below default SAGE.Increasing λ to 2 or 5 lowers all four benchmark scores below Baseline; removing the shared expert still leaves SAGE ahead on both benchmarks.Source: Table 3 · p. 9, §4.2.3 · p. 9

5 CONCLUSION AND FUTURE WORK

The paper identifies routing-induced subspace leakage in MoE diffusion models and proposes SAGE to align unconditional writes with the conditional subspace. SAGE improves CFG stability in experiments, while evaluation beyond 1B-parameter models and other settings remains future work.

5 CONCLUSION AND FUTURE WORK

  • SAGE improves CFG stability from controlled toy experiments to a 1B-parameter text-to-image model.It penalizes the unconditional write's component outside the conditional subspace during training, with no architectural change or inference cost.Source: §5 · p. 9
  • Routing misalignment lets unconditional writes leave the conditional subspace, and CFG amplifies that residual linearly with guidance scale.The paper identifies this as a subspace leak in MoE diffusion models.Source: §5 · p. 9
  • Evaluation reaches 1B parameters; scaling beyond 1B, deeper architectures, video generation, and more complex CFG settings remain future directions.The paper specifically mentions T2AV and R2V tasks that often use double CFG.Source: §5 · p. 9

A THEORETICAL ANALYSIS

The analysis formalizes how independent CFG routing produces a leak outside the realized conditional activation subspace, and how SAGE's loss relates to that leak. It also distinguishes exclusive and shared experts and clarifies the scope of the subspace and the guarantees.

A THEORETICAL ANALYSIS

  • The analysis represents CFG as combining conditional and unconditional layer writes, with shared experts included in the MoE output.The conditional and unconditional branches route independently.Source: A THEORETICAL ANALYSIS · p. 13, §A.1 · p. 13

A.1 INDEPENDENT ROUTING, DIFFERENT SUBSPACES

  • Independent routing can place the unconditional write outside the realized conditional subspace, producing a leak that is exactly linear in guidance scale.The analysis does not claim the realized subspace contains every conditional-batch activation when n > m.Source: §A.1 · p. 13

A.2 WHAT SAGE MINIMIZES

  • SAGE minimizes the relative energy of unconditional activations outside the conditional subspace, which controls leak energy at any parameter value.The analysis explicitly makes no Markov bound or λ^-1/2 rate claim; λ trades the SAGE loss against flow-matching loss during optimization.Source: §A.2 · p. 13

A.3 EXCLUSIVE EXPERTS

  • Exclusive unconditional experts provide a MoE-specific extra write, while dense FFNs have no exclusive routed experts.The exclusive-versus-shared routing split does not bound the residual outside the conditional subspace; dense CFG may still fail for other reasons.Source: §A.3 · p. 14

A.4 SHARED EXPERT

  • A shared expert does not ensure the unconditional write lies in the conditional subspace, because its branch activations differ and exclusive routed writes remain.The 1B model applies SAGE to the full MoE output and uses a rank-m section as a proper subspace when n ≫ dh.Source: §A.4 · p. 14

B.1 DPG-BENCH EVALUATION AND CATEGORY BREAKDOWN

On DPG-Bench, SAGE exceeds the baseline across reported top-level categories at guidance scales of at least 3 and across all second-level categories at scales 3 and 5. The reported gains span multiple aspects of compositional generation.

B.1 DPG-BENCH EVALUATION AND CATEGORY BREAKDOWN

  • At every tested s ≥3, SAGE exceeds the baseline in all five DPG-Bench top-level categories.At s=3, the largest gains are Other (+6.0 pp) and Entity (+4.7 pp); at s=10, they are Entity (+10.8 pp) and Relation (+8.1 pp).Source: §B.1 · p. 14
  • At s=3 and s=5, SAGE improves over the baseline in all 13 DPG-Bench subcategories.Counting gains are +6.4 and +12.4 pp, and whole-entity presence gains are +5.1 and +8.6 pp, respectively.Source: §B.1 · p. 14
  • The reported subcategory gains also cover size, texture, and spatial relations, spanning multiple aspects of compositional generation.The evaluation compares Baseline, SAGE, and a KL routing constraint across guidance scales from 1 to 10.Source: §B.1 · p. 14

B.2 GENEVAL PER-TASK BREAKDOWN

On GenEval, SAGE leads the baseline on Overall at every guidance scale s ≥3, while its category-level gains and declines vary at the shared operating point s=3.

B.2 GENEVAL PER-TASK BREAKDOWN

  • SAGE leads the baseline on GenEval Overall at every guidance scale s ≥3.Overall is the mean over the six tasks.Source: §B.2 · p. 14
  • At s=3, SAGE improves four of six categories, led by color-attribute binding (+11.5 pp), two-object composition (+3.5 pp), and single-object presence (+3.4 pp).The increases are accompanied by decreases in counting (−1.9 pp) and colors (−2.4 pp).Source: §B.2 · p. 14
  • The KL routing constraint degrades every GenEval category at every tested guidance scale.The passage relates this pattern to the loss of expert specialization.Source: §B.2 · p. 14

B.3 ADDITIONAL ABLATIONS

The ablations test intermittent SAGE activation and compare shared-expert and ProMoE configurations. They find that intermittent activation retains nearly all of the from-scratch gain, while SAGE outperforms the alternatives in the shared-expert comparison.

B.3 ADDITIONAL ABLATIONS

  • The ablations compare DPG-Bench and GenEval scores at s∈{3, 5}, while the activation-frequency table covers every-step, every-second-step, and every-fifth-step activation.The shared-expert and ProMoE comparison is reported at both guidance scales.Source: §B.3 · pp. 14–15, Table 7 · p. 16, Table 8 · p. 16
  • Intermittent SAGE activation every second or fifth training step retains nearly all of the from-scratch gain.The authors attribute this to the minibatch subspace constraint reshaping branch writes across steps rather than enforcing a per-token identity.Source: §B.3 · p. 15, Table 7 · p. 16
  • At s=3, the Baseline still benefits from SAGE by +5.73 DPG, with the gap widening dramatically at s=7 to 10.The comparison focuses on how shared experts and ProMoE affect performance at s=3 and s=5.Source: §B.3 · pp. 14, 16
  • Removing the shared expert worsens Baseline performance, while SAGE recovers most of the gap; ProMoE provides partial relief but does not match SAGE.ProMoE reserves one expert for the unconditional branch without aligning the remaining routed activations.Source: §B.3 · p. 16, Table 8 · p. 16

B.4 SAGE PRESERVES ROUTING DIVERSITY

SAGE reduces the subspace residual while leaving conditional and unconditional routing behavior essentially unchanged, unlike the KL routing constraint. Across tested guidance scales, SAGE also achieves better FID than the Baseline.

B.4 SAGE PRESERVES ROUTING DIVERSITY

  • Across guidance scales, SAGE's routing similarity measures are statistically indistinguishable from Baseline, while KL drives Top-2 Jaccard and Top-1 agreement toward 0.95 and JS divergence toward zero.SAGE's JS divergence is marginally higher than Baseline's, indicating it does not push the branches toward common experts.Source: Figure 5 · p. 17, §B.4 · p. 17
  • SAGE reduces the residual outside the conditional subspace while preserving routing diversity, whereas the KL constraint nearly eliminates branch load differences by enforcing routing agreement.At s=3, mean absolute per-expert load differences are 0.0144 for Baseline, 0.0148 for SAGE, and 0.0034 for KL.Source: Figure 6 · p. 17, §B.4 · pp. 17–18

B.5 FID ACROSS CFG SCALES

  • Across CFG scales s∈{1, 3, 5, 7, 10}, SAGE achieves better (lower) FID than baseline; from s=1 to s=10, baseline worsens by +31.5 while SAGE worsens by +23.94.Source: Table 9 · p. 18, §B.5 · p. 18

C TRAINING DYNAMICS AND HEALTH

SAGE keeps the alignment residual close to zero while preserving the primary optimization and routing-health statistics. Baseline and SAGE have nearly identical flow-matching loss trajectories, closely matched load-balancing loss, and fewer than 5% dead experts throughout training.

C TRAINING DYNAMICS AND HEALTH

Loading 2609.34525v1…