QuantWAMs: Calibrating at the Right Granularity for... | AI Research

QuantWAMs is a post-training quantization (PTQ) framework designed to make World Action Models (WAMs) efficient for real-robot deployment. WAMs are computationally expensive because they perform iterative denoising and operate in a closed loop, where every action influences future observations. QuantWAMs addresses the limitations of existing quantization methods, which often fail for WAMs because they rely on open-loop objectives and ignore the specific structural and distributional requirements of coupled video-action models.

Aligning Quantization with Model Context

The researchers, Jiacheng Zhou et al., argue that quantization decisions must be calibrated based on the specific context of the model. They identify three axes for this calibration: the model structure, the rollout distribution, and the task objective. QuantWAMs introduces three strategies to align these factors:

• Shared-basis outlier calibration: This method pools activation evidence only across modules that share the same coordinate system. By using a shared Hadamard rotation, the framework identifies and preserves critical channels at high precision while quantizing others, ensuring that the quantization mask remains stable across different model components.

• Co-training-objective saliency: Instead of using a single objective, this approach computes empirical-Fisher scores from the joint video-action gradient. It assigns weight precision at a layer-level granularity, prioritizing the most important layers based on their contribution to both video prediction and action generation.

• Fixed-intervention rollout auditing: To address the closed-loop nature of WAMs, this strategy uses real-world rollout states to revise denoising-step schedules. By replaying recorded FP16 trajectories under an unprotected intervention, the framework identifies which denoising steps are most sensitive to quantization error and protects them, all while maintaining the original precision budget.

Performance and Deployment

The authors evaluated QuantWAMs on Fast-WAM and LingBot-VA architectures using the RoboTwin 2.0 and LIBERO benchmarks, as well as real-world manipulation tasks with an AgiBot G2 robot. Under a W4A4-dominant setting, the simulation results showed that QuantWAMs maintained performance within 0.2–0.7 percentage points of the full-precision FP16 model.
In terms of efficiency, the framework reduced peak weight-and-activation memory to approximately 29% of the FP16 baseline and achieved 1.4–1.6× speedups at the block level. These results indicate that the method is feasible for real-robot manipulation tasks where computational resources are constrained.

Limitations and Considerations

The researchers note that their optimality claims are restricted to the specific energy surrogates and calibration data used in their experiments. Because the framework relies on finite-sample estimates, the stability of the quantization mask depends on the quality of the calibration set. Furthermore, the authors emphasize that the sensitivity profiles generated by their auditing process should not be interpreted as direct estimates of marginal task gain, but rather as a tool for repairing and optimizing the precision schedule. The framework requires access to the original co-training targets and a backward pass to compute saliency scores, though it does not require updating the pretrained weights.