LingXi-Image-MoE: From Alternating to Full-Layer
A Comparative Study of Mixture-of-Experts Architectures for Text-to-Image Diffusion Models

🌐 中文文档 (Chinese README) 🌐 English README

🔥 News

• [Coming Soon] ProMoE‑L (1.68B) joint training is underway with 2.66M mixed data; a unified framework for text-to-image generation and image editing is coming soon!
• [2026-07-14] The paper "From Alternating to Full-Layer: A Comparative Study of Mixture-of-Experts Architectures for Text-to-Image Diffusion Models" has been released on arXiv!
• [2026-07-14] Model weights are now available on HuggingFace and ModelScope
• [2026-07-14] Full code open-sourced on GitHub; training records and routing logs are included in the HuggingFace and ModelScope repositories

---

About This Project

This project is built upon the ProMoE (ICLR 2026) framework. We adapt it from class-conditioned (ImageNet) to text-conditioned generation, delivering a unified framework that supports both text-to-image generation and image editing. The Full-Layer MoE architecture reaches the generation quality of Alternating MoE at over 800K steps after only 400K steps, achieving a 2–2.5× convergence speedup.

---

👤 A Word from the Author

I am an independent researcher, working alone with a single GPU, paying out of my own pocket to explore the frontiers of visual MoE.

I am currently building and running an AI company called "LingXi Qihang Intelligent Technology" — yes, product development, R&D, and operations are all handled by me alone. The company is expected to officially launch next month, and at that time it will include:

• 🎨 LingXi AIGC (WeChat Mini Program): AI painting, video generation
• 🔒 LingTong AI (Android app): Offline AI assistant with privacy-first voice interaction

The official website is also under construction. LingXi-Image-MoE-1.0 is my independent research achievement in the direction of visual MoE, and it will serve as the foundation model for LingXi Qihang's future image generation capabilities.

---

📌 Roadmap

• ✅ LingXi-Image-MoE-1.0 (ProMoE‑B, 0.47B): Full-Layer MoE text-to-image model, DPG‑Bench score 63.86, surpassing SD v1.5
• 🚧 ProMoE‑L (1.68B): Joint training in progress with 2.66M mixed data, unified framework for text-to-image + image editing
• 🧪 Unified Editing Model: 7:3 mixed training to achieve strong editing capability while maintaining generation quality
• 📈 Larger-Scale Exploration: 3B+ parameter MoE with cross-attention routing and instruction-level routing
• 📚 Diagnostic Paper: Systematic Loss Spike fixes, U-shaped layer-wise specialization, progressive awakening (arXiv:2605.19378)

---

Main Results

Full-Layer MoE vs Alternating MoE: Generation Quality Comparison over 1M Steps

Evaluation Dimension	Alternating MoE (400K steps)	Full-Layer MoE (400K steps)
Background Purity	Obvious gray "haze", blurred object edges	Clean background, sharp subject edges
Multi-person Interaction	Melted faces, fused limbs	Clear facial features, natural expressions
Text Rendering	Strokes blurred into noise, unreadable	Sharp strokes, separated characters
Two-panel Comic	Blurred boundaries, mixed content	Clear panel division, independent left/right themes
Convergence Speed	Baseline	2–2.5× speedup

Core Finding: Full-Layer MoE reaches the generation quality of Alternating MoE at over 800K steps after only 400K steps. The Dense layers are the structural root of multi-task conflict and information loss in the alternating architecture.

DPG-Bench Evaluation (2.85M-step checkpoint)

Model	Global	Entity	Attribute	Relation	Other	Overall↑
SD v1.5 (0.86B)	74.63	74.23	75.39	73.49	67.81	63.18
LingXi-Image-MoE (0.47B)	81.90	75.28	80.12	87.02	54.00	63.86
PixArt-α	74.97	79.32	78.60	82.57	76.96	71.11
SDXL	83.27	82.43	80.91	86.76	80.41	74.65

Key Sub-scores:

• Spatial Relation: 87.83 — surpasses SDXL, approaching SOTA
• Text Rendering: 86.00 — historical best (with dedicated text enhancement data)
• Texture: 82.74 — surpasses SD v1.5 and PixArt-α
• Entity Completeness (face/details): 75.28 — needs improvement (limited by VAE capacity and training stability)
• Counting Ability: 46.00 — current biggest weakness

---

Core Features

ProMoE Two-Stage Routing

