PaperlayerView exact revision on arXiv

Source-linked AI summary

Jacobian-guided Noise Injection for Quantization Robustness in Large Language Models

Deepanshu Pandey, Arnav Chavan, Nahush Lele, Sankalp Dayal, Deepak Gupta

arXiv:2608.20988v1cs.LGcs.AI

TL;DR

Low-bit quantization can degrade LLM performance because softmax sensitivity amplifies errors from outlier logits. The paper introduces Jacobian-guided Gaussian noise injection with locally calibrated variance, reporting consistent robustness improvements without inference overhead. Its scope includes multiple architectures and quantization settings, while the authors note computational constraints for direct second-order Jacobian regularization.

  • Problem

    Low-bit quantization is vulnerable to softmax-mediated amplification of approximation and rounding errors, especially when logits contain outliers or large magnitudes.

  • Method

    Jacobian-guided noise injection adds Gaussian logit perturbations with variance calibrated from an approximate softmax Jacobian Frobenius norm.

  • Results

    The method improves quantized performance consistently across Llama and Qwen models and multiple quantization levels, recovering up to 37% accuracy in W4A4 with zero inference overhead.

  • Takeaways & Limitations

    The approach learns representations that are more robust to discretization error while preserving full-precision accuracy without significant degradation.

  • Takeaways & Limitations

    Direct Jacobian regularization requires continuous second-order derivatives and is computationally intractable for large Transformers, motivating a first-order approximation.

Abstract

from arXiv · show

Quantization of Large Language Models (LLMs) is often hindered by the sensitivity of the self-attention mechanism to discretization errors. We identify the softmax operator as a bottleneck for quantization stability due to its sensitivity to outliers and state-dependent Jacobian. We theoretically establish that suppressing the norm of this Jacobian helps in bounding quantization-induced performance degradation. Based on this, we propose Jacobian-Guided Noise Injection, a training strategy that injects zero-mean Gaussian noise into pre-attention logits, with variance derived directly from the Jacobian Frobenius norm. Unlike prior approaches that rely on heuristic or penalise jacobian directly, our method provides a way to identify the optimal noise variance based on the local attention sensitivity. We evaluate the method on SOTA LLM architectures, where it demonstrates improved robustness over popular PTQ methods. Empirical analysis reveals that the proposed method gives up to +37% relative gains on Top-1 accuracy on ImageNet-1K for SigLIP and improves relative perplexity by upto 40% on WikiText for language models in low bit quantisation settings, proving the efficacy of the approach.

1. Introduction

Quantization improves deployment efficiency but can severely damage model performance when low-precision errors are amplified by sensitive softmax logits. The paper proposes Jacobian-guided noise injection to bound this amplification and improve robustness.

  • 4-bit weight quantization can reduce memory footprint by 4× versus FP16 and achieve over 3× inference speedup.These savings can enable deployment of 70B-parameter models on mobile GPUs.
  • Low-bit quantization causes performance degradation because approximation noise and rounding errors perturb intermediate computations, with degradation worsening at 4-bit and below.Activation outliers and reduced representational capacity contribute to the problem.
  • Large or outlier logits amplify quantization errors exponentially through softmax, so small perturbations can cause disproportionately large probability changes.The resulting error propagation is unpredictable for standard quantization techniques.
  • The paper analyzes softmax quantization-error propagation and derives conditions for bounding it using the softmax Jacobian norm.It focuses on the relationship between spectral Jacobian norm and error amplification, showing that minimizing expected loss under logit perturbations implicitly regularizes this norm.
  • Jacobian-guided noise injection calibrates Gaussian noise variance using an approximate Jacobian Frobenius norm, replacing heuristic variance selection.The method is presented as a simpler alternative to heuristic approaches.
  • 37% accuracy is recovered in W4A4 with zero inference overhead across multiple LLM architectures and quantization settings.The reported robustness improvements span multiple architectures and settings.

2. Methodology

