wan2.2-i2v-a14b-physalign-lora

Wan2.2-I2V-A14B · PhysAlign LoRA (faithful reproduction, pure-Gram)

LoRA post-training of Wan2.2-I2V-A14B with a V-JEPA2 feature-alignment (Gram) loss on synthetic physics videos (PhyCo Kubric), reproducing the PhysAlign recipe (arXiv:2603.13770) without the 3D depth loss. Improves the intrinsic physical consistency of generated image-to-video.

Results (Physics-IQ, official scorer, full 198)

Pure Gram alignment (no 3D) reaches near paper-level gain, confirming the paper's ablation that the geometry loss is not the bottleneck.

What's in here

The A14B is a dual-expert MoE (high-noise + low-noise). PhysAlign requires training both experts, each on its own timestep range:

• low_noise_expert/pytorch_lora_weights.safetensors — LoRA for transformer_2 (σ < 0.9)
• high_noise_expert/pytorch_lora_weights.safetensors — LoRA for transformer (σ ≥ 0.9)
• config_*.yaml — the exact training configs

LoRA rank 32 on q/k/v/o + ffn.0/ffn.2.

The 4 faithful-reproduction fixes (vs a naive single-expert run, +0.58 → +5.29)

1. Train BOTH MoE experts separately (the high-noise expert governs coarse motion/physics) — the dominant fix.
2. Multi-block alignment at DiT blocks 12/16/20/24 (not a single block).
3. Teacher input keeps aspect ratio (V-JEPA2 resize to height 160, not a 256² square stretch).
4. Student→teacher resampling, Gram computed at the teacher's grid resolution.

Pure Gram loss: L = L_FM + 0.25 · L_Gram, margin 0.1, constant LR 1e-4, 3000 steps/expert.

Usage (sketch)

# load the low-noise LoRA onto transformer_2 and the high-noise LoRA onto transformer
pipe.transformer_2.load_lora_adapter("low_noise_expert/pytorch_lora_weights.safetensors")
pipe.transformer.load_lora_adapter("high_noise_expert/pytorch_lora_weights.safetensors")
# then run I2V as usual (432×768, 121 frames @24fps in our eval)

Teacher (V-JEPA2) and the projector are training-only and not needed at inference — only the base model + these two LoRAs.