• Conditional Routing: Divides tokens into conditional (samples not dropped in CFG) and unconditional (CFG null-text samples) groups
• Prototype Routing: Assigns experts via cosine similarity between learnable prototype vectors and tokens; each token activates Top-K routed experts
• Routing Contrastive Loss: Explicitly enhances intra-class compactness and inter-class diversity of prototype routing
• Unconditional Expert Relocation: Preserves ProMoE's two-stage routing structure while reassigning the unconditional expert (Expert 12) to specifically handle CFG null-text samples

MM-DiT Text-Image Fusion

• Text sequence is mapped to the model's hidden space and concatenated with image tokens along the sequence dimension
• Self-attention processes text and images simultaneously, without any cross-attention module
• Text global vector is injected into the model via AdaLN modulation

Dynamic Resolution (Continuous Rotary Position Encoding)

• Image tokens: 2D RoPE (Rotary Position Embedding), coordinates normalized to [0,1]
• Text tokens: 1D RoPE encodes sequence positions
• Advantage: A single set of weights supports arbitrary resolutions (256/512/768/1024+), no interpolation or fine-tuning needed
• Inspired by FiT (NeurIPS 2024)'s continuous position encoding idea

Unified Symbiosis Framework

• Text-to-Image (T2I): Source latent is a zero black image
• Image Editing: Source latent is the original encoding of the target image
• One model, one training process; the two tasks are distinguished via source_latent

---

Environment Setup

conda create -n promoe python=3.10 -y
conda activate promoe
pip install -r requirements.txt

Data Preprocessing

We use T5-base to encode text and sd-vae-ft-mse to encode images, caching them as preprocessed files to accelerate training.

Below we use the hf/ma-xu-fine-t2i dataset as an example to illustrate the full data preparation pipeline.

1. Download Data

Download tar shards from ModelScope to your local machine:

bash preprocess/download.sh

This will concurrently download train-000000.tar ~ train-001310.tar, saved to /data/coding/ by default.

2. Extract Data

python preprocess/extracted_data.py

After extraction, the directory structure is:

curated_extracted_data/
└── synthetic_original_prompt_square_resolution/
    ├── 000000000.jpg / 000000000.txt
    ├── 000000001.jpg / 000000001.txt
    └── ...

Each image corresponds to a .txt file with the same name containing the text description.

3. Extract Features

Use preprocess_ma-xu-fine-t2i.py to batch extract VAE latents and T5 text features:

python preprocess/preprocess_ma-xu-fine-t2i.py \
  --image_root /path/to/curated_extracted_data \
  --output_dir /path/to/preprocessed_output \
  --t5_path /path/to/t5-base \
  --vae_path /path/to/sd-vae-ft-mse \
  --image_size 512 \
  --batch_size 8

The script automatically scans all images under image_root, looks for the matching .txt file, and outputs the following structure:

preprocessed_output/
├── latents/
│   ├── target_512/            # 512-resolution latents
│   └── target/                # 256-resolution latents
└── t5_text_features/
    ├── 000000000_seq.npy      # T5 sequence features (512, 768)
    ├── 000000000_mask.npy     # Attention mask (512,)
    └── ...

After extraction, the script automatically removes the original images and text files to save space.

4. Configure Training

Set the preprocessing output paths in the YAML config file:

latent_root: "/path/to/preprocessed_output/latents"
text_root: "/path/to/preprocessed_output/t5_text_features"

The UnifiedSymbiosisDataset in the training script train_GradientAccumulationSteps.py will automatically load data from these two directories.

Dataset Composition

The actual training used about 2.38M image-text pairs. The planned total dataset is about 3.4M images (some later stages not yet executed):

Subset	Size	Description
Fine-T2I (enhanced + original)	2.09M	Synthetic data covering long and short texts
Fine-T2I (curated)	168K	Real professional photos to remove "AI plastic feel"
Face/Text Enhancement Data	192K	Improves face detail and text rendering quality
Pico-Banana-400K	400K	Real image editing (later stage)

Four-Stage Progressive Training

bash train.sh                  # Stage 2a → 4 (default)
bash train.sh --full           # Stage 1  → 4
bash train.sh --stage 3        # Stage 3  → 4
bash train.sh --accum 4        # 4-step gradient accumulation (default 1)

Default uses no gradient accumulation. The script automatically orchestrates all stages, and checkpoints seamlessly connect between stages.

Stage Configuration

