PaperlayerView exact revision on arXiv

Source-linked AI summary

LatentQuant: Preserving the Policy-Facing Latent Contract under NVFP4 VAE Quantization

Ziye Deng, Lufang Chen, Shuyu Feng, Zhenwei Duan, Zicong Ye, Yu Sun, Xiaofan Li, Ruyi Gan, Hao Wang, Hao Zhang

arXiv:2610.03959v1cs.AI

TL;DR

World-action-model VAEs must preserve the latents expected by frozen action policies, not only reconstruction fidelity. LatentQuant aligns the quantized encoder to its high-precision reference, freezes that interface, and adapts the decoder; across Wan2.1 and Wan2.2 it retains near-baseline control.

  • Problem

    VAE quantization for world action models must preserve the encoder representation consumed by a frozen policy, a requirement not captured by reconstruction-centric evaluation.

    Source: §1 · p. 1

  • Method

    LatentQuant aligns the NVFP4 encoder to its high-precision reference, then freezes it while adapting only the NVFP4 decoder.

    Source: §1 · p. 2

  • Results

    95.75% versus 95.50% success on LIBERO, and 68.8% versus 69.8% on RoboTwin, show near-baseline control for LatentQuant across Wan2.1 and Wan2.2.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support preserving downstream-consumed representations, rather than reconstructed outputs alone, when quantizing shared generative components.

    Source: §7 · pp. 9–10

    Evaluation covers two WAM stacks, LIBERO and RoboTwin, NVIDIA B300 hardware, and convolutional W4A4; broader models, hardware, and operators remain open.

    Source: §7 · p. 10

Abstract

from arXiv · show

Recent world action models (WAMs) reuse pretrained video VAEs whose encoder latents directly condition downstream action policies. Quantization must therefore preserve not only reconstruction fidelity but also the policy-facing latent contract expected by the frozen policy. Direct NVFP4 leaves W4A4 quantization error uncompensated, whereas joint quantization-aware training (QAT) can recover reconstruction by moving this representation. On Wan2.1, joint QAT nearly matches FP32 VBench-7 (0.7403 versus 0.7409), yet LIBERO success collapses from 95.5% to 10.5%. Controlled decoder-only experiments show that activation quantize-dequantize operations alter the reconstruction signal and decoder Jacobian, redirecting the gradient returned to the encoder and inducing persistent latent drift. Based on this mechanism, we introduce LatentQuant, a two-stage NVFP4 QAT framework that first aligns the quantized encoder with its high-precision counterpart, then freezes it while adapting the decoder. Across Wan2.1 and Wan2.2, LatentQuant preserves near-baseline control and high reconstruction quality, achieving 95.75% success on LIBERO and 68.8% on RoboTwin. On NVIDIA B300 GPUs, NVFP4 execution achieves 1.17x-1.26x end-to-end VAE speedups over BF16 cuDNN.

1 INTRODUCTION

World-action-model VAEs must preserve the encoder representation expected by a frozen action policy, not just reconstruction quality. The paper traces joint-QAT latent drift to decoder activation quantization and introduces a two-stage method that preserves near-baseline control.

1 INTRODUCTION

  • A fixed downstream policy may not be invariant to latent reparameterization, so preserving the encoder-decoder reconstruction alone does not ensure preservation of the representation it consumes.The paper names this compatibility requirement the policy-facing latent contract.Source: §1 · p. 1
  • Joint QAT nearly matches FP32 VBench-7 (0.7403 versus 0.7409), yet LIBERO success falls from 95.5% to 10.5%; latent regularization recovers only 86.8%.Reconstruction fidelity can therefore conceal severe incompatibility with the frozen policy.Source: §1 · p. 2
  • Decoder activation quantization changes the reconstruction signal and decoder Jacobian, redirects gradients returned to the encoder, and consolidates them into persistent latent drift.The failure is a learned representation shift, not merely an instantaneous encoder-output perturbation.Source: §1 · p. 2
  • LatentQuant first aligns the NVFP4 encoder with its high-precision reference, then freezes that interface while adapting only the decoder.This boundary prevents reconstruction gradients from modifying the policy-facing latent.Source: §1 · p. 2, Figure 1 · p. 2
  • 95.75% versus 95.50% success on LIBERO, and 68.8% versus 69.8% on RoboTwin, show LatentQuant retains near-baseline control across Wan2.1 and Wan2.2.The method also retains high reconstruction quality, and the native NVFP4 backend demonstrates end-to-end VAE acceleration on NVIDIA B300 GPUs.Source: §1 · p. 2
  • The paper formulates the policy-facing latent contract for quantizing VAEs consumed by fixed downstream policies, where reconstruction fidelity alone can conceal policy incompatibility.Closed-loop control also imposes a tight per-step latency budget.Source: §1 · pp. 1–2

