OJBKQ: Objective-Joint Babai-Klein Quantization

1 Introduction

Large Language Models (LLMs) [Achiam et al., 2023; Touvron et al., 2023; Comanici et al., 2025] pose substantial deployment challenges due to their massive memory footprint and inference costs [Chen et al., 2023]. Layerwise post-training quantization (PTQ) [Frantar et al., 2023; Lin et al., 2024; Chee et al., 2024] has become a widely adopted compression standard, enabling efficient deployment by sequentially quantizing parameters without retraining. In this work, we focus on standard and practically relevant 4bit and 3-bit settings (e.g., group size 128) and study how to further improve quantization quality beyond strong PTQ baselines.

Despite its empirical success, layer-wise PTQ is a discrete optimization problem. At its core, each layer-wise step tries to solve an integer quadratic programming problem, typically handled using sequential rounding-type solvers. Recent theoretical analyses [Chen et al., 2025] reveal that standard solvers (e.g., GPTQ) admit a lattice-decoding interpretation: they are closely related to Babai's nearest-plane algorithm [Babai, 1986], which finds the Babai point, a suboptimal solution to the closest vector problem (CVP), also known as the ordinary integer least-squares (OILS) problem. More precisely, the underlying optimization problem in PTQ is a box-constrained integer least-squares (BILS) problem. However, when the lattice basis matrix is ill-conditioned or when the dimension is large, the residual associated with the Babai point may be significantly larger than the optimal residual. Thus, it is desirable to find a better sub-optimal solution. Furthermore, minimizing a layer-local proxy loss does not necessarily lead to better end-to-end quantization due to error propagation and input distribution drift. Consequently, the selection objective used to compare and choose among candidate solutions is as crucial as the solver itself, yet remains inconsistent across existing methods.

∗A full version with appendix is available at https://arxiv.org/pdf/ 2602.08376

‡Corresponding author.

In this paper we propose OJBKQ (Objective-Joint BabaiKlein Quantization with (K+1)-Best Sampling), a unified framework that addresses both issues. (i) Joint Target Alignment (JTA): To better align layer-wise decisions with end-toend behavior under error propagation, we introduce a unified candidate scoring objective (Eq. (7)) with a single continuous knob. JTA interpolates between two fundamental alignment targets: matching the runtime quantized activations induced by the partially-quantized network, and matching the corresponding full-precision reference outputs [Arai and Ichikawa, 2025]. The formulated quantization problems at each layer are multiple-right-hand-side BILS problems. (ii) RandomK quantization: For each column of the weight matrix in a BILS problem, we apply the box-constrained Babai algorithm to obtain one sub-optimal solution (a low-bit vector candidate). For the same column, we then extend Klein's Ktime randomized Babai algorithm [Klein, 2000] to the boxconstrained case to generate K additional sub-optimal solutions independently. Among these K + 1 candidates, we select the one which has the minimum residual, referred to as the best Babai-Klein point. There are two levels of parallelization. One across the columns of the weight matrix and the other across the K points produced by the boxconstrained Klein algorithm. Both can be implemented efficiently.

In OJBKQ, Random-K expands the quality of the candidate set, while JTA improves the selection of the optimal candidate.

Our contributions are summarized as follows:

• Quantization Formulation. We revisit layer-wise PTQ from a CVP perspective and reformulate it as an explicit lattice-decoding problem. This enables a faithful Babaistyle implementation that operates via triangular factors and back-substitution, without materializing any matrix inverse.
• Superior Sub-optimal Solution. We extend Klein's randomized algorithm to the box-constrained case and apply it to PTQ. By combining the Babai solution with K parallel runs of the extended Klein randomized Babai algorithm, we generate K + 1 integer candidates and select the one with the smallest residual, yielding low perplexity in both 4-bit and 3-bit configurations. In designing and implementation of our numerical algorithm, we take both accuracy and efficiency into account. For example, unlike GPTQ, no inverse of a matrix is computed. To reduce the overhead of multiple independent solves, we develop a GPU-efficient algorithm presented as Algorithm 2 in the appendix of the full version.
• JTA Candidate Selection. We propose the Joint Target Alignment (JTA) objective, a unified scoring criterion that subsumes commonly used PTQ objectives as special cases by varying the alignment target between runtime-quantized and full-precision references, providing a more fine-grained criterion for candidate selection under error propagation.