Stage	Core Task	Dataset	Resolution	Learning Rate	Steps	Status
1	Routing convergence & multi-task alignment	OpenGPT-4o-Image (80K)	256	1×10⁻⁴	2200K	✅ Done
2a	T2I Scaling (low resolution)	Mixed data (~2.38M)	256	2.5×10⁻⁵	200K (actual to 2.85M)	✅ Done
2b	T2I Scaling (high resolution)	Same as above	512	5×10⁻⁵	100K	⏳ Not executed
3	Editing specialization	Pico-Banana-400K (400K)	512	5×10⁻⁵	100K	⏳ Not executed
4	Unified joint fine-tuning	Mixed 7:3 (T2I + Editing)	512	5×10⁻⁵	50K	⏳ Not executed

Note: The 2.09M data in Stage 2 were randomly sampled and merged from the synthetic_enhanced_prompt_square_resolution (1.54M) and synthetic_original_prompt_square_resolution (1.3M) subsets of the Fine-T2I dataset. After actual preprocessing, we obtained about 2.09M image-text pairs. The learning rate for Stage 2a has been corrected to the actually used value of 2.5×10⁻⁵.

Stage Configuration Files

File	Purpose
configs/stage2a.yaml	Stage 2a: Fine-T2I, 256, lr=2.5e-5, 200K
configs/stage2b.yaml	Stage 2b: Fine-T2I, 512, lr=5e-5, 100K
configs/stage3.yaml	Stage 3: Pico-Banana-400K, 512, lr=5e-5, 100K
configs/stage4.yaml	Stage 4: Mixed, 512, lr=5e-5, 50K

Running a Single Stage Manually

python train_GradientAccumulationSteps.py \
  --config configs/stage2a.yaml
# Default is no gradient accumulation; if needed: --grad_accum_steps 4

Key Configuration Parameters

Parameter	Default	Description
image_size	256/512	Training resolution
latent_root	—	VAE latent path (contains target/ subdirectory)
text_root	—	T5 text features path
total_train_batch_size	1	Total batch size (used with grad_accum_steps)
lr	1e-4	Learning rate
num_steps	Stage endpoint steps	Total training steps
resume_checkpoint_step	—	Resume from checkpoint at given step
grad_accum_steps	1 (CLI)	Gradient accumulation steps
routing_contrastive_lam	0	Routing contrastive loss weight (set to 0 to disable)
load_balance_loss_coef	0.01	Load balancing loss weight

Notes

• When optimizer states cannot be loaded due to model structure changes, they are automatically reset, without affecting model weight loading.
• Checkpoint loading uses strict=False; old pos_embed / text_pos_embed are safely ignored.
• Stage 2b requires switching latents to 512 resolution; the script automatically creates symbolic links.
• Data for Stage 3/4 (Pico-Banana-400K, mixed data) must be preprocessed in advance; the script auto-executes when it detects the directories, otherwise skips.

---

Quick Inference (inference.py)

A lightweight standalone inference script that only needs command-line arguments, without a full YAML config:

# Single prompt
python inference.py \
  --config configs/004_ProMoE_B.yaml \
  --ckpt E:/AImodel/ProMoE/ProMoE_B/ProMoE_TC_B/004_ProMoE_B/checkpoints/ckpt_step_2850000.pth \
  --prompt "A cat holds a poster" \
  --output outputs \
  --vae_path E:/AImodel/ProMoE/sd-vae-ft-mse \
  --t5_path E:/AImodel/t5-base \
  --guide_scale 4.0 \
  --seed 42
  
# a spectacular display of fireworks illuminates the night sky with bursts of red, white, and blue. the vibrant colors reflect off a nearby lake, creating a mirror image of the aerial spectacle. in the foreground, silhouettes of a small crowd can be seen gathered to watch the show, with some individuals pointing upwards in awe.

# Batch processing from a prompt directory (each .txt file is one prompt, output has the same name as .png)
python inference.py \
  --config configs/004_ProMoE_B.yaml \
  --ckpt path/to/checkpoint.pth \
  --prompt_dir /path/to/prompt/dir \
  --output outputs

Parameter	Default	Description
--config	Required	YAML config file (only model architecture parameters are read)
--ckpt	Required	Model checkpoint path
--prompt	—	Single text prompt
--prompt_dir	—	Directory containing .txt files; each file provides one prompt
--output	outputs	Output directory
--image_size	512	Generated image size
--guide_scale	4.0	CFG guidance strength
--sample_steps	50	Sampling steps
--seed	0	Random seed (auto-increments for each prompt in batch mode)

