Technical Summary: DISK – Dynamic Inference SKipping for World Models

1. Problem Statement

Autoregressive world models, particularly those based on Diffusion Transformers (DiTs), have demonstrated strong capabilities in predicting future scene content and ego-motion for closed-loop driving. However, these models suffer from significant inference latency and computational costs. In autoregressive settings, the cost compounds over time because each predicted frame and trajectory conditions the next step. Long-horizon rollouts require thousands of model evaluations across coupled vision and trajectory branches.

Existing acceleration methods fail to address this specific challenge effectively:

• Fixed-step acceleration or global step reduction degrades fidelity and stability, especially when small errors accumulate in closed-loop rollouts.
• Training-time accelerators (e.g., consistency models, progressive distillation) require expensive retraining or distillation, breaking the "plug-and-play" flexibility needed to accelerate existing foundation models without altering their learned behavior.
• Inference-time adaptations (e.g., deterministic formulations, higher-order ODE solvers) do not adapt to per-step difficulty and fail to coordinate visual and trajectory predictions in coupled branches.
• Single-branch optimizers (e.g., early exiting, feature caching) optimize computation within a single network but do not coordinate cross-modal branches, risking desynchronization between appearance and motion.

The core challenge is to reduce inference cost dynamically based on scene complexity (e.g., static highway cruising vs. critical maneuvers) without retraining the model or compromising the consistency between the generated video and the predicted trajectory.

2. Methodology

The authors propose DISK (Dynamic Inference SKipping), a training-free adaptive inference module that operates purely at test time. DISK coordinates two coupled diffusion transformers—one for video (Vision DiT) and one for ego-trajectory (Trajectory DiT)—via dual-branch controllers with cross-modal skip decisions.

2.1. Adaptive Per-Step Skipping

For each diffusion step $k$ in the reverse-time grid, the controller decides whether to compute a fresh prediction or reuse a cached one. The decision is based on a local smoothness test:

• Difficulty Proxy ($d_k$): The mean absolute change in the latent state between consecutive diffusion steps.
• Curvature Check ($\Delta_k$): A second-order finite-difference criterion measuring the local linearity of the latent evolution. If the latent changes are locally affine (low curvature), the step is deemed predictable.
• Skip Decision: If $\Delta_k \leq \theta \cdot d_{k+1}$ (where $\theta$ is a tolerance threshold), the step is skipped, and the cached prediction is reused. Otherwise, a full forward pass is computed.

Safety Guards: To prevent error accumulation and degenerate behavior, three guards are enforced:

1. Warm-up ($W$): Mandatory computation for the first $W$ steps.
2. Consecutive Skip Cap ($C_{max}$): Limits the number of consecutive skips to prevent drift.
3. Stall Threshold ($\epsilon$): Forces computation if the latent change is numerically tiny (near zero), preventing infinite skipping on stagnant states.

2.2. Cross-Modal Coordination

Individually skipping branches can desynchronize appearance and motion. DISK introduces a unidirectional safety gate:

• The Trajectory Controller monitors its own safety guards. If the trajectory branch encounters a complex maneuver (triggering a safety condition), a signal $\sigma^{(t)}_k$ is raised.
• Unidirectional Coupling: If $\sigma^{(t)}_k = 1$, the Vision branch is forced to compute, regardless of its own smoothness test. This ensures that complex ego-motions are rendered with corresponding visual changes.
• Asymmetry: Visual complexity does not force trajectory computation. This prioritizes throughput, as the trajectory branch is typically faster, while ensuring the slower video branch does not lag behind critical maneuvers.

2.3. Autoregressive Stability

To stabilize long-horizon rollouts, DISK propagates controller statistics through the autoregressive chain. After each rollout step $s$, summaries (compute/skip counts and safety fractions) are passed to the next step. A cross-modal difficulty score ($\beta_s$) is computed from the maximum safety fraction of both branches. This score dynamically adjusts the warm-up window and the skip threshold for the subsequent step, expanding the compute window after difficult maneuvers and relaxing it once the scene stabilizes.

3. Key Contributions

1. Training-Free Adaptive Framework: A method that orchestrates visual and trajectory DiTs using adaptive per-step skip decisions, reducing computational cost without retraining or altering model weights.
2. Cross-Modal Skip Coordination: A mechanism that harmonizes the vision and trajectory branches via a unidirectional safety gate and the propagation of controller statistics, ensuring motion-appearance consistency during long-horizon rollouts.
3. Preservation of Chain-of-Forward Self-Conditioning: The method explicitly targets the challenges of closed-loop driving, maintaining the integrity of the autoregressive loop while skipping steps.

4. Experimental Results

The method was evaluated on 1500 samples from the NuPlan and NuScenes datasets using an NVIDIA L40S GPU, integrated into the Epona world model (a 2.5B parameter model).

4.1. Inference Speed

• Trajectory Branch (TrajDiT): Achieved a ~2× speedup (reduction from ~20ms to ~10ms per step).
• Vision Branch (VisDiT): Achieved a ~1.6× speedup (reduction from ~244ms to ~155ms per step).
• Overall: Significant reduction in total per-step latency while maintaining the same model architecture.

4.2. Planning Performance

• Open-Loop: DISK maintained L2 displacement errors nearly identical to the baseline (e.g., L2@3s: 0.3007 vs. 0.3011).
• Closed-Loop (NAVSIM): DISK achieved competitive results on the Predictive Driver Model Score (PDMS).
  • Baseline PDMS: 85.4
  • DISK PDMS: 83.6
  • Sub-metrics (NC, DAC, TTC, Comf, EP) showed only modest drops, indicating preserved safety and drivability.

4.3. Visual Quality

• Metrics: FID increased slightly from 7.5 to 12.5; FVD from 82.8 to 88.1.
• Qualitative: Despite the metric shift, visual comparisons showed the generated videos were "visually indistinguishable" from the baseline, with high temporal consistency.

4.4. Ablation Studies

• Threshold ($\theta$): $\theta=0.01$ was identified as the optimal balance between speedup (~2.03×) and quality.
• Max Consecutive Skips ($C_{max}$): $C_{max}=4$ provided the best trade-off. Higher values (e.g., 8) increased speed but degraded visual stability (FID 12.65).
• Warm-up ($W$): $W=3$ was sufficient; $W=0$ offered similar metrics but risked early instability, while $W=5$ reduced efficiency without quality gains.

5. Significance and Claims

The paper claims that DISK demonstrates practical long-horizon video-and-trajectory prediction at substantially reduced cost. Its primary significance lies in:

• Plug-and-Play Efficiency: It offers a way to accelerate existing foundation models without the prohibitive cost of retraining or distillation.
• Dynamic Adaptation: Unlike fixed-step methods, DISK adapts compute allocation to scene complexity, saving resources on simple frames while preserving precision during critical maneuvers.
• Closed-Loop Viability: It successfully maintains the stability required for closed-loop driving, where error accumulation is a critical failure mode, by coordinating cross-modal branches and propagating difficulty statistics.

The authors conclude that while there is a minor trade-off in visual metrics (FID/FVD), the substantial speedup and preserved planning accuracy make DISK a viable solution for reducing the energy and latency budgets of real-time autonomous driving simulations and planning pipelines.