MPQ: A Message-Passing View of Post-Training Quantization

AI summary

Overview

Research area: Post-training quantization (PTQ) of large language models, specifically the discrete optimization ("solver") stage that assigns weights to low-bit codewords, analyzed through the lens of message passing and posterior inference.

Technical level: Advanced. The paper combines vector approximate message passing (VAMP), state-evolution analysis, random matrix / Hessian spectrum arguments, and large-scale LLM benchmarking; the math is dense even though the core intuition is simple.

Scope: The paper reframes row-wise PTQ as posterior inference over a discrete codebook, shows that GPTQ, LDLQ, and block coordinate descent are zero-temperature hard-decision limits of that view, and introduces MPQ, a two-stage finite-temperature solver evaluated on 2-bit lattice and 3-/4-bit scalar quantization across Llama, Qwen, and Mistral models.

What This Paper Is About

Quantizing a language model to 2–4 bits requires choosing, for each block of weights, one codeword from a small discrete set — and that choice for one block changes the error seen by every other block. Existing solvers (GPTQ, LDLQ, block coordinate descent) pick one codeword per block immediately and then greedily compensate, so an early bad choice can never be revisited. The authors ask what happens if the solver instead keeps a soft distribution over plausible codewords for all blocks, lets those distributions influence each other, and only commits to hard codewords at the very end.

Key Contributions

1. A unifying probabilistic view of PTQ solvers. The paper casts row-wise PTQ as MAP inference under a finite-temperature Gibbs posterior with a Gaussian likelihood and codebook-supported prior, and shows that GPTQ, LDLQ, and block coordinate descent correspond to the zero-temperature (β → ∞) hard-decision limit of this posterior.

2. The MPQ solver. A two-stage method: (i) finite-temperature VAMP, implemented in rank space via a compact SVD, computes approximate blockwise posteriors using a blockwise discrete denoiser plus a Gaussian linear module, with an empirical-Bayes inverse temperature adapted across iterations; (ii) the posterior is projected to a hard assignment by blockwise marginal MAP and refined by exact block coordinate descent.

3. Theoretical analysis of the warm start and of preprocessing. Theorem 5.1 bounds the per-block MAP projection error probability by a sum of exponentials in the fixed-point denoiser precision γ₁*, and Proposition 5.2 gives an exact derivative/integral formula for how diagonal preprocessing (such as AWQ) changes the linear-module error term E_lin.

4. Empirical evaluation across model families and bit widths. Perplexity and zero-shot results on Llama, Qwen, and Mistral models at 2-bit E8P12 lattice and 3-/4-bit scalar quantization, plus solver-time, memory, and mechanism diagnostics.

Main Findings

• Gains grow as bits shrink. On Llama3.2-1B 4-bit scalar, MPQ reaches 10.38 WikiText-2 / 13.63 C4 versus 10.43 / 13.68 for OJBKQ and a BF16 reference of 9.75 / 12.77. At 3-bit scalar the same model gives 13.65 / 17.42 for MPQ versus 14.26 / 18.00 for OJBKQ and 15.45 / 19.25 for GPTQ. At 2-bit E8P12, MPQ gives 31.67 / 31.48 versus 41.79 / 42.69 for LDLQ+CD and 39.28 / 39.72 for OJBKQ.

• The advantage persists at larger scale, but the absolute gap narrows. On Llama3-8B at 2-bit E8P12, MPQ reaches 11.63 / 14.77 versus 12.76 / 15.96 for LDLQ+CD. On Llama2-70B at 2-bit, MPQ is 4.49 / 6.30 versus 4.58 / 6.36 for LDLQ and 4.59 / 6.36 for LDLQ+CD.

• Refinement budget is matched, so the gain is not just extra polishing. The baselines GPTQ+CD and LDLQ+CD receive the same three CD passes as MPQ; MPQ remains ahead. The authors report that in the 2-bit lattice regime the hard projection can start from a higher objective than the greedy initialization yet reach a lower objective after identical CD refinement.

• Zero-shot accuracy improves too. Averaged over five tasks, MPQ beats the matched CD-polished greedy baseline in every model and bit width, with gains of 0.4 to 1.7 percentage points. Against QuantEase and OJBKQ, MPQ is best or tied best in nine of ten model/bit-width cells. The single exception is Mistral 7B at 3-bit scalar, where QuantEase leads by 0.2 points and MPQ ties OJBKQ for second. At 2-bit E8P12, MPQ is best for all five models, improving over LDLQ+CD by 0.4 to 1.4 points.