2 RelatedWork

Layer-wise post-training quantization (PTQ). Posttraining quantization is a practical standard for compressing large language models without retraining. Following recent PTQ taxonomies [Zhao et al., 2025], mainstream weightonly PTQ methods can be broadly grouped by their dominant mechanism. Compensation-based approaches, exemplified by GPTQ [Frantar et al., 2023], quantize weights sequentially while updating the remaining weights to compensate for accumulated quantization errors, often leveraging secondorder curvature surrogates estimated from calibration data. Rotation-based methods such as QuIP [Chee et al., 2024] apply structured transformations to reshape weight distributions and improve quantizability under low-bit constraints. Salience-based methods, represented by AWQ [Lin et al., 2024], reduce quantization error by selecting scaling factors according to activation/weight importance, avoiding mixedprecision deployment. In addition, optimization-based approaches (e.g., OmniQuant and PoTPTQ) [Shao et al., 2024; Wang et al., 2025] directly optimize quantization parameters with lightweight objectives while keeping pretrained weights frozen. Overall, existing PTQ pipelines mainly differ along two orthogonal axes: (i) the approximate solver used to handle the discrete subproblem, and (ii) the objective used to evaluate candidates and guide selection.

Although it is classified as PTQ in the survey, it still requires training on the calibration dataset to find scales.Then we exclude it from our baselines

Integer least squares and Babai-type solution. CVP or OILS arises in lattice decoding [Agrell et al., 2002], latticebased cryptography [Micciancio and Goldwasser, 2002], and GPS [Teunissen, 1993; Teunissen, 1995; Chang et al., 2005]. BILS mainly arises in MIMO detection [Damen et al., 2003; Chang and Han, 2008]. The most practical approach for OILS is the Schnorr-Euchner algorithm [Schnorr and Euchner, 1994], which enumerates integer points in an ellipsoid to find the optimal solution and has been extended for BILS. The MATLAB package MILES provides solvers for various ILS problems [Chang, 2022]. Since OILS and BILS are NP-hard [Emde-Boas, 1981; Verd´u, 1989], in some applications efficient sub-optimal solutions are often applied. A widely used method is Babai's nearest-plane algorithm [Babai, 1986], which computes the Babai point by back substitution on an upper-triangular linear system, where rounding is performed at each step. Actually, the algorithm is exactly the sizereduction procedure of the LLL algorithm [Lenstra et al., 1982], and it has been extended to the box-constrained case. While computationally efficient, when the coefficient matrix is ill-conditioned, or when the dimension of the integer vector is large, the Babai point may have much larger residual than the optimal solution. There are other sub-optimal algorithms, such as those proposed in [Barbero and Thompson, 2006; Chang et al., 2024], but they are not suitable for GPU computing. Fortunately, there is a randomized version of the Babai algorithm proposed in [Klein, 2000]. In each step of the Babai algorithm, instead of rounding the real solution to the nearest integer, Klein's algorithm rounds it to a nearby integer according to a predefined distribution. When the algorithm is run K times, it generates K integer points. Then the best one, which gives the minimum residual, is chosen as the final sub-optimal solution. This algorithm is for the unconstrained case and has not yet been widely used. Recently, the connection between sequential PTQ solvers and the Babai point has gained attention [Chen et al., 2025]. Some similar work can be found in the GPS literature [Teunissen, 2003; Chang et al., 2005]. This paper emphasizes that the PTQ problems can be modeled as BILS problems so that suboptimal BILS algorithms can be applied.

