Quantized One-Step Generative Modeling via Cell-Margin Quantization-Aware Training

Low-bit quantization and one-step sampling can substantially reduce the cost of flow-matching models, yet existing quantization-aware training methods typically minimize velocity-field reconstruction error, although generation quality is ultimately determined by the terminal distribution after ODE integration. We introduce CM-QAT, which trains a quantized student for one-step sampling with one function evaluation (1-NFE) to preserve the teacher's terminal geometry in a frozen feature space. Specifically, CM-QAT constructs a public nested partition and uses a differentiable prototype-margin objective to encourage each source sample to reach the correct terminal cell. Our counterexamples show that small velocity-$L_2$ error is neither necessary for terminal cell correctness nor sufficient for a distribution-free cell-risk guarantee without boundary-margin control. We derive conditional feature-space $W_1$ bounds from recursive cell-histogram discrepancies and relate cell-crossing probability to endpoint displacement and boundary margins. Experiments on Fashion-MNIST, CIFAR-10, and class-conditioned ImageNet-256 show that CM-QAT outperforms recent low-bit flow-generation methods and enables generation with models quantized to as low as 2-bit & 3-bit weights and sampled in a single step (1-NFE).