2 RELATED WORK

Prior work studies quantization of generative-model backbones, VAE efficiency, and downstream representation degradation. This paper focuses on low-bit VAE QAT when a pretrained VAE's latent is consumed by a frozen action policy.

WAMs and policy-consuming latents

  • WAMs consume pretrained video-VAE latents in downstream action policies, making latent preservation stricter when the policy is frozen and cannot adapt after quantization.Related video-VAE work also uses latent alignment for compatibility with a pretrained generator, but this setting fixes the downstream action policy.Source: §2 · p. 3

Low-bit quantization for generative models

  • Low-bit generative-model quantization has focused mainly on Transformer backbones; VAE work is limited, and existing examples target PTQ or compression rather than low-bit weight-and-activation execution.The paper motivates a VAE-specific NVFP4 QAT method for practical low-bit acceleration.Source: §2 · p. 3

Interface-aware quantization

  • Prior interface-aware quantization work addresses intermediate-representation degradation, while this paper studies persistent shifts in a pretrained video VAE's policy-facing encoder representation during joint QAT.The paper frames latent preservation as a representation contract during low-bit VAE adaptation.Source: §2 · p. 3

3 PROBLEM SETTING

The problem setting treats the VAE encoder output as an interface to a fixed action policy, making reconstruction preservation alone insufficient. NVFP4 W4A4 applies QDQ to selected layer operands while leaving specified boundaries and operations in higher precision.

3.1 POLICY-FACING VAE QUANTIZATION

  • The quantized VAE must preserve both reconstruction and compatibility with the representation expected by the fixed downstream policy.The action policy remains frozen, and the paper evaluates the encoder representation against the policy-facing contract.Source: §3.1 · pp. 3–4

3.2 NVFP4 QUANTIZATION

  • The paper studies Blackwell NVFP4 W4A4, applying QDQ to convolutional and linear weights and activations while keeping normalization, nonlinearities, residual additions, and input and latent boundaries higher precision.During QAT, master parameters and optimizer states remain FP32, and the forward pass uses a straight-through estimator for backpropagation.Source: §3.2 · p. 4

4 MECHANISM OF LATENT-CONTRACT VIOLATION UNDER JOINT QAT

Joint QAT can preserve reconstruction while changing the encoder representation. Decoder activation quantization changes the reconstruction signal and backward map, redirecting gradients into encoder updates and producing persistent latent drift.

4 MECHANISM OF LATENT-CONTRACT VIOLATION UNDER JOINT QAT

  • Joint QAT lets encoder and decoder co-adapt to quantization error, modifying the policy-facing latent representation while optimizing reconstruction.The analysis traces this pathway through latent drift, decoder-induced gradient perturbation, and persistent encoder representation change.Source: §4 · p. 4

4.1 PRECISION-DEPENDENT LATENT DRIFT

  • 1.445× SD ratio and 0.714 sample-wise deviation under NVFP4, versus 1.020× and 0.084 in FP32 after 400 matched samples.NVFP4 QDQ also contracts decoder activation magnitudes across most measured sites, but the analysis does not attribute latent drift to shrinkage alone.Source: §4.1 · p. 5