Objectives under error propagation and distribution drift. A recurring challenge in layer-wise PTQ is that minimizing a layer-local reconstruction proxy does not necessarily translate to improved end-to-end quality, due to quantization error propagation and the resulting input distribution drift across layers. Consequently, the choice of the alignment target-whether to match outputs under runtime (partially-quantized) activations or to match full-precision reference outputs-plays a central role in practical PTQ performance. Recent work revisits this issue from an errorpropagation viewpoint and advocates objectives that better reflect downstream behavior under partial quantization [Arai and Ichikawa, 2025]. Our JTA objective follows this perspective but provides a unified scoring criterion with a single continuous knob that interpolates between runtime-quantized activations and full-precision reference targets, enabling consistent candidate selection across layers.

3 Methodology

3.1 Layer-wise Objectives and Joint Target Alignment (JTA)

We now analyze commonly used layer-wise PTQ objectives. Consider a linear layer with full-precision weight matrix W ∈ Rm×n . Let X ∈ Rp×m be full-precision calibration activations, ˜X ∈ Rp×m be the corresponding runtime activations produced by a partially quantized network (i.e., upstream layers already quantized), Yfp := XW denote the full-precision reference output, and Yrt := ˜XW denote the runtime-consistent reference output under partially quantized activations. At runtime, a candidate ̂W produces the output ̂Y := ˜X̂W.

Many existing objectives can be interpreted as choosing different alignment targets:

In these optimization problems, each column of ̂W takes values in a finite subset of a translated lattice. Specifically, ̂W has a special form; see Eq. (9). Later we will see these are in fact BILS problems. Eq. (1) directly reflects runtime behavior but aligns to a target computed on quantized activations, which may be noisy; this objective is adopted by GPTQ and QUIP. All three formulations above admit a convenient least-squares structure in their continuous relaxation (i.e., before the integrality constraints are imposed), but suffers from train-run mismatch when earlier layers are quantized, and is optimized by AWQ. Eq. (4) partially alleviates this mismatch by evaluating candidates on runtime activations while preserving the full-precision reference; QEP [Arai and Ichikawa, 2026] employs this objective as a corrective patch within a layer-wise quantization pipeline, although it may over-penalize candidates when error drift accumulates across layers. These trade-offs motivate a unified selection objective that interpolates between runtime-consistent and full-precision references, while optionally controlling weight drift.

We propose an expressive Joint Target Alignment (JTA) objective with a single continuous knob to score and select candidates generated by a random decoding. We define an interpolated target

and formulate the optimization problem:

The knob µ interpolates between matching full-precision and runtime-consistent references, while λ optionally regularizes weight displacement. In our framework, Random-(K + 1) quantization (Sec. 3.4) finds a sub-optimal solution to (7), and JTA (Eq. (7)) provides a consistent curvature-aware criterion for selecting the best candidate under error propagation. When µ=1 and λ=0, Eq. (7) reduces to the runtimeconsistent objective in Eq. (1); when µ=0 and λ=0, it aligns to the full-precision reference in Eq. (4).

End-to-end layer-wise procedure. For each layer, we (i) form ˜X by a partially quantized forward pass, (ii) find a suboptimal optimal solution to Eq. (7), to be seen in the next sections. This closes the loop between discrete candidate generation and objective-aligned selection.

3.2 FromLayer-wise Quantization to Integer Least Squares

From (7) we obtain

The task of finding an optimal quantized weight ̂W to minimize the selection objective in Sec. 3.1 can be formally mapped to the classical Integer Least Squares (ILS) problem.

For the derivation, we first introduce additional notation. Let

denote the box constraint corresponding to the set of admissible integer values for a wbit-bit quantized weight. Let

denote the quantized version of the full-precision weight matrixW. We further define

as the scale matrix and zero-point matrix, respectively. Both S and Z, which have some structures (each sj and each zj have several groups and the elements in the same group have the same value), are pre-computed using standard statistical calibration methods (e.g., absmax calibration).

Finally, the dequantized weight matrix is given by

where ⊙ denotes element-wise multiplication.

To simplify notation in Eq. (8), we further define

It is easy to see that (8) can be rewritten as

