PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning Low-Bit Models with Gradient in Quantized Code Space

Shiguang Wu, Zhouchen Lin, Quanming Yao

arXiv:2608.30908v1cs.LG

TL;DR

Strict fully low-bit fine-tuning must optimize quantization codes and scales while keeping adapted checkpoints deployable in the target representation. GradCodeS provides code-coordinate surrogate gradients to guide candidate search, then selects among valid low-bit states by realized loss. Across reasoning, instruction-following, and structured language-understanding tasks, it consistently improves fully 4-bit adaptation across supported parameterizations and datatypes.

  • Problem

    Strict fully low-bit deployment requires optimizing quantization codes and admissible scales without retaining a high-precision residual adapter or relying on post-training requantization.

  • Method

    GradCodeS derives a code-coordinate surrogate gradient, samples nearby deployable code candidates around its reference, and selects the lowest-loss candidate.

  • Results

    GradCodeS consistently improves fully low-bit adaptation across arithmetic reasoning, instruction following, and structured language understanding, with gains across quantization datatypes and parameterizations.

  • Takeaways & Limitations

    Gradient guidance can make discrete low-bit search more efficient while realized-loss selection preserves deployment faithfulness.

  • Takeaways & Limitations

    The setup leaves parameters outside the targeted quantized linear layers unquantized.

Abstract

from arXiv · show

Fine-tuning Low-bit models aims to adapt a quantized model while keeping the final deployed checkpoint in the same low-bit form. This setting is practically important as it reduces memory and inference cost for storage and deployment. Under this constraint, adaptation becomes an optimization problem over quantization codes and scales. Existing continuous low-bit training is efficient, but it can be distorted by straight through estimation error or by post-quantize gap; discrete search is deployment-faithful, but it is often too inefficient under a finite training budget. We propose code surrogate gradient as the first order signal in deployable code space to acceleate optimization, and performing guided search to preserve deployment faithfulness. Experiments across arithmetic reasoning, instruction following, and structured language understanding show that GradCodes consistently improves fine-tuning low-bit models across different quantization datatypes. Code is provided at https://github.com/ovo67/GradCodes.

1 INTRODUCTION

Fully low-bit fine-tuning must optimize quantization codes and scales while keeping every adapted checkpoint deployable in the target representation. GradCodeS combines code-space gradient guidance with loss-based candidate selection, and is reported to improve fully low-bit adaptation across tasks and datatypes.

  • Strict fully low-bit deployment keeps targeted weights quantized, permits datatype-specific scale metadata, and disallows high-precision residual adapters at inference.
  • Continuous adapter training can violate this constraint or create a gap between the optimized state and the merged, requantized deployment state.
  • Direct discrete search preserves deployment faithfulness, but random or population proposals may require many forward evaluations in high-dimensional code spaces.
  • GradCodeS maps code-space geometry into a surrogate gradient, samples nearby code candidates around its reference, and selects updates by realized deployed loss.
  • GradCodeS derives a locally steepest-descent code-coordinate signal and uses guided candidate sampling with loss-based selection over deployable low-bit candidates.
  • Across three task families, multiple datatypes, and full or structured parameterizations, GradCodeS is reported to provide consistent fully low-bit gains.

2 RELATED WORK

Prior work either keeps continuous high-precision adapters or searches directly in quantized models. GradCodeS addresses the resulting deployment-faithfulness and proposal-efficiency trade-off with analytical guidance followed by loss-based selection.

  • Continuous low-bit adaptation is effective when mixed-precision inference or a separate conversion stage is acceptable, but optimized and deployed states can differ under strict fully low-bit deployment.
  • Zeroth-order and evolution-strategy methods evaluate quantized or nearby candidates, yet proposal efficiency remains difficult in high-dimensional spaces under finite evaluation budgets.
  • GradCodeS uses one backward-pass analytical guidance to focus discrete search while selecting among deployable low-bit candidates by realized loss.
  • Gradient-informed discrete methods motivate separating proposal construction from state selection, but do not address quantization-specific code geometry or finite-candidate loss minimization.

3 PROBLEM SETUP

The problem is formulated over feasible quantization codes and admissible scales for targeted layers, with deployment weights recovered directly from those variables. This excludes post-training conversion from the optimization objective.

  • A quantization datatype specifies bitwidth, codebook, grouping structure, and admissible scale space for feasible code-and-scale states.
  • Code matrices select codebook values entrywise, while each weight receives the scale metadata of its assigned group.
  • The representation covers NF4, INT4, and MXFP4 through datatype-specific codebooks, grouping, and scale components.
  • Targeted quantized linear layers use collections of codes and scales, while parameters outside those layers remain unquantized.
  • Fully low-bit fine-tuning is defined directly over deployable states, with discrete codes and continuously optimized but admissibly constrained scales.