4.2 DECODER QUANTIZATION REDIRECTS LATENT GRADIENTS

  • 67.45° latent-gradient rotation under decoder-activation NVFP4 QDQ, despite only a 1.126× increase in gradient norm.The matched reconstruction loss increases from 0.008834 to 0.017591; restricting QDQ to the activation before the first decoder convolution yields an 11.857% final-reconstruction difference.Source: §4.2 · p. 5

4.3 GRADIENT PERTURBATION LEADS TO PERSISTENT LATENT DRIFT

  • The 1.126× latent-gradient norm ratio becomes a 1.514× encoder-gradient norm ratio, transferring decoder-induced perturbation to trainable encoder updates.The latent-gradient change includes both reconstruction-signal and backward-map changes.Source: §4.3 · p. 5, §4.2 · p. 5
  • With encoder quantization disabled and QDQ removed at evaluation, decoder-QDQ yields master-latent SD ratio 1.04585 versus 1.01127 in the high-precision branch, indicating learned encoder drift.Decoder-side quantization can violate the latent contract when reconstruction optimization is allowed to rewrite the encoder representation consumed by the frozen policy.Source: §4.3 · p. 6

5 FROM MECHANISM TO METHOD: LATENTQUANT

LatentQuant treats reconstruction adaptation as subject to preservation of the established policy-facing latent contract. It first aligns the low-bit encoder to its reference representation, then freezes that encoder while adapting the decoder.

5 FROM MECHANISM TO METHOD: LATENTQUANT

  • LatentQuant replaces joint encoder–decoder adaptation with an aligned low-bit encoder followed by decoder-only adaptation, explicitly blocking further reconstruction-driven encoder updates.This hard optimization boundary contrasts with a joint alignment objective that still subjects encoder coordinates to reconstruction and alignment updates.Source: §5 · p. 6

5.1 STAGE I: SAMPLE-WISE ENCODER ALIGNMENT

  • Stage I updates only the NVFP4 encoder to align each output with the fixed high-precision reference representation expected by the frozen policy.The alignment error is normalized by reference channel-wise standard deviation, and matching is sample-wise rather than distribution-level; the reference encoder requires no action labels or policy gradients.Source: §5.1 · pp. 6–7

5.2 STAGE II: DECODER ADAPTATION UNDER A FIXED CONTRACT

  • Stage II freezes the aligned encoder under deployment-matched NVFP4 execution and confines reconstruction adaptation to the decoder, recovering reconstruction on a fixed latent contract.The decoder observes the quantized latent distribution, while fixed encoder parameters prevent reconstruction gradients from introducing additional learned encoder drift.Source: §5.2 · p. 7

6 EXPERIMENTS

The experiments compare quantization strategies on two frozen-policy WAM stacks, testing reconstruction, interface compatibility, control, physical transfer, and deployment efficiency. Their results show that reconstruction alone does not predict policy compatibility, while LatentQuant preserves the interface and enables native NVFP4 acceleration.

6.1 EXPERIMENTAL SETUP

  • The evaluation uses WALL-WM with Wan2.1 on LIBERO and Wan2.2 on RoboTwin, plus paired replay on fixed physical observations; policies remain frozen across quantization strategies.The comparisons include FP32/BF16 references, Direct NVFP4, Joint QAT, Joint+KL, and LatentQuant; success is measured over 800 LIBERO and 500 RoboTwin trials.Source: §6.1 · p. 7

6.2 RECONSTRUCTION DOES NOT DETERMINE POLICY COMPATIBILITY

  • 95.75% versus 95.50% LIBERO success and 68.80% versus 69.80% RoboTwin success: LatentQuant retains near-reference control while improving Robot VBench-7 over Direct NVFP4 on both stacks.LatentQuant raises Robot VBench-7 from 0.7266 to 0.7340 on Wan2.1 and from 0.7325 to 0.7414 on Wan2.2.Source: Table 1 · p. 8, §6.2 · p. 8
  • On Wan2.1, Joint QAT nearly matches high-precision Robot VBench-7 (0.7403 versus 0.7409) while LIBERO success falls from 95.50% to 10.50%; on Wan2.2, it raises Robot VBench-7 over Direct NVFP4 (0.7402 versus 0.7325) while RoboTwin success falls from 69.80% to 0%.Direct NVFP4 retains 94.38% and 68.00% success despite greater reconstruction degradation, showing the reconstruction–control mismatch across both WAM stacks.Source: §6.2 · pp. 7–8, Table 1 · p. 8