where Dj := diag(sj) ∈ Rm×m.

As a result, the optimization problem in Eq. (7) decomposes into n independent BILS problems:

Since each subproblem is independent, we simplify notation by focusing on a single per-column subproblem:

However, for computational efficiency, we need to keep in mind that A is shared across the n BILS problems, while D depends on each individual.

3.3 Babai-Based Quantization

The standard way to compute the box-constrained Babai point, a sub-optimal solution to (11), is via the QR factorization of AD. However, in the quantization application, the dimensions of A are large, so we instead use the Cholesky factorization to reduce the computational cost.

Let ¯q ∈ Rm denote the unconstrained real least-squares solution of (11). The corresponding normal equations are given by

or equivalently

Denote v := D(¯q - z). To obtain v from (12), we first compute the Cholesky factorization A⊤A = R⊤R, where R is upper triangular. Then we solve two triangular linear systems R⊤u = A⊤t and Rv = u.

After v is founded, we obtain the real least-squares solution of (11):

where ⊘ denotes element-wise division.

Denote ¯R := RD. It is not difficult to verify that the BILS problem (11) is equivalent to

Then, we can find the optimal solution or some sub-optimal solutions.

The algorithm which finds the box-constrained Babai point of (13) is described in Algorithm 1. In the algorithm,

i.e., the point in the interval [u, v] closest to x, Moreover, c(i) is the real-valued quantity computed for q(i) at step i, assuming that q(m), q(m-1), . . . , q(i+1) have been fixed during back-substitution, and q(i) is then chosen as the integer in the constraint B closest to c(i).

3.4 (K+1)-Best Babai-Klein-Based Quantization

Babai's nearest plane algorithm is deterministic and makes greedy local decisions during back-substitution, so an early incorrect decision cannot be corrected later. Consequently, the algorithm may return a lattice point far from the closest

Algorithm 1 Babai quantization with JTA

• 1: Compute s and z
• 2: Compute the Cholesky factorization: ˜X⊤ ˜X + λ2 I = R⊤R
• 3: Solve the lower triangular systemR⊤u = A⊤t and the upper triangular system Rv = u.
• 4: Compute ¯q = v⊘ s +z
• 5: Compute ¯R=RDwith D= diag(s)
• 6: c(m) = ¯q(m)
• 7: q(m) = clamp(⌊c(m)⌉; 0, 2wbit-1 )
• 8: for i =m-1 : -1 : 1 do
• 9: c(i)=¯q(i)+(∑mj=i+1 ¯R(i, j)(¯q(j)-q(j))) / ¯R(i, i)
• 10: q(i) = clamp(⌊c(i)⌉; 0, 2wbit -1)
• 11: end for

lattice point, especially when the lattice basis matrix is illconditioned (or poorly reduced). The problem becomes more serious when the dimension of the lattice basis matrix is large. To address this limitation, Klein's sampling algorithm [Klein, 2000] introduces controlled randomness into the rounding process, enabling the algorithm to explore a richer set of lattice points while maintaining polynomial-time complexity.

Klein randomized rounding. Like Babai's original algorithm, Klein's original algorithm is for the unconstrained case. However, we can easily extend it to deal with the box constraint in our quantization application. In each backsubstitution step, instead of deterministic rounding such as lines 7 and 10 in Algorithm 1, we sample q(i) from a distribution concentrated around c(i):

where α > 0 controls the sampling temperature. Larger α approaches greedy Babai; smaller α encourages exploration. For this paper, we follow [Liu et al., 2011] design choices where α = ln(ρ)/ min1≤i≤m r2 ii and ρ is the solution of K = (eρ)(2m/ρ) . We refer to the integer vector obtained by the extended Klein algorithm as a box-constrained Klein point.

By executing the extended Klein algorithm K times independently, we obtain K box-constrained Klein points. We then choose, among these K points and the box-constrained Babai point, the one with the minimum residual as a final suboptimal solution to the BILS problem (13). We refer to suboptimal solution as the (K + 1)-best box-constrained BabaiKlein point.

