Vision-language MoE workloads change with image resolution, image count, dynamic tiling, and prompt length. A standard mixed token-level auxiliary loss can look balanced at one image-text ratio while the image and text loads remain large and point toward different experts. ReBA follows two measured routing boundaries:

1. image and text occupy distinct router-input regions, so ReBA uses separate modality balance terms;
2. patches from one image form a correlated routing instance, so ReBA averages within each image and balances image profiles across the batch.

Standard balancing can learn image and text load deviations that cancel only near the training composition. ReBA reduces the image-text load gap.

Image and text tokens occupy distinct router-input regions in both split and released native vision-language MoEs.

Visual tokens also retain a clear image boundary. This observation motivates one auxiliary routing profile per image.

The repository provides model conversion, training, routing probes, evaluation, and plotting code. Model weights and datasets are downloaded separately.

Python 3.10 or 3.11 is recommended. Install a PyTorch build compatible with the machine's CUDA driver before running the setup:

git clone https://github.com/ZiangWu-77/ReBA.git ReBA
cd ReBA
python3 -m venv .venv
source .venv/bin/activate

# Example only. Select the wheel index for your CUDA driver.
pip install torch torchvision \
  --index-url https://download.pytorch.org/whl/cu128

./setup.sh

setup.sh installs pinned Transformers, ms-swift, and VLMEvalKit revisions under third_party/src/ and applies the ReBA patches.

python scripts/make_split_moe.py \
  --src /path/to/Qwen3-VL-4B-Instruct \
  --dst /path/to/Qwen3-VL-4B-SplitMoE-e4-top2 \
  --num-total-experts 4 \
  --num-base-chunks 4 \
  --top-k 2 \
  --num-shared-experts 0

The converter partitions matching FFN gate/up rows and down columns. It adds a router but does not copy the original FFN blocks.

The training file must follow an ms-swift-compatible JSON or JSONL format. Image paths can be relative to ROOT_IMAGE_DIR.

NPROC_PER_NODE=8 \
ROUTER_AUX_LOSS_TYPE=reba \
ROUTER_AUX_LOSS_COEF=0.1 \
./scripts/train_reba.sh \
  /path/to/split-checkpoint \
  /path/to/train.jsonl \
  /path/to/image-root \
  /path/to/output

The public method selector is ROUTER_AUX_LOSS_TYPE=reba. ReBA forms one routing instance per image and applies separate image and text balance terms. The training script uses this selector by default.

For a short plumbing test, append --max_steps 2.

python scripts/run_probe.py \
  --model /path/to/checkpoint \
  --data examples/probe_example.jsonl \
  --out /path/to/probe-output \
  --arch qwen3vl_split \
  --condition reba \
  --max-samples 100 \
  --attn sdpa

Replace the placeholder image path in the example file. The probe writes metrics.json and optional compressed routing arrays to the selected output directory. Generated outputs are ignored by Git.

Compare standard balancing and ReBA:

python analysis/plot_probe_complementarity.py \
  --std /path/to/std-probe/metrics.json \
  --reba /path/to/reba-probe/metrics.json \
  --output /path/to/complementarity.pdf

For InternVL dynamic tiling, use --arch internvl and set --internvl-max-num to the desired tile cap.

evaluate_load_accuracy.py runs batched forwarding on LMUData-style TSV files. The script collects true top-k expert dispatch counts and supports capacity-constrained log-probability scoring:

python scripts/evaluate_load_accuracy.py \
  --model /path/to/checkpoint \
  --data /path/to/benchmark.tsv \
  --bench pope \
  --arch qwen3vl_split \
  --batch-sizes 1 4 16 \
  --cf 0 1.0 \
  --out /path/to/evaluation.json

The patched VLMEvalKit checkout provides standard benchmark scorers. Dataset downloads and judge API credentials remain external.

The generic plotting utility accepts a CSV generated from your own sweeps:

method,alpha,rms_cv
Std,<image-token fraction>,<measured RMS-CV>
ReBA,<image-token fraction>,<measured RMS-CV>

python analysis/plot_composition_shift.py \
  --input /path/to/your_sweep.csv \
  --output /path/to/composition_shift.pdf

Original repository code is released under the Apache License 2.0. Patched third-party files remain subject to their upstream Apache 2.0 licenses, which are included under third_party_licenses/.

Citation metadata will be added after the anonymous review period.