🤖 machine learning

MoECa: Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers

The paper proposes MoECa, a fine-grained caching framework that accelerates Diffusion Transformers with Mixture-of-Experts by performing feature reuse at the expert-branch level rather than the token level, achieving up to 2.93 $\times$ speedups while preserving generation quality.

Original authors: Maoliang Li, Haojing Chen, Jiayu Chen, Zihao Zheng, Xinhao Sun, Hailong Zou, Xiang Chen

Published 2026-08-05

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Maoliang Li, Haojing Chen, Jiayu Chen, Zihao Zheng, Xinhao Sun, Hailong Zou, Xiang Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a masterpiece, but instead of using a single brush, you have a magical team of hundreds of tiny artists, each specializing in a different style—some are masters of texture, others of color, and some of lighting. This is how modern AI image generators, known as Diffusion Transformers, work. They start with a screen full of static noise and slowly "denoise" it, step by step, until a clear picture emerges. To make these images look amazing, the AI uses a clever trick called "Mixture-of-Experts" (MoE). Think of this as a smart manager who, for every tiny patch of the image, only wakes up the specific team of artists best suited for that job, rather than asking everyone to work at once. This saves energy and allows the AI to be incredibly smart.

However, there's a catch. Creating an image takes hundreds of these "denoising steps," and the AI has to do a massive amount of math for every single step. It's like having to repaint the entire canvas from scratch every time you add a new layer of detail. To speed this up, scientists have tried a technique called "feature caching." This is like saying, "Hey, this part of the painting hasn't changed much since the last step, so let's just reuse the old paint instead of mixing new colors." The problem is, the old way of doing this was too clumsy. It treated the whole tiny patch of the image as one single unit. If even one tiny detail in that patch needed a fresh coat of paint, the old method would force the AI to repaint the entire patch, wasting time. If it was too conservative, it would reuse the whole patch even when parts of it had changed, leading to a blurry mess.

This is where a new study called MoECa comes in. The researchers realized that because the AI uses specialized "expert" artists, the caching strategy needs to be just as specialized. Instead of treating a patch of the image as a single block, MoECa looks inside the block and sees the individual experts working on it. They found that while one expert might be busy changing the texture of a leaf, another expert working on the same patch might be perfectly happy reusing their old work on the sky. By breaking the image down to the level of these individual expert branches, MoECa can decide exactly which parts to repaint and which parts to reuse.

The paper shows that this "fine-grained" approach is a game-changer. By aligning the caching strategy with how the experts actually work, the AI can skip redundant calculations without losing quality. In their tests, the new method made the image generation process up to 2.93 times faster while keeping the images looking just as sharp and detailed as the slow, original version. It's like upgrading from a construction crew that repaints the whole house every time a window needs fixing, to a team that only paints the specific window that needs it, leaving the rest of the house untouched and ready to go.

The Problem: The "All-or-Nothing" Mistake

To understand why MoECa is needed, we have to look at how these AI models think. In the past, when trying to speed up image generation, researchers used a "token-level" caching method. Imagine a token as a tiny square on your canvas. The old rule was simple: "If this square changed even a tiny bit, repaint the whole square. If it didn't change, leave the whole square alone."

But in models with "Mixture-of-Experts," a single square isn't just one thing; it's a collaboration between several different "expert" branches. The paper's analysis revealed a flaw in the old rule: these experts don't all change at the same time. One expert might be working on the rough texture of a tree bark, which changes rapidly as the image forms. Another expert on the same square might be working on the smooth green leaves, which stay very stable.

If you use the old "all-or-nothing" rule, you face a dilemma. If you decide to repaint the whole square because the bark changed, you waste time repainting the stable leaves. If you decide to reuse the whole square to save time, you end up with blurry, wrong leaves because the bark wasn't updated. The paper argues that this mismatch between the "whole square" caching and the "split expert" reality is the main reason why current speed-up methods aren't working as well as they could.

The Solution: MoECa's "Expert-Level" Strategy

The authors propose MoECa (Mixture-of-Experts Caching), a system that changes the rules of the game. Instead of asking, "Should we repaint this whole square?", MoECa asks, "Should we repaint this specific expert's work on this square?"

Here is how MoECa works in practice:

