Addresses the issue of gradual drift in camera ego-motion during long-duration generation with diffusion video models: first, linear probes and causal interventions are used to identify which layers and denoising stages in Wan2.2 encode camera motion (the authors found that this encoding rises and then falls as denoising progresses); then, the representations from these layers are aligned to the "inter-frame change" representations provided by a 3D foundation model, training LoRA only on the selected layers and timesteps.
Likes 13 · Code Stars 7
Reasons
- The authors claim that after fine-tuning Wan2.2 on DL3DV and RealEstate10K, motion accuracy, geometric consistency, and temporal smoothness all improved, with geometric consistency increasing by more than 20% compared to the base model, while image quality remained unchanged.
- It is transferable to embodied tasks: using the fine-tuned video model as a world action model improved the success rate on RoboCasa365 mobile manipulation tasks by 10 percentage points (as reported by the authors).
- The repository contains a complete three-stage training implementation, with recipes for layer numbers, timesteps, and LoRA rank hardcoded in the release configuration, along with consistency verification scripts.
Usage
- Code: github.com/PKU-VaLuE-Lab/CausalAlign (Apache-2.0), Python 3.10/3.11, PyTorch 2.4+, CUDA 12.x, pip install -e '.[wan,test]'; additionally requires Wan-AI/Wan2.2-I2V-A14B-Diffusers, LingBot-Map checkpoints, and DL3DV / RealEstate10K data.
- Workflow: tools/prepare_data.py creates the manifest → extract teacher features and cache Wan inputs → run stage1 / stage2 / stage3 sequentially using torchrun → tools/export_lora.py exports the merged model; before starting the actual training, use tools/validate_release.py to check the configuration and data.
- Embodied extension experiments are located in the wam/ directory and use Wan2.2-TI2V-5B, which is not compatible with the default A14B artifacts.
Risks
- No pre-trained LoRA weights are released; users must train them themselves. The training commands in the README are written for a single machine with 8 GPUs, and teacher features must be extracted first, resulting in high computational requirements.
- The implementation is only for Wan2.2, and datasets must be applied for and downloaded separately; all figures are reported by the authors, and since the demo videos have not been viewed, no assessment of image quality is provided.
Open · Code · Paper Page