🤖 machine learning

Compander-Aligned Query Geometry for Quantized Zeroth-Order Optimization

This paper introduces CAQ-ZO, a zeroth-order optimization method that aligns query geometry with nonuniform companding quantizers by performing perturbations in the transformed latent space to eliminate endpoint-rounding residuals and improve fine-tuning performance on low-bit models.

Original authors: Yao Shu, Zilin Zhu

Published 2026-05-12

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yao Shu, Zilin Zhu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to tune a very complex, high-tech radio (a large AI model) to get the clearest signal possible. You can't see the internal knobs or dials; you can only listen to the sound (the "loss") and guess which way to turn the knobs to make it better. This is called Zeroth-Order Optimization.

Usually, to figure out which way to turn a knob, you wiggle it slightly left and right, listen to the difference in sound, and then move the knob in the direction that improved the sound.

The Problem: The "Low-Bit" Radio

Now, imagine this radio is built with a special, memory-saving constraint: it only understands "low-resolution" settings. It's like the radio has a digital dial with only a few specific "clicks" or "steps" it can land on, rather than a smooth, continuous rotation.

In the old way of doing things (which the paper calls Weight-Space Queries), you would:

Decide to wiggle the knob slightly to the left and right in your mind (continuous movement).
Ask the radio to play the sound at those two new positions.
The Glitch: Because the radio only understands specific "clicks," it snaps your continuous wiggle to the nearest available click.
- If the "clicks" are very close together in some areas (dense), your tiny wiggle might get snapped to the same click on both sides. The radio hears no difference, and you get zero information.
- If the "clicks" are far apart in other areas (sparse), your tiny wiggle might get snapped to a click that is actually very far away from where you intended. You get a distorted signal.

The paper argues that this "snapping" creates a confusing, distorted map for the optimizer, making it hard to learn.

The Solution: CAQ-ZO (The "Map-Aligned" Approach)

The authors, Yao Shu and Zilin Zhu, propose a new way to ask the radio questions. They call it CAQ-ZO (Compander-Aligned Queries).

Instead of thinking about the knobs in the original "continuous" world, they change the perspective. They realize that the radio's "clicks" are actually perfectly evenly spaced if you look at them through a special funnel (a mathematical transformation called a compander).

The Analogy of the Funnel:
Imagine the radio's settings are a bumpy, uneven road.

The Old Way: You try to walk a fixed distance (say, 1 meter) on this bumpy road. Sometimes you walk 1 meter through a deep mud pit (where you barely move), and sometimes you walk 1 meter over a smooth hill (where you fly forward). Your steps are inconsistent.
The CAQ-ZO Way: You look at the road through a special funnel lens. Through this lens, the bumpy road looks like a perfectly flat, straight road with evenly spaced tiles.
- Now, instead of walking 1 meter on the bumpy road, you simply walk one tile on the flat, tiled road.
- You ask the radio: "What does the sound look like if I move exactly one tile to the left? And one tile to the right?"
- Because you are moving exactly from tile to tile, the radio's "snapping" mechanism doesn't distort your movement at all. It lands exactly where you told it to.

What They Found

The Theory: They proved mathematically that if you stick to the "tiles" in the special lens view (the compander coordinate), the measurement you get is perfectly accurate. There is no "snapping error."
The Experiments:
- Synthetic Tests: They created fake math problems to isolate this "snapping" error. They showed that their new method (CAQ-ZO) eliminated the error completely, while the old method kept getting confused by it.
- Real AI Models: They tested this on real Large Language Models (like Qwen and Llama) using a specific low-precision format called NF4. They found that by using their "tile-walking" method, the models learned better and achieved higher accuracy than models using the old "bumpy road" method, even though they used the exact same memory and hardware.

The Bottom Line

The paper doesn't invent a new type of radio or a new way to store data. Instead, it fixes how we ask questions to a low-precision radio.

Old Method: "Wiggle the knob a little bit." (Result: The radio snaps the wiggle to a random spot, giving you bad data).
CAQ-ZO Method: "Move exactly to the next click mark." (Result: The radio hears exactly what you asked, giving you perfect data).

By aligning their questions with the radio's internal "clicks," they made low-memory AI training much more effective.

Technical Summary: Compander-Aligned Query Geometry for Quantized Zeroth-Order Optimization

1. Problem Statement

The paper addresses a fundamental limitation in Zeroth-Order (ZO) Optimization when applied to low-precision (quantized) large language models (LLMs). While ZO methods (like MeZO) enable memory-efficient adaptation by relying solely on forward passes and scalar loss values, existing approaches struggle when the model is stored and queried in low-bit precision (e.g., NF4).

The core issue is a measurement mismatch caused by nonuniform quantization:

Standard ZO: Perturbs weights in the continuous weight space ( $x \pm \mu u$ ) and measures the loss difference.
Quantized ZO: The engine evaluates the rounded endpoints $Q(x \pm \mu u)$ .
The Conflict: For nonuniform companding quantizers (e.g., NF4, $\mu$ $μ$ -law), the quantization cells are not uniform. A fixed weight-space perturbation radius $\mu$ $μ$ may:
1. Collapse: Fall entirely within a single dense quantization cell, resulting in zero loss difference (no gradient signal).
2. Over-span: Cross multiple sparse cells, measuring a nonlocal secant rather than a local gradient.
3. Misalign: Create a "rounded chord" $Q(x+\mu u) - Q(x-\mu u)$ that is not parallel to the intended update direction $u$ .