To get some idea about how the value of K affects the residual of the (K + 1)-best box-constrained Babai-Klein point, we give Figure 1, in which K0 means only the boxconstrained Babai is used as the sub-optimal solution. We observe that typically larger K makes the objective smaller.

Parallel-Isolated (K+1)-best Klein-Babai

We propose Parallel-Isolated (K+1)-best Klein-Babai (PIKlein-Babai), a GPU kernel that evaluates K randomized Klein-Babai instances together with one deterministic greedy Babai instance in a single parallel pass, while keeping all K+1 residuals strictly isolated. Each instance owns a private residual slot along a batch dimension, and the lookahead update from already-quantized coordinates to the remaining ones is fused across all K+1 instances into a single batched matrix multiplication. This recovers the throughput of a shared-buffer implementation without its correctness violation. Reserving the extra slot for the deterministic greedy Babai solution guarantees its inclusion among the K+1 candidates, and the final quantization is selected using the unified JTA score in Sec. 3.1. Klein-style JTA quantization, (K+1)- best JTA quantization, and detailed kernel design, buffer layout, and blocked update rules are shown in the full version

4 Experiments

To evaluate the effectiveness of our method, we conduct comprehensive experiments comparing it against commonly used layer-wise quantization baselines. We first describe the baseline methods, the large language models being quantized, as well as the evaluation metrics and datasets. We then present the experimental results, followed by a detailed analysis. To ensure fair comparisons, we adopt the default configurations of each baseline's repo whenever possible.

Baselines We evaluate our method against state-of-the-art quantization approaches, including AWQ [Lin et al., 2024], GPTQ [Chen et al., 2023], and QUIP [Chee et al., 2024], and additionally report round-to-nearest as a na¨ıve baseline for perplexity. Although theortically, the K can be any positive integer number, considering the trade-off between resource usage and performance gain Figure2, we set the K as 5 in our experiments. To ensure strong baselines, we enable activation ordering for GPTQ. Calibration is performed using 128 samples from the C4 [Raffel et al., 2020a] dataset with a sequence length of 2048 tokens. For all baselines, only minimal modifications are made to accommodate the architectural structures of newer LLMs, while their original algorithms and configurations are kept unchanged.

Models We consider twomajor LLMfamilies, Qwen [Yang et al., 2025]and LLaMA [Touvron et al., 2023; Grattafiori et al., 2024], together with Mistral-7B [Jiang et al., 2023]. This selection covers diverse model families, reasoning capabilities, and parameter sizes.

Metrics andDatasets To test the basic text generation ability, we use the perplexity to measure each quantized model on C4 [Raffel et al., 2020a] and Wikitext2 [Raffel et al., 2020b]. Then we further evaluate on the zero shot accuary on ARC [Clark et al., 2018], boolq [Wang et al., 2019], HellaSwag [Zellers et al., 2019], PIQA [Bisk et al., 2020] and WinoGrande [Sakaguchi et al., 2021] For the last part, we test the reasoning ablility on GSM-8K [Cobbe et al., 2021], GPQAdiamond [Rein et al., 2023] and MBPP [Austin et al., 2021]. Except the evalution of PPL, all other tasks are evaluated on LM-harness library [Gao et al., 2024]

Perplexity Table 1 reports perplexity (PPL) under groupwise weight-only quantization with 16-bit activation. At 4bit, GPTQ and AWQ perform reasonably on large LLaMA

models but degrade on harder architectures such as Qwen0.6B and LLaMA-3-8B, while RTN frequently diverges. Our method consistently achieves the lowest or near-lowest perplexity across all model families, with clear gains in challenging regimes. This advantage is more pronounced at 3-bit, where prior methods often suffer sharp drops or catastrophic failures, whereas ours remains stable with smooth degradation relative to BF16. When group quantization is disabled

(group size = 0), baselines exhibit increased quantization error and numerical instability, especially on smaller models, while our method remains stable and preserves low perplexity across all architectures.

