LingXi-Image-MoE

<div align="center">

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

<a href="https://arxiv.org/abs/xxxx.xxxxx" target="_blank"><img src="https://img.shields.io/badge/Paper-b5212f.svg?logo=arxiv" height="22px"></a> <a href="https://huggingface.co/shyai/LingXi-Image-MoE" target="_blank"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model-d96902.svg" height="22px"></a> <a href="https://github.com/Lingxi-Qihang/LingXi-Image-MoE" target="_blank"><img src="https://img.shields.io/badge/Code-181717.svg?logo=github" height="22px"></a> <a href="https://modelscope.cn/models/haohanxingcheng/LingXi-Image-MoE" target="_blank"><img src="https://img.shields.io/badge/ModelScope-624aff.svg" height="22px"></a>

</div>

🌐 中文文档 (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

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)

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):

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

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

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

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

---

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:

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.