The methodology identifies the self-attention softmax as a quantization bottleneck and develops Jacobian-guided Gaussian logit noise to regularize sensitivity adaptively.

  • Softmax sensitivity: The self-attention softmax is treated as the bottleneck because its state-dependent Jacobian can amplify logit perturbations from quantization.Sensitive regions occur when attention mass is distributed across multiple tokens, whereas concentrated attention approaches saturated low-Jacobian regions.
  • Jacobian constraints: Bounding quantization error requires constraining the softmax Jacobian spectral norm, which can be related to its Frobenius norm.The resulting regularization objective penalizes the trace of the logit Hessian.
  • Practical approximation: Computing the ideal Hessian-trace penalty directly is computationally intractable for large Transformers, motivating its first-order approximation through noise injection.The methodology therefore replaces continuous second-order derivative computation with an efficient stochastic training procedure.
  • Noise-based regularization: Gaussian noise applied directly to logits implicitly regularizes the Hessian trace because zero-mean perturbations eliminate the first-order Taylor term.For isotropic noise, the variance acts as the regularization strength with λ = σ^2/2.
  • Adaptive noise injection: The proposed method derives noise variance from the local Jacobian Frobenius norm rather than using a fixed global value.This adaptive scheme targets positions according to their local attention sensitivity while avoiding explicit Jacobian materialization.
  • Adaptive noise injection: Rowwise noise applies stronger perturbations to high-sensitivity attention positions and weaker perturbations to saturated positions, with parameters updated periodically.Algorithm 1 perturbs pre-softmax logits before computing the noisy attention probabilities and training loss.

3. Experiments

Experiments across language and vision-language models evaluate Jacobian-guided noise injection under PTQ and QAT, finding consistent robustness gains and sensitivity to noise scaling.

  • Evaluation setup: Experiments cover Llama-3.2-3B, Qwen2.5-3B, and SigLIP across multiple low-bit quantization settings and PTQ methods.Language models are fine-tuned on Alpaca, while SigLIP is evaluated on ImageNet-1K classification.
  • Generalization: Consistent gains across three model families and four PTQ methods indicate robustness beyond a specific architecture or quantization scheme.The evaluated PTQ methods include AWQ, GPTQ, ERQ, and RepQViT.
  • Attention analysis: Noise-trained models exhibit more diffuse attention patterns across key positions, consistent with smaller Jacobian norms.The diffusion phenomenon is observed across model layers.
  • Noise scaling: α = 0.5 achieves 31.37% Top-1 accuracy versus 30.31% for the SigLIP W4A4 QAT baseline.Performance is best for moderate scaling, α ∈[0.5, 1.0], and degrades rapidly at excessive noise.

4. Related Work

Prior work connects Jacobian control and noise-based stability to robustness, while this paper applies those ideas specifically to attention softmax quantization.

  • Attention analysis: Noise-trained SigLIP attention maps are more diffuse across key positions, providing the visual behavior examined in the paper’s analysis.The figure concerns SigLIP layer 10.
  • Quantization methods: PTQ methods such as GPTQ and AWQ compress models to 4–8 bits using second-order information, salient-weight protection, or outlier handling.QAT instead incorporates quantization into the training loop, but both approaches typically treat layers uniformly.
  • Jacobian robustness: Jacobian regularization has been used to learn locally invariant features and to bound sensitivity to input perturbations.Related work also links Gaussian-noise stability to reduced softmax curvature.
  • Paper positioning: The proposed approach targets attention softmax by deriving noise variance from the Jacobian norm rather than applying uniform perturbation.Noise adapts per position to concentrate regularization where quantization errors are most amplified.

5. Conclusion

The paper concludes that Jacobian-guided noise injection adaptively regularizes attention softmax sensitivity and improves low-bit quantization robustness without inference overhead.

  • Method: Noise variance derived from the Softmax Jacobian Frobenius norm targets sensitive regions of attention softmax.The method is presented as adaptive regularization for quantization robustness.
  • Evidence: Experiments across multiple PTQ and QAT methods and model families demonstrate consistent improvements.The paper also reports qualitative evidence linking noise injection to Jacobian sensitivity.
  • Practical implication: The method adds no inference overhead, matches FP16 baseline performance, and improves quantization robustness.The conclusion frames these properties as practical for resource-constrained deployment.

Impact Statement

The impact statement presents robust low-bit quantization as a way to reduce deployment costs while positioning the method as general-purpose and capability-neutral.

  • Deployment impact: Robust quantization can reduce the computational cost, memory footprint, and energy required to deploy LLMs.The statement connects these benefits to aggressive low-bit quantization with reduced quality degradation.
  • Scope: The proposed methods are described as general-purpose techniques applicable to any neural network.The statement says they do not introduce novel capabilities.

B.1. Training Configuration