Branch-Level Caching: It treats each expert branch as its own little unit. It remembers the output of the "bark expert" and the "leaf expert" separately.
Smart Matching: When the AI moves to the next step, MoECa checks which experts are active. If the "bark expert" is still the same one as before, it checks if its work has changed enough to matter. If the "leaf expert" is also the same, it checks that one too.
Adaptive Control: The system is smart enough to know that some experts are more sensitive than others. For example, an expert that focuses on fine details (like the edge of a coffee pot) is very sensitive to changes. MoECa gives these experts a lower "reuse threshold," meaning it's more likely to repaint them to ensure quality. An expert working on a blurry background might get a higher threshold, allowing it to be reused more often.
Synchronized Updates: A crucial part of MoECa is making sure the "attention" part of the AI (which looks at the whole image to understand context) stays in sync with the "expert" part. If the experts update but the attention part doesn't, the image gets confused. MoECa ensures they update together, preventing the AI from getting "drunk" on old information.

The Results: Faster Without the Blur

The researchers tested MoECa on several popular AI models, including the DSMoE family and DiT-MoE. The results were impressive.

Speed: In the best cases, MoECa made the image generation 2.93 times faster (specifically on the HiDream-I1 model for text-to-image tasks). On other models, it achieved speedups of around 2.14x to 2.83x.
Quality: Despite the massive speed increase, the quality of the images remained nearly identical to the original, slow method. The paper measured this using metrics like FID (which measures how realistic the images look) and PSNR (which measures how close the accelerated image is to the original). For instance, on the DSMoE-L-E48 model, MoECa achieved a speedup of 2.83x while keeping the FID score at 9.54, which is almost the same as the original model's 9.20.
Comparison: When compared to other speed-up methods like ToCa, DuCa, and TeaCache, MoECa consistently produced the fastest results with the lowest computational cost (FLOPs), all while maintaining or slightly improving image quality.

The paper also looked at why this works so well. They found that in models with more experts (like the E48 configuration with 48 experts), there are more opportunities to find "stable" branches to reuse. This is why the E48 models saw slightly higher speedups (up to 2.83x) compared to models with fewer experts. The more specialized the team, the more you can pick and choose who works and who rests.

What This Means for the Future

The paper concludes that the key to speeding up these powerful AI models isn't just to make them work faster in general, but to understand how they work internally. By realizing that the "experts" inside the AI have different rhythms and needs, we can build smarter caching systems that respect those differences.

MoECa doesn't require retraining the AI models; it works with the models as they are, simply changing how they process information during the image creation. While the current study focuses on single-computer image generation, the authors suggest that this "fine-grained" approach could be the key to unlocking even faster generation for video and more complex tasks in the future. For now, it proves that sometimes, the fastest way to paint a masterpiece is to let the right artists rest while the others keep working.

Technical Summary: MoECa – Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers

1. Problem Statement

Diffusion Transformers (DiTs) have become the dominant paradigm for visual generation, and the integration of Mixture-of-Experts (MoE) architectures (DiT-MoE) has further enhanced model capacity through sparse activation. However, inference efficiency remains a bottleneck due to the iterative denoising process. While feature caching has emerged as a primary acceleration strategy by reusing intermediate features across timesteps, existing methods operate primarily at the token level.

The authors identify a fundamental structural mismatch in applying token-level caching to DiT-MoE models:

Internal Decomposition: In DiT-MoE, a single token's feature update is decomposed into multiple routed expert branches.
Heterogeneity: Different expert branches within the same token exhibit distinct temporal evolution patterns. Some branches remain stable across timesteps, while others change rapidly.
Suboptimality: Treating the entire token as an atomic caching unit forces a "all-or-nothing" decision. This leads to either unnecessary recomputation of stable branches (redundancy) or excessive reuse of rapidly changing branches (quality degradation).

Consequently, the granularity of existing caching mechanisms is misaligned with the intrinsic feature evolution structure of MoE models.

2. Methodology: MoECa

To address this misalignment, the authors propose MoECa, a fine-grained, feature- and structure-aware caching framework that shifts the caching unit from the token level to the expert-branch level.

2.1 Expert Cache Architecture

MoECa maintains a cache of intermediate states for each expert branch rather than the whole token.

State Definition: For a token $i$ at timestep $t$ , the system tracks the current routing result (selected experts and weights) and compares it against the cached reference state from the last refresh step $t_0$ .
Partitioning: The current branches are partitioned into two sets:
1. Recomputation Set ( $R_i^t$ ): Branches that require fresh computation.
2. Reuse Set: Branches that can safely reuse cached features.
Synchronization: To prevent representation drift, MoECa synchronizes updates between the MoE path and the Attention path. If any branch is recomputed, the attention output is also updated in the same step; otherwise, both paths reuse their cached states.