4 METHOD

GradCodeS separates gradient guidance from discrete state selection: a code surrogate gradient centers proposals, while realized deployed loss determines accepted low-bit updates. Its analysis connects local surrogate descent to feasible one-level code moves and expected progress under guided sampling.

  • 4.2 FIRST-ORDER PROPERTIES: The code surrogate gradient is locally the steepest-descent direction of the active continuous surrogate under the Frobenius norm.This interpretation is local and does not by itself guarantee a valid discrete code update.
  • 4.2 FIRST-ORDER PROPERTIES: For feasible gradient-aligned one-level code moves, the surrogate linear term exactly matches the first-order deployed-weight change.This provides the discrete connection between the surrogate score and actual quantized parameter changes.
  • 4.2 FIRST-ORDER PROPERTIES: The realized loss is bounded by the surrogate prediction plus a curvature penalty, so sufficiently dominant negative surrogate terms guarantee loss decrease.Candidate evaluation remains necessary to account for nonlinear and higher-order effects.
  • 4.3 GRADIENT-GUIDED DISCRETE OPTIMIZATION: The guide–sample–evaluate–select pipeline samples M much smaller than the full code space, concentrating proposals near a real-valued reference point before retaining the lowest-loss candidate.The reference is generally non-integer and serves only to center proposals over valid codes.
  • 4.3 GRADIENT-GUIDED DISCRETE OPTIMIZATION: GradCodeS alternates projected scale updates with code search, using a surrogate gradient to guide proposals and deployed loss to select valid low-bit states.The procedure updates differentiable scales first, then samples and evaluates code candidates at the updated scales.
  • 4.3 GRADIENT-GUIDED DISCRETE OPTIMIZATION: Expected progress depends on proposal concentration and candidate count through the probability 1 −(1 −qt)^M that at least one candidate enters the specified near-reference set.The result is conditional on the stated smoothness, distance, and proposal assumptions.

5 EXPERIMENTS

Experiments evaluate GradCodeS across tasks, models, deployment regimes, data sizes, datatypes, and search components. Results consistently favor fully 4-bit GradCodeS while showing trade-offs in data scale and search cost.

  • Main results: Across both backbones and all three tasks, GradCodeS is the strongest fully 4-bit method.It provides the most consistent recovery from direct 4-bit quantization across reported metrics.
  • Main results: GradCodeS exceeds 16-bit SFT on GSM8K accuracy and MASSIVE results, while 16-bit SFT remains strongest on Qwen AlpacaEval and marginally higher on Llama AlpacaEval.The paper interprets these mixed results as a possible regularization effect rather than universal superiority of lower precision.
  • Main results: GradCodeS gives the best fully 4-bit results on both structured-prediction backbones and all three reported metrics.Exact Match requires the entire semantic frame to be correct, while Slot F1 measures predicted slot-structure quality.
  • Robustness across data: With 7.0k examples GradCodeS (LoRA) performs best, but with 167k examples 16-bit SFT becomes strongest.The result is consistent with, but does not prove, a regularization benefit from low-rank code parameterization in the smaller-data regime.
  • Robustness across datatypes: GradCodeS applies to NF4, INT4, and MXFP4, with NF4 and MXFP4 stronger than INT4 in this setting.The common search rule uses datatype-specific codebooks, scale grouping, and recovery maps.
  • Search efficiency: GradCodeS reaches a stronger result within the displayed time range despite using one backward pass and multiple candidate evaluations.The targeted evaluations support a trade-off between gradient-computation cost and wasted discrete search.
  • Candidate-set quality: Candidate loss tends to be lower near the guided reference, but realized discrete losses determine the accepted update among several nearby candidates.The guided sampler is selected more often than Gaussian or one-hop sampling, and the code surrogate gradient guides better candidates than mapped weight gradients.

6 CONCLUSION

GradCodeS uses quantization geometry to guide candidate sampling in code space while selecting updates by realized discrete loss. Across tested tasks and settings, it consistently improves fully 4-bit adaptation, with larger-backbone and richer-sampler scaling left for future work.

  • 6 CONCLUSION: GradCodeS finds a code-space steepest-descent direction to guide candidate sampling while preserving deployment-faithful selection by realized discrete loss.Every selected state remains a valid low-bit checkpoint.
  • 6 CONCLUSION: Across reasoning, instruction-following, and structured language-understanding tasks, GradCodeS consistently improves fully 4-bit adaptation.It remains compatible with different codebooks, low-rank code parameterizations, and alternating scale updates.
  • 6 CONCLUSION: Scaling to larger backbones and richer structured candidate-sampling families is identified as an important next step.The update principle is intended to remain deployment-faithful during this scaling.