• Cost profile depends on the bit width. On a single RTX 5090, in the 3-bit scalar regime MPQ matches GPTQ+CD3 across four Llama scales (1B to 13B) but is several times slower than plain GPTQ. In the 2-bit E8P12 regime MPQ stays within 3× of LDLQ and is faster than the strongest baseline OJBKQ across all scales.

• VAMP alone suffices at 3 bits; CD is the dial at 2 bits. For 3-bit scalar on Llama-3-8B, 4 VAMP iterations with no CD already reach 7.8 WikiText-2 PPL, below both OJBKQ and QuantEase, and extra CD sweeps barely move the curve. For 2-bit E8P12, VAMP-only configurations clip above 16 PPL, but 16 VAMP iterations with three CD passes reach around 11.6 PPL versus 12.76 for LDLQ+CD3.

• Matched-time comparisons favor MPQ at 2 bits. On Llama-3-8B, 8 VAMP iterations plus 3 CD sweeps give 12.370 WikiText-2 PPL in 2134 seconds, versus 12.756 in 2076 seconds for LDLQ+CD3. At 16 VAMP iterations plus 3 CD sweeps, MPQ reaches 11.630 in 2951 seconds, versus 12.728 in 3057 seconds for LDLQ+CD6.

• Mechanism diagnostics match the theory's prediction. The measured terminal VAMP precision proxy is high in the 3-bit scalar regime with negligible CD cleanup, much lower in the 2-bit scalar regime with a far higher cleanup ratio, and lower still for the dense 2-bit E8P12 codebook with a larger fraction of blocks changed by CD — consistent with Theorem 5.1 being informative when codewords are well separated and uninformative when many near-tied competitors exist.

• AWQ moves the linear proxy in the predicted direction. On Meta-Llama-3-8B under the AWQ+RHT scalar setup, E_lin decreases for most scale groups, with the few increases concentrated in qkv and gate_up; the sign of the leading-order integrand from Proposition 5.2 agrees with the observed change on nearly all groups.

• Random Hadamard transforms leave the analysis intact. Because E_lin depends on the Hessian only through its eigenvalues, conjugating by the RHT leaves E_lin unchanged, so the AWQ analysis applies to the AWQ+RHT pipeline.

Methodology in Plain English

The layer-wise PTQ objective for one output row is a quadratic reconstruction error, but the calibration Hessian XᵀX couples every block to every other block, so the problem is combinatorial over the product of all codebooks.

The authors' move is to soften it. Instead of searching for the single best discrete assignment, they define a probability distribution over assignments in which better-fitting assignments get exponentially more weight, at an inverse temperature β. As β goes to infinity this distribution collapses onto the exact minimizers of the original problem — which is precisely what greedy solvers approximate when they lock in one codeword per block without looking ahead.

MPQ then runs the softened problem at finite temperature. A VAMP loop alternates between two modules: a denoiser that, for each block independently, converts a Gaussian "cavity" into a full probability distribution over that block's codewords (with a posterior mean and covariance), and a linear module that solves a regularized least-squares subproblem and passes the result back. Because the data only enters through the rank-r SVD factors of the calibration matrix, everything runs in a compressed rank space. The inverse temperature itself is adapted across iterations by an empirical-Bayes rule.

After a fixed number of iterations, the soft posteriors are collapsed by taking each block's marginal mode, producing a hard starting point. That starting point is then polished by exact block coordinate descent on the original objective: for each block, try every codeword, keep the best, update the residual, and repeat for a bounded number of sweeps with early stopping. The VAMP stage is not doing the fine-grained work — it is choosing a better basin for local search to start in.

Empirically, calibration uses 128 C4 training-set sequences of length 2048, activations stay at 16 bits, and evaluation covers WikiText-2 and C4 perplexity plus zero-shot normalized accuracy on ARC-Challenge, ARC-Easy, HellaSwag, PIQA, and Winogrande via the lm-eval harness. Scalar experiments use group size 128 with activation-aware scaling followed by RHT; 2-bit experiments use the E8P12 codebook (d = 8, M_b = 2^16) with RHT and LDLQ as the matched greedy baseline. Defaults are 16 VAMP iterations and 3 CD passes, selected from a pilot sweep that chose (T_vamp, damping, τ) = (16, 0.5, 0.2).