2.2 Expert-Level Reuse Decision

The decision to reuse or recompute a specific branch is based on a branch-wise scoring mechanism that evaluates three signals:

Routing Weight ( $w$ ): The current contribution of the branch to the token representation.
Temporal Drift ( $\Delta$ ): The change in routing weight between the current step and the cached step.
Historical Contribution ( $o$ ): The relative magnitude of the branch's output feature in the cached state.

The recomputation score $s_{i,p}^t$ is calculated as:
$s_{i,p}^t = w_{i,p}^t + \lambda_1 \Delta_{i,p}^t + \lambda_2 o_{i,p}^{t_0}$
Branches exceeding a threshold $\tau$ are selected for recomputation.

2.3 Expert-Aware Adaptive Control

Recognizing that different experts exhibit varying spatial response patterns (specialization), MoECa introduces an adaptive threshold mechanism:

Spatial Entropy: The system estimates the spatial response entropy of each expert online using an Exponential Moving Average (EMA) of activation statistics.
Adaptive Thresholding: Experts with low entropy (concentrated, detail-oriented responses) are assigned lower thresholds, making them more likely to be recomputed to preserve local details. Experts with high entropy (broad, background responses) are assigned higher thresholds, favoring reuse.
Formula: $\tau_e^t = \tau_0 (1 + \alpha(2\bar{H}_e^t - 1))$ , where $\bar{H}$ is the normalized entropy.

3. Key Contributions

Structural Analysis: The paper provides an analysis revealing that cross-timestep redundancy in DiT-MoE is better characterized at the expert-branch level than the token level, highlighting the structural mismatch of existing token-level caching.
MoECa Framework: The proposal of a fine-grained caching framework that aligns feature reuse with the MoE decomposition structure, enabling independent reuse decisions for different branches within the same token.
Adaptive Control: The introduction of an expert-aware adaptive thresholding mechanism that accounts for the heterogeneity of expert specialization, balancing acceleration with generation quality.
Empirical Validation: Comprehensive evaluation across multiple DiT-MoE models (DSMoE, DiT-MoE, HiDream-I1) demonstrating favorable speed-quality trade-offs.

4. Experimental Results

The authors evaluated MoECa on class-conditional generation (ImageNet) and text-to-image generation (DrawBench200) using models ranging from DSMoE-S-E16 to DSMoE-L-E48.

Speedup: MoECa achieves significant speedups compared to dense baselines and other caching methods.
- On DSMoE-L-E48, it achieves a 2.14× latency speedup and a 2.83× FLOPs speedup while maintaining an FID of 9.54.
- On HiDream-I1 (1024px), it achieves a 2.93× FLOPs speedup (5.61s latency vs. 14.29s baseline).
Quality Preservation: MoECa consistently outperforms or matches other caching baselines (ToCa, DuCa, TeaCache) in terms of FID, Inception Score, PSNR, and SSIM.
- In qualitative comparisons, MoECa preserves high-frequency details (e.g., textures, edges) better than token-level methods, which tend to exhibit structural blurring.
Efficiency: The method reduces FLOPs more effectively than step-reduction strategies (e.g., 33% steps) which often suffer from significant quality degradation.
Overhead: The control logic (scoring, routing, cache updates) adds minimal overhead, accounting for only 4.8% of the end-to-end inference time on DSMoE-S-E16. Memory overhead is moderate, storing branch states without exceeding the memory footprint of other token-level caching methods by a large margin.

5. Significance and Claims

The paper claims that MoECa represents a necessary evolution in diffusion acceleration strategies for MoE architectures. By shifting the caching granularity to align with the expert decomposition, MoECa resolves the suboptimality inherent in token-level reuse.

The authors emphasize that their approach is training-free, requiring no modification to the pretrained model's weights or routing graph. The significance lies in demonstrating that fine-grained, structure-aware reuse can unlock substantial computational savings (up to 2.93×) without compromising the generation quality, a trade-off that coarse-grained methods (step reduction) or misaligned fine-grained methods (token-level caching) fail to achieve simultaneously.

The study concludes that while the current work focuses on single-GPU image generation, the principles of aligning caching granularity with internal model decomposition are applicable to broader contexts, though video generation and expert-parallel serving remain future work.