6.3 WHY THE OPTIMIZATION BOUNDARY MATTERS

  • Stage I alignment removes the initial interface mismatch, while freezing the encoder in Stage II prevents reconstruction updates from reintroducing latent drift.The matched soft-alignment sweep shows interface improvements at λ ≥ 0.03, but larger weights reduce held-out PSNR from 29.77 to 29.19 dB, motivating separate alignment and reconstruction stages.Source: §6.3 · p. 8

6.4 PHYSICAL-OBSERVATION TRANSFER AND NATIVE DEPLOYMENT

  • On fixed physical observations, LatentQuant has interface deviations dz = 0.068 and da = 0.036; across 15 real-robot tasks, score differences versus FP32 had median 0, with 7 lower, 2 tied, and 6 higher.Source: §6.4 · p. 8, Figure 3 · p. 9
  • Native NVFP4 execution reduces packed model memory by 60.4% and 69.4%, and full-VAE DRAM traffic by 23.5% and 58.0%, for Wan2.1 and Wan2.2, respectively.Direct NVFP4 and LatentQuant use the same backend, so the reported systems gains come from native execution.Source: §6.4 · pp. 8–9, Figure 4 · p. 9

7 CONCLUSION

The paper concludes that quantizing a WAM VAE must preserve both reconstructed outputs and the latent representation consumed by the frozen policy. LatentQuant separates encoder alignment from decoder adaptation, while evaluation beyond the tested models, hardware, and operators remains open.

7 CONCLUSION

  • Joint-QAT failure is traced to decoder activation QDQ redirecting encoder updates and producing persistent latent drift.LatentQuant aligns the quantized encoder to its high-precision reference, freezes that interface, and then adapts the decoder.Source: §7 · p. 9
  • Evaluation covers two WAM stacks, LIBERO and RoboTwin, an NVIDIA B300, and convolutional W4A4; broader models, hardware, and operators remain open.The conclusion calls for low-bit optimization of shared generative components to preserve downstream-consumed representations, not only reconstructed outputs.Source: §7 · pp. 9–10

AI USE STATEMENT

These appendix materials specify the AI-use statement, NVFP4 mixed-precision and deployment contract, two-stage training, alignment sweep, and reconstruction and interface metrics.

AI USE STATEMENT

  • Generative AI tools assisted literature search, planning, coding, reasoning, figure preparation, drafting, and editing; authors reviewed and verified the outputs.The statement includes cited literature, mathematical claims, code, and reported results.Source: AI USE STATEMENT · p. 10

A.1 OPERAND LAYOUT AND MIXED-PRECISION SCOPE

  • Activations use E4M3 scales per 1×16 channel block, while convolution weights are grouped by output channel and kernel tap; incomplete groups are zero padded.The layout applies to Conv2D, Conv3D, and linear operands, with zero tensors and underflowed blocks handled explicitly.Source: §A.1 · p. 13
  • NVFP4 W4A4 denotes mixed-precision execution: selected convolutional and linear operands use NVFP4 QDQ, while normalization, nonlinearities, residual additions, and latent boundaries remain higher precision.The outgoing latent projection belongs to the encoder scope and the incoming projection to the decoder scope.Source: §A.1 · p. 13

A.2 TRAINING AND DEPLOYMENT STATE

  • Training retains FP32 master parameters and optimizer states while using strict NVFP4 QDQ in the forward graph; deployment stores supported convolution weights packed with scales and packs activations online.Direct NVFP4, joint QAT, and LatentQuant share arithmetic, operand layout, and layer coverage.Source: §A.2 · p. 13