This creates a query-time quantization-perturbation coupling, introducing a bias channel that distorts the descent direction and prevents convergence to a stationary point, even if the codebook itself is optimal for reconstruction.

2. Methodology: CAQ-ZO

The authors propose CAQ-ZO (Compander-Aligned Queries for Zeroth-Order Optimization), which shifts the query geometry from the weight space to the compander coordinate space.

Theoretical Framework

The paper models scalar nonuniform quantization as a companding process:
$Q = \phi^{-1} \circ U \circ \phi$
Where:

$\phi$ : A monotone coordinate map (compressor).
$U$ : A uniform quantizer with grid spacing $\Delta$ in the transformed space $z = \phi(x)$ .
$\phi^{-1}$ : The expander map back to weight space.

The Algorithm

Instead of perturbing $x$ in weight space, CAQ-ZO operates in the $z$ -space:

Coordinate Transformation: Maintain the iterate in the compander coordinate $z_t = \phi(x_t)$ .
On-Grid Stencil: Generate perturbations directly in $z$ -space using a one-grid-step Rademacher stencil: $z_t \pm \Delta r$ , where $r \in \{-1, +1\}^d$ .
Query: Map endpoints back to weight space: $x_{query} = \phi^{-1}(z_t \pm \Delta r)$ .
Measurement: Query the low-precision oracle. Since the endpoints $z_t \pm \Delta r$ lie exactly on the uniform grid of $U$ , the quantizer $Q$ acts as the identity at the query points (no rounding error occurs at the endpoints).
Update: Compute the gradient estimate in $z$ -space and update $z_{t+1}$ , projecting the result back onto the grid to maintain the invariant for future queries.

Key Distinction: CAQ-ZO does not change the codebook, the optimizer wrapper, or the forward oracle. It only changes the geometry of the finite-difference stencil to align with the quantizer's internal grid.

3. Key Contributions

The paper makes four primary contributions:

Formulation of Query Geometry: It reframes quantized ZO from a codebook comparison problem into a measurement-geometry problem. It isolates the codebook, query coordinate, direction distribution, smoothing radius, and endpoint grid alignment as separate control variables.
Theoretical Analysis of Mismatch:
- Theorem 1: Proves that no fixed weight-space radius $\mu$ can globally calibrate nonuniform cells; the effective grid span $\rho$ depends on the local compander slope $\nabla \phi$ .
- Theorem 2: Decomposes estimator error, showing that standard weight-space queries incur a residual proportional to $\Delta^2/\mu^2$ (endpoint rounding), whereas CAQ-ZO makes this query-time residual exactly zero.
- Theorem 3 & 4: Provides stationarity bounds demonstrating that weight-space queries suffer from a residual floor due to rounding bias, while CAQ-ZO achieves convergence governed only by standard finite-difference bias and stochastic variance.
Algorithm Design: Introduces the CAQ-ZO algorithm (Algorithm 1) which utilizes on-grid Rademacher finite differences in the compander coordinate and enforces update-time projection to preserve grid alignment.
Empirical Validation:
- Synthetic Experiments: Isolates the endpoint-rounding channel, showing CAQ-ZO eliminates the residual channel that plagues weight-space baselines (Vanilla Gaussian ZO, QuZO).
- LLM Fine-Tuning: Demonstrates improvements on Qwen-2.5-1.5B and Llama-2-7B using NF4 quantization. CAQ-ZO outperforms the weight-space baseline (QuZO) across multiple tasks (SST-2, RTE, CB, SQuAD, etc.) under the same quantizer and evaluation budget.

4. Results

Synthetic: In high-dimensional ( $d=10,000$ ) synthetic experiments with $\mu$ -law and NF4-style companders, CAQ-ZO achieves the lowest final-to-start loss gap ratio. It specifically eliminates the "endpoint-rounding residual" observed in weight-space methods, which remains non-zero even with increased iterations.
LLM Fine-Tuning: In matched NF4 fine-tuning experiments:
- CAQ-ZO consistently improves over QuZO (the weight-space baseline) on all metrics for both models.
- Notable gains include 18.8/21.4 accuracy points on CB and 24.0/28.1 F1 on SQuAD for Qwen and Llama respectively.
- The results confirm that changing the finite-difference endpoints (while keeping the codebook fixed) improves the trained quantized model.

5. Significance and Scope

The paper claims that quantized ZO is not determined solely by the codebook; the interaction between the finite-difference stencil and the quantizer cells is equally critical.

Core Insight: A codebook optimal for static reconstruction (e.g., NF4) may be suboptimal for weight-space perturbations if the query geometry is misaligned. CAQ-ZO resolves this by aligning the query geometry with the compander.
Modesty of Claims: The authors explicitly scope their claims:
- The theory applies to scalar monotone companders (including NF4 via monotone interpolation).
- The evidence is limited to controlled synthetic experiments and matched NF4 LLM fine-tuning.
- The paper does not claim universal transferability to all non-scalar codebooks (e.g., vector quantizers) or broader large-model claims without separate validation.
- It is positioned as a specific solution for the "low-precision-forward-pass" setting, distinct from methods that modify the codebook or use straight-through estimators.

In summary, CAQ-ZO provides a theoretically grounded and empirically validated method to remove the bias introduced by quantization rounding in zeroth-order optimization, enabling more effective memory-efficient adaptation of quantized LLMs.