---

Sampling and Evaluation

python sample.py --config configs/004_ProMoE_B.yaml

Custom parameters:

python sample.py \
  --config configs/004_ProMoE_B.yaml \
  --step_list_for_sample 2850000 \
  --guide_scale_list 4.0,5.0,7.0 \
  --sample_prompts "a cat sitting on a chair"

---

Training Stability: Loss Spike Diagnosis and Fixes

After 2.8M steps, the training repeatedly encountered Loss Spikes that surged from 0.7 to 3.6. After nearly a month of investigation, we identified and fixed the following key issues:

Issue	Symptom	Root Cause	Fix
Lack of learning rate decay	Severe oscillations in later training	Fixed lr prevents convergence	Introduced CosineAnnealingLR
Improper VAE latent handling	Introduced numerical noise	Mistakenly used variance sampling	Directly take mean, single scaling
Insufficient data integrity	Occasional gradient explosion	Corrupted files contain NaN/Inf	Dataset-level filtering + Loss anomaly skip
Optimizer state contamination	Easy to collapse again even after recovery	Momentum contaminated by abnormal gradients	Reset optimizer state at each healthy resume
VRAM fragmentation and mis-kill	Occasional OOM, normal samples skipped	Loss threshold too low	Raised threshold to 1.5, enabled expandable_segments

These lessons have been fully consolidated into a reproducible engineering baseline and transferred to the next-generation ProMoE-L model.

---

Known Limitations

• Face Details: Limited by the 4-channel VAE and T5-base capacity; occasional collapses occur. Requires larger models and dedicated data.
• Counting Ability: DPG-Bench count score only 46; the training set lacks explicit counting samples.
• Editing Capability: Joint training has not been completed; currently only supports text-to-image generation (inference with zero black image source).
• Ultra-long training stability: The 0.47B parameters are approaching the design limit of ProMoE-B; upgrading to ProMoE-L is recommended.

---

Model Architecture

ProMoE-TC-B configuration (0.47B parameters):

• 12 routed experts + 1 unconditional expert + 1 shared expert
• 12 Transformer layers, all using MoE
• Hidden size 768, 12 attention heads, head_dim=64
• Top-2 routing, cosine similarity routing weights
• AdaLN-Zero conditional normalization
• Sequence-concatenated text injection

File Structure

├── configs/                     # YAML configuration files
│   ├── 004_ProMoE_B.yaml       # Stage 1 config (reference)
│   ├── stage2a.yaml             # Stage 2a: T2I Scaling @ 256
│   ├── stage2b.yaml             # Stage 2b: T2I Scaling @ 512
│   ├── stage3.yaml              # Stage 3: Editing Specialization
│   └── stage4.yaml              # Stage 4: Unified Joint Fine-tuning
├── train.sh                     # Four-stage training pipeline script
├── models/
│   ├── modules.py               # Shared modules (Attention, RoPE, MoE)
│   └── models_ProMoE_TC.py     # ProMoE text-conditional model
├── preprocess/                  # Data preprocessing scripts
│   ├── download.sh              # ma-xu-fine-t2i dataset download
│   ├── extracted_data.py        # tar extraction tool
│   ├── preprocess_ma-xu-fine-t2i.py  # ma-xu-fine-t2i feature extraction
│   ├── preprocess_all.py        # Batch preprocessing (VAE + T5)
│   ├── preprocess_vae.py        # VAE latent encoding
│   └── extract_t5_features.py   # T5 text feature extraction
├── train_GradientAccumulationSteps.py  # Training script (gradient accumulation)
├── inference.py                 # Lightweight inference script
├── sample.py                    # Sampling & evaluation script
└── evaluation/                  # Evaluation scripts

---

Citation

If this project is helpful to your research, please cite:

@misc{sha2025fullayer,
  title={从交替到全层: 混合专家架构在文生图扩散模型中的对比研究},
  author={Haiying Sha},
  year={2025},
  note={LingXi-Image-MoE-1.0},
  url={https://github.com/Lingxi-Qihang/LingXi-Image-MoE}
}

---

References

[1] DiT: Scalable Diffusion Models with Transformers. ICCV 2023.
[2] ProMoE: Routing Matters in MoE. ICLR 2026.
[3] FiT: Flexible Vision Transformer for Diffusion Model. NeurIPS 2024.
[4] PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis. ICLR 2024.