The training configuration combines low-precision evaluation across language and vision models with Jacobian-guided, rowwise noise injection and distributed fine-tuning. Experiments cover multiple quantization methods, bit-widths, and zero-shot benchmarks.

  • Optimization: AdamW training uses learning rate 5 × 10^-7, weight decay 0.01, gradient clipping at 1.0, cosine scheduling, and 500 warmup steps.
  • Training setup: Models are fine-tuned for 5 epochs on Alpaca with maximum sequence length 2048 tokens, using evaluation every 2000 steps.
  • Noise injection: Noise scaling uses α = 0.5, with noise standard deviation clamped to [0.005, 10.0] for numerical stability.
  • Noise injection: Jacobian-guided noise uses rowwise, per-position calibration so high-sensitivity positions receive targeted noise rather than uniform layerwise or global noise.
  • Quantization: Experiments evaluate WxAy configurations with per-channel weight quantization and per-token activation quantization for language models.
  • Evaluation: Language models are evaluated zero-shot on HellaSwag, Winogrande, PIQA, TruthfulQA, WikiText, BoolQ, Arc-Easy, and Arc-Challenge, while vision models use ImageNet-1K classification.

C. SigLIP Post-Training Quantization Results

SigLIP post-training quantization experiments use the base 16-384 vision encoder and compare ERQ and RepQViT across low-bit configurations. The evaluation measures ImageNet-1K zero-shot Top-1 and Top-5 accuracy against baseline results.

  • Setup: SigLIP base 16-384 processes 384×384 images with 16×16 patches and is evaluated on ImageNet-1K using a partitioned 95-5 test split.
  • Quantization methods: The vision experiments compare ERQ, which minimizes propagated reconstruction error, with RepQViT, which decouples quantization scales for hardware-friendly deployment.
  • Evaluation: ImageNet-1K zero-shot results report Top-1 accuracy for SigLIP base 16-384 under post-training quantization, using baseline and noise-injection variants.
  • Bit-width settings: The evaluated configurations include aggressive W4A4 quantization, moderate W5A5 quantization, conservative W6A6 quantization, and near-lossless W8A8 quantization.
  • Training protocol: The training protocol sets the noise scaling coefficient α to 0.5 based on the sensitivity analysis in Figure 3.

D. Full-Precision Results With and Without Noise

Full-precision evaluation examines whether Jacobian-guided noise injection affects Qwen2.5-3B performance before quantization. The reported comparison uses FP16 baseline and noise-injection results.

  • Full-precision comparison: Table 3 compares Qwen2.5-3B in FP16 with and without Jacobian-guided noise injection during training.
  • Reported finding: The comparison is intended to demonstrate that noise injection does not degrade full-precision performance while providing quantization robustness.

E. Detailed zero-shot PTQ results for Llama and Qwen

Detailed zero-shot results evaluate noise injection for Llama and Qwen models across seven benchmarks. The reported pattern is improved overall quantized performance with comparable or better results across the benchmark suite.

  • Benchmark coverage: Tables 4, 5, and 6 report detailed results for Llama and Qwen models across all seven zero-shot benchmarks.
  • Reported results: Noise injection improves overall model performance while maintaining comparable or better performance across the seven benchmarks.

F. Activation Similarity Analysis on Hardware

Hardware-focused activation analysis compares quantized and unquantized behavior across layers and evaluates how Jacobian-guided noise injection affects numerical stability. The noise-trained models preserve substantially higher activation similarity and reduce layer-wise divergence under low-bit and edge-hardware deployment settings.

  • Deployment setup: 15-bit edge deployment measures layer-wise cosine similarity between quantized and unquantized activations to assess preservation of full-precision computational behavior.The setup uses fixed-point arithmetic on edge accelerator hardware, whose rounding behavior can differ from GPU training.
  • Baseline behavior: Baseline activations fluctuate dramatically across layers, with multiple similarity scores approaching zero.These errors accumulate and amplify through the network, causing internal representations to diverge from the full-precision model.
  • Noise-trained behavior: >0.95 cosine similarity is maintained across all layers by the noise-trained model.This indicates substantially more consistent alignment between quantized and unquantized activations than the baseline.
  • PTQ reproduction: W4A4 RTN PTQ analysis likewise finds that noise training reduces layer-wise activation divergence and maintains higher similarity to actual activations.The result is reproduced for the SigLIP base 384 model.
  • Deployment implications: Softmax instability is linked to deployment failures because near-zero intermediate similarity propagates and compounds through subsequent attention operations.The proposed training regularizes attention so models maintain more consistent behavior across hardware with varying numerical precision.
  • Attention patterns: Attention maps for SigLIP layer 10 become more diffused in noise-trained models than in the baseline.The supplied figure passage identifies the qualitative difference but does not report a numerical comparison.
Loading 2608.20988v1…