Why This Matters

The solver stage of PTQ has advanced incrementally while most recent progress has gone into preprocessing, codebooks, and objective design. This paper argues that the discrete solver is an underexplored axis and gives a principled reason why: early commitment is structurally wrong when many codewords are near-tied and the Hessian couples blocks strongly. That reframing is useful independent of the specific algorithm, because it explains when soft inference should help (low bits, dense codebooks) and when it should not (4-bit, where baselines are already near the BF16 ceiling).

Real-world applications implied by the paper's framing:

• On-device and personal-device deployment, where the paper motivates PTQ by the need to run models under tight memory and compute budgets on consumer hardware.
• 2-bit weight-only serving, the regime where the reported perplexity and accuracy gaps are largest and where greedy baselines degrade sharply.
• Drop-in acceleration of existing pipelines, since MPQ quantizes weights offline into GPTQ-compatible scalar or QuIP# E8P12 formats and uses the corresponding inference kernels with matched grouping and packing.
• Calibration-constrained settings, since all solver-side cost is incurred offline during PTQ and does not affect inference latency or memory.

Industry relevance: the method targets exactly the bit widths that decide whether a model fits a given memory budget, composes with existing preprocessing (AWQ, RHT) rather than replacing it, and produces standard export formats — all of which lower the barrier to adoption. The cost is offline solver time, which the paper quantifies rather than hides.

Future Directions

• Extending beyond row-wise, layer-wise PTQ. The analysis is per output row and per layer; the authors' related-work discussion notes that inter-layer error propagation methods (YAQA, Qronos) address a different axis. Combining soft inference with cross-layer objectives is a natural next step.

• Tighter theory at finite temperature. The analysis is explicitly a qualitative surrogate: it uses a right-orthogonally-invariant model, the unanchored τ → ∞ limit, and idealized uniform codebook priors, while MPQ runs at finite τ on a finite RHT-conditioned Hessian. Closing that gap is open.

• Better handling of dense codebooks. Theorem 5.1's bound becomes uninformative when many near-tied competitors exist, and the diagnostics show the E8P12 regime leans heavily on CD cleanup. Improving the message-passing stage for dense lattices — rather than delegating to local search — is the clear weak point.

• Reducing solver cost at 3 bits. VAMP alone already saturates scalar quality at 3 bits, which suggests the CD stage could be dropped or made adaptive; conversely, MPQ remains several times slower than plain GPTQ in that regime, so scheduling and convergence-speed improvements matter for practical adoption.

Target Audience

Researchers and engineers working on LLM compression and efficient inference, especially those already familiar with GPTQ, QuIP#, or block coordinate descent and interested in the discrete-optimization side of quantization. The probabilistic reframing and the VAMP/state-evolution machinery will also appeal to readers from signal processing and statistical physics who work on approximate message passing. Readers wanting only the practical takeaway can focus on the tables and the speed–quality tradeoff, but the paper assumes comfort with SVDs, Hessian spectra, and Gaussian posterior arguments for the theoretical sections.

Authors’ abstract

Post-training quantization (PTQ) is a key technique for efficiently deploying large language models, as it compresses weights into low-bit representations without retraining. Existing PTQ solvers, including GPTQ, LDLQ, and block coordinate descent, commit each weight block to a single codeword before its neighbors are considered, leaving near-optimal alternatives unexplored. In this work, we cast row-wise PTQ as posterior inference over discrete codebooks, in which existing solvers correspond to a zero-temperature, hard-decision limit retaining only the top codeword per block. We propose Message-Passing Quantization (MPQ), which preserves soft posterior information across blocks during inference and commits to hard codewords only at the output stage. MPQ first uses vector approximate message passing to compute approximate blockwise posteriors, then refines a hard-decision initialization by exact block coordinate descent. Across 2-bit lattice and 3- and 4-bit scalar quantization on dense language models from Llama, Qwen, and Mistral families, MPQ improves perplexity and zero-shot accuracy over strong hard-decision baselines, and composes favorably with preprocessing techniques such as activation-aware scaling and incoherence processing. Our implementation is available at https://github.com/poseidonchan/MPQ.