A.1 MODELS AND EVALUATION TASKS

The evaluation uses compact instruction-tuned decoder-only models and three task families, comparing unquantized, mixed-precision, and fully 4-bit deployment regimes. GradCodeS is evaluated with both full-matrix and low-rank code updates under matched protocols.

  • Models: The study evaluates Qwen3-0.6B and Llama-3.2-1B and 3B-Instruct, using compact backbones relevant to memory-constrained deployment.The compact models also make optimization-method differences easier to attribute than in very large overparameterized models.
  • Evaluation tasks: GSM8K, AlpacaEval, and MASSIVE en-US cover arithmetic reasoning, instruction following, and structured semantic parsing.The evaluation reports answer accuracy, evaluator win rate, and structured prediction metrics respectively.
  • Quantization setup: All transformer-block linear parameters are quantized while token embeddings and the LM head remain unquantized.The reported 4-bit scores reflect the actually deployed quantized modules.
  • Baselines: Baselines are grouped into unquantized references, mixed methods retaining high-precision inference adapters, and fully low-bit deployable methods.A method is fully low-bit deployable only when quantized modules require no high-precision adapters or residual branches at inference.
  • GradCodeS variants: GradCodeS searches low-bit code states directly, and evaluates both full-matrix and low-rank code parameterizations.Every selected update is deployable without a 16-bit adapter or high-precision residual branch.

B DATATYPE-SPECIFIC INSTANTIATIONS OF THE UNIFIED QUANTIZATION MODEL

The unified quantization representation uses datatype-specific codebooks and scale groupings while sharing a common update rule. Feasible code moves are recovered into deployed weights through the appropriate datatype-specific map.

  • Datatype-specific components: The common update rule depends on each datatype through its ordered codebook, scale grouping, and recovery map.Table 6 summarizes these datatype-specific components.
  • Unified representation: The deployed weight is represented as Wq(Z, s) = Sq(s) ⊙ Cq(Z), subject to datatype-specific code and scale constraints.Z denotes code coordinates and s denotes scale variables in the recovered-weight representation.
  • Feasible updates: The effective step incorporates the datatype-specific codebook gap and scale automatically.Datatype-agnostic describes the shared optimization rule, not identical feasible spaces.

C DERIVATION OF THE CODE SURROGATE GRADIENT

The derivation replaces the unavailable codebook derivative with adjacent quantization gaps, producing a directional surrogate that maps weight gradients to code moves.

  • The codebook lookup lacks a standard derivative, so actual code-coordinate changes must be related to deployed-weight changes directly.
  • The local continuation uses forward and backward adjacent gaps as directional slopes for one-level code moves.This continuation is defined for code updates satisfying ∥∆Z∥∞≤1.
  • A negative weight-gradient entry selects a forward code move, while a positive entry selects a backward code move.The direction is encoded using the sign of the deployed-weight gradient.
  • The resulting effective step D combines the selected directional gap with the scale factor S(s).For uniform codebooks and interior coordinates, D reduces to the quantization gap multiplied by S(s).
  • Applying the chain rule to the local continuation yields the code-space surrogate gradient used for optimization.

D PROOFS FOR SECTION 4.2

The proofs establish that the surrogate gradient gives a locally steepest continuous direction, exactly matches aligned one-step code changes, and guides deployment-faithful candidate search.

  • The local directional surrogate extends the discrete loss to real-valued interpolated code updates.
  • −∇ZL is the steepest descent direction of the surrogate at the current code under the Frobenius norm.This holds whenever ∇ZL(Z, s) ≠ 0.
  • For feasible one-step moves aligned with the deployed-weight gradient, the surrogate linear term matches the exact first-order deployed-weight change.The move satisfies ∆Z ∈ {−1, 0, 1}^dout×din and entrywise gradient alignment.
  • The proof obtains this consistency by selecting the forward or backward adjacent gap according to each gradient sign.
  • Candidate sampling independently perturbs entries around a guided reference, with closer code states receiving greater probability mass.The neighborhood is controlled by the integer search radius ρ.
  • Under smoothness on admissible sampled moves and sufficient near-reference proposal mass, best-of-M search yields a one-step descent inequality.The accept-if-improving rule guarantees non-increasing loss even when no sampled candidate belongs to the good-update set.
  • The search rule also supports integer low-rank factorization, reducing searched coordinates from doutdin to R(dout + din).The factors are merged into the deployed code state rather than retained as inference-time residual adapters.
Loading 2608.30908v1…