Model-Preserving Adaptive Rounding

2025-05-29

AI Summary

Post-training quantization (PTQ) of large language models (LLMs) often minimizes local activation errors, leading to distortion in the global output distribution. Method: This paper proposes a PTQ method that directly optimizes the full-model KL divergence between the quantized and original model outputs. Its core innovation is the first incorporation of a Kronecker-factored approximation of the layer-wise Hessian matrix into the adaptive rounding process—providing theoretical grounding for rounding decisions without reliance on specific quantizer designs—and the development of a quantizer-agnostic, model-level KL divergence optimization framework. Results: Evaluated on models ranging from 10B to 100B parameters, the method reduces average KL divergence by approximately 30% and achieves state-of-the-art performance on downstream tasks.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: (Large) Language Models
• Search and Optimization: Learning to Search

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

The main goal of post-training quantization (PTQ) is to produced a compressed model whose output distribution is as close to the original model's as possible. To do this tractably, almost all LLM PTQ algorithms quantize linear layers by independently minimizing the immediate activation error. However, this localized objective ignores the effect of subsequent layers, so reducing it does not necessarily give a closer model. In this work, we introduce Yet Another Quantization Algorithm (YAQA), an adaptive rounding algorithm that uses Kronecker-factored approximations of each linear layer's Hessian with respect to the extit{full model} KL divergence. YAQA consists of two components: Kronecker-factored sketches of the full layerwise Hessian that can be tractably computed for hundred-billion parameter LLMs, and a quantizer-independent rounding algorithm that uses these sketches and comes with theoretical guarantees. Across a wide range of models and quantizers, YAQA empirically reduces the KL divergence to the original model by $approx 30%$ while achieving state of the art performance on downstream tasks.

Problem

• Minimize output distribution gap in post-training quantization
• Address localized error ignoring subsequent layer effects
• Improve KL divergence reduction with adaptive rounding

Innovation

• Adaptive rounding with Kronecker-factored Hessian approximations
• Tractable Kronecker-factored sketches for large models
• Quantizer-independent rounding with theoretical guarantees