Zero-shot accuracy Table 2 compares BF16 and low-bit PTQ methods across three model families under both 4-bit and 3-bit settings. At 4 bits, all methods remain relatively close to BF16, but our approach consistently achieves the strongest or second-strongest performance across most tasks and models, yielding the highest average accuracy for all three backbones. Notably, the gains are more pronounced on reasoning-heavy benchmarks such as ARC-C, ARC-E, and Hella, indicating improved robustness beyond simple lexical or pattern-matching tasks. When moving to the more challenging 3-bit regime, performance gaps widen substantially: GPTQ and QUIP exhibit significant degradation, particularly on ARC and Hella, while our method degrades more gracefully and maintains clear advantages in average performance. This trend is consistent across model scales, from Qwen3-4B to Qwen-3-8B, suggesting that the proposed approach scales favorably and is less sensitive to aggressive quantization. Overall, the results demonstrate that our method not only preserves accuracy under moderate compression but also provides superior stability in extreme low-bit settings, highlighting its effectiveness as a general and robust PTQ solution.

Reasoning Table 3 reports reasoning accuracy on GSM8K, GPQA, and MBPP under 4-bit quantization across three model families. Compared to GPTQ, AWQ, and QUIP, our method consistently achieves the highest average accuracy for all backbones, with particularly strong gains on GSM8K and GPQA, which require multi-step numerical and factual reasoning. While AWQ remains competitive on MBPP for some models, its performance is less consistent across tasks, leading to lower overall averages. QUIP exhibits severe degradation on reasoning benchmarks, including complete failure on MBPP for Qwen models, highlighting its instability under reasoning-intensive workloads. In contrast, our approach maintains balanced performance across arithmetic, knowledge-intensive, and program synthesis tasks, closely tracking BF16 accuracy despite aggressive quantization. These results indicate that the proposed method better preserves structured reasoning capabilities under low-bit constraints, rather than overfitting to a single task type or model family.

Ablations We investigate the impact of the hyperparameter K by varying it within {1, 5, . . . , 50} under the setting of 4-bit Llama-3-8B with a group size of 128. As illustrated in Fig. 2, the perplexity on both WikiText-2 and C4 exhibits a significant drop at K = 5. Beyond this point, while the performance continues to improve up to K = 50, the marginal gains become negligible. Based on these observations, we select K = 5 as the default configuration for our method.

We analyze the impact of µ and λ on WikiText-2 in Fig. 3. For µ, performance is optimized at 0.6 but degrades at the boundaries (0.1, 1.0). This confirms that neither Eq. 1 nor Eq. 4 alone is sufficient; a balanced combination is essential. For λ, despite higher variance, setting λ = 0.6 yields a robust minimum. Consequently, we adopt (µ = 0.6, λ = 0.6) as the default configuration for 3 bits setting. Similarly, we adopt (µ = 0.1, λ = 0.2) for 4 bits setting.

5 Conclusion

We presented a unified post-training quantization framework that reformulates layer-wise weight quantization as a structured integer optimization problem and solves it using a joint objective with Babai-type decoding. By integrating compensation-aware targets and controlled randomness, the proposed method consistently outperforms existing PTQ approaches across a wide range of models, tasks, and bit-widths. Extensive experiments on both general evaluation benchmarks and reasoning-intensive tasks demonstrate that our approach preserves accuracy more effectively under aggressive low-bit quantization, while remaining stable across model scales. These results suggest that principled optimizationbased formulations provide a robust foundation for future low-bit quantization methods, bridging classical lattice decoding techniques and modern large language model compression.

Limitations Our current framework does not yet incorporate weight permutation (e.g., GPTQ) or dynamic scaling, which are promising for future integration. Additionally, we presently use fixed hyperparameters (µ, λ) and candidate size K for all layers. Future work will explore layer-wise adaptive strategies, specifically assigning distinct µ, λ and varying the number of random candidates per layer to better match the sensitivity of different model components.

Contribution

Xinyu Wang and Ziyu Zhao contributed equally.