B.1 TWO-STAGE OPTIMIZATION

  • Stage I updates only the encoder using sample-wise latent alignment; Stage II freezes the aligned encoder and updates only the decoder with the reconstruction objective.Runs start from pretrained Wan VAEs with the video backbone and action model frozen; Stage I checkpoints use validation alignment loss and Stage II checkpoints use held-out reconstruction.Source: Table 3 · p. 13, §B.1 · p. 13

B.2 SOFT ALIGNMENT SWEEP

  • In the matched joint-QAT sweep, λ = 0 gives (dz, da) = (1.115, 0.071), while λ ∈ [0.03, 10] gives dz ∈ [0.158, 0.222] and da ∈ [0.023, 0.040].The sweep uses a shared Wan2.1 initialization and 1,000 steps; it exposes the soft-constraint trade-off, not downstream success.Source: §B.2 · p. 14, Table 4 · p. 14, Figure 5 · p. 14

B.3 RECONSTRUCTION AND INTERFACE METRICS

  • VBench-7 is the unweighted mean of seven reconstruction dimensions, while paired dz and da compare low-bit and high-precision latents and standardized actions under matched inputs and frozen-policy state.Paired replay metrics isolate interface compatibility before closed-loop trajectories diverge.Source: §B.3 · pp. 14–15

B.4 DOWNSTREAM EVALUATION PROTOCOLS

The downstream protocols use fixed trial counts for LIBERO and RoboTwin and hold the physical evaluation setup constant across VAE conditions.

B.4 DOWNSTREAM EVALUATION PROTOCOLS

  • LIBERO uses 800 trials per method and RoboTwin 500; physical evaluation fixes the action model, calibration, control settings, resets, and instructions across VAE conditions.Fixed-observation replay is separated from interactive robot evaluation to avoid confounding interface deviations with diverging trajectories; task progress sums milestone points reached before timeout.Source: §B.4 · p. 15

C COMPLETE RECONSTRUCTION RESULTS

Supplementary evaluations compare reconstruction, frozen-policy compatibility, and physical-robot performance, while decoder interventions examine how activation quantization changes the gradient reaching the latent. The two-stage method remains near the high-precision reference on the reported control evaluations.

C COMPLETE RECONSTRUCTION RESULTS

  • Table 5 presents the complete seven-dimension VBench reconstruction profile for high-precision, Direct NVFP4, and LatentQuant settings.Each setting contains 128 reconstructed clips, and the mean is unweighted; the individual dimensions need not all improve for the main claims.Source: Table 5 · p. 15, C COMPLETE RECONSTRUCTION RESULTS · p. 15

D.1 MATCHED INTERFACE REPLAY

  • Across both VAE families, JOINT produces the largest latent displacement and frozen-action-output change, while OURS has the smallest low-bit dz and da.OURS also has lower deviations than DIRECT; ordinary relative ℓ2 differences show the same ordering.Source: §D.1 · p. 15, Table 6 · p. 16

D.2 LIBERO SUITE BREAKDOWN

  • Across all four LIBERO suites, joint-QAT failure is not confined to one suite, while OURS remains near the high-precision reference.Source: §D.2 · p. 15, Table 7 · p. 16

D.3 ROBOTWIN TASK BREAKDOWN

  • On RoboTwin's 50 tasks, HP scored 349/500 (69.8%), DIRECT 340/500 (68.0%), and OURS 344/500 (68.8%).Each task used 10 rollouts under the same task list, episode budget, and success evaluator.Source: §D.3 · p. 16, Table 8 · p. 17

D.4 OPEN-LOOP ACTION TRACES

  • Across five paired Wan2.2 events, high-precision and NVFP4 action L1 errors were 0.121864 and 0.121618, with mean prediction cosine 0.995265.This diagnostic measures paired action compatibility, not closed-loop success.Source: Figure 6 · p. 16, §D.4 · p. 16

D.5 PHYSICAL-EVALUATION BREAKDOWN

  • Across 15 task-wise three-round means, LatentQuant was lower on 7 tasks, higher on 6, and unchanged on 2; the authors interpret physical results as parity.The median paired difference was 0.00 points and the mean paired difference was +0.43 points.Source: §D.5 · p. 17

E.1 LAYERWISE DECODER TRANSPORT

  • The layerwise decoder sweep uses signed adjacent-branch increments that sum to the full-FP32 versus full-QDQ difference, retaining interactions rather than treating sites independently.Full decoder QDQ yields an 11.04% reconstruction difference; latent-gradient cosine is 0.296 and its norm ratio is 1.375× versus FP32.Source: E ADDITIONAL MECHANISM ANALYSIS · p. 17, §E.1 · p. 18, Figure 7 · p. 19

E.2 ONE-SITE TRANSPORT AND REVERSE-REFERENCE DECOMPOSITION

  • The one-site intervention quantizes only the activation entering the first decoder convolution, with FP32 weights and encoder activations and a frozen decoder.The component allocation between loss-gradient transport and backward-map shift depends on reference choice, so the causal claim relies on the total differential signal.Source: §E.2 · p. 18, Table 10 · p. 19

E.3 MATCHED TRAINING TRAJECTORIES

Matched training diagnostics show persistent latent divergence, while a temporal-length control checks reconstruction beyond the 17-frame training clips. The native NVFP4 backend accelerates end-to-end VAE execution, and the appendix also details stage-level performance, resource accounting, and correctness safeguards.

E.3 MATCHED TRAINING TRAJECTORIES

  • Matched FP32 and strict-NVFP4 trajectories diverge persistently in latent standard deviation and sample-wise drift, even when the instantaneous radial-force proxy fluctuates around zero.The directional projection captures 58.8% of paired-difference energy and is a visualization, not a causal-effect estimate.Source: §E.3 · p. 19, Figure 8 · p. 20, Figure 9 · p. 20
  • Channel-wise diagnostics show latent growth distributed across channels, while gradient-norm and RMS-directed-force maps vary over time and are not independent per-channel causal effects.Source: §E.3 · pp. 19–20, Figure 10 · p. 21

E.4 TEMPORAL-LENGTH CONTROL

  • From 17 to 65 frames, the NVFP4–high-precision PSNR gap widens by only 0.273 dB for Wan2.1 and 0.108 dB for Wan2.2.The paired VBench diagnostic shows no systematic widening; this checks reconstruction across input lengths, not error accumulation across sequential task chains.Source: §E.4 · pp. 20–21

F.1 END-TO-END TIMING

  • 56.499 versus 47.629 ms for Wan2.1 and 65.755 versus 51.499 ms for Wan2.2 are the headline 17-frame encode–decode latencies on NVIDIA B300.The values come from dedicated clean-device repeats at the same 1800 MHz clock; weights are packed offline, while activation scaling and packing remain in the measured path.Source: §F.1 · p. 21

F.2 STAGE AND OPERATOR BREAKDOWN

  • Both encoder and decoder stages are faster at every tested temporal length, with the larger gain in the decoder.The operator diagnostics report 2.044× Wan2.1 Conv2D-region speedup, 1.305× Wan2.1 Conv3D-region speedup, and 1.458× call-count-weighted Wan2.2 convolution speedup; 169 of 188 Wan2.2 workloads are faster.Source: Figure 11 · p. 22, §F.2 · p. 22

F.3 RESIDENT STATE AND DRAM TRAFFIC

  • Resident-state counts exclude shape-plan activation and output scratch, while DRAM traffic measures transferred volume across one encode–decode graph rather than physical bandwidth.The ideal packed payload is 4.5 bits per value, distinct from measured resident-state reduction because higher-precision boundaries and unsupported operators remain.Source: §F.3 · p. 22

F.4 CORRECTNESS AND FALLBACK POLICY

  • Native convolution dispatch is enabled only for validated shapes; unsupported workloads use an explicit higher-precision fallback.Measured latent and reconstruction differences remain within configured numerical tolerances in the reported correctness checks.Source: §F.4 · pp. 22–23

Related papers

Loading 2610.03959v1…