PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Scalar Sensitivity: Activation-Aware Mixed-Precision LLM Quantization with Cross-Layer Refinement

Akihiro Yoshida, Yuma Ichikawa

arXiv:2609.25916v1cs.LGstat.APstat.MEstat.ML

TL;DR

Scalar sensitivity proxies discard Hessian directionality, and additive MCKP formulations omit cross-layer interactions. CASA combines activation-aware allocation with cross-layer local search; experiments on 7B–14B LLMs report consistent gains over scalar proxies, especially at ultra-low bit-widths.

  • Problem

    Existing scalar proxies reduce module sensitivity to one number, while MCKP-style allocators treat modules independently despite cross-layer interactions.

    Source: §1 · pp. 1–2

  • Method

    CASA replaces scalar costs with an activation-aware quadratic derived from the Kronecker-factored Hessian and refines the allocation through cross-layer local search.

    Source: §1 · p. 2, §3.1 · p. 4

  • Results

    On 7B–14B LLMs, CASA consistently outperforms scalar-based proxies, with cross-layer refinement providing further gains particularly in the ultra-low-bit regime.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Empirical improvement correlates with the per-model distortion bound, consistent with the prediction that scalar-proxy distortion tracks CASA's gains.

    Source: §6 · p. 9

    The cross-term approximation covers only adjacent layer pairs, and its quadratic proxy ranking may diverge from true perplexity ranking in extreme regimes.

    Source: §6 · p. 9

Abstract

from arXiv · show

Mixed-precision weight quantization is commonly formulated as a Multiple-Choice Knapsack Problem (MCKP), yet existing solvers rely on scalar sensitivity proxies that collapse each weight matrix's Hessian into a single number and treat every module independently. We prove that even the optimal scalar proxy incurs multiplicative distortion up to $\sqrt{κ(\mathbf{A})κ(\mathbf{B})}$ relative to the full activation-aware quadratic, where $κ(\mathbf{A})$ and $κ(\mathbf{B})$ denote the condition numbers of the input- and output-side Hessian factors. This bound varies from $10^1$ to $10^{13}$ for typical LLM modules, making inter-module sensitivity ranking unreliable. To address these limitations, we propose Cross-layer Activation-aware Sensitivity Allocation (CASA), a two-phase method. In Stage 1, the scalar proxy is replaced by an activation-aware metric derived from the Kronecker-factored Hessian, reducing the MCKP to a form whose continuous relaxation admits a closed-form solution. In Stage 2, a cross-layer-aware local search evaluates bit-width updates using the end-to-end model loss. Experiments on multiple LLMs across different bit budgets show that CASA achieves lower perplexity than the latest scalar-proxy baselines, especially at ultra-low bit-widths ($<3$ bits per weight). Moreover, the performance gain in zero-shot accuracy tracks the per-model average condition-number over modules, confirming the distortion bound as a practical indicator of scalar-proxy failure.

1 Introduction

Scalar sensitivity proxies discard directional Hessian information, while MCKP allocators also assume modules are independent despite cross-layer error interactions. CASA separates activation-aware allocation inside the MCKP from cross-layer refinement outside it, and experiments report better performance than scalar-based proxies across models and quantization settings.

1 Introduction

  • Scalar proxies discard Hessian directionality and can misrank modules, causing the MCKP to allocate fewer bits to modules that matter most.The lost information includes activation-covariance anisotropy, heterogeneous output sensitivities, and channel structure in quantization perturbations.Source: §1 · p. 1
  • MCKP-style allocators assume layer-wise independence, but residual-stream error propagation creates cross-layer interactions that additive objectives cannot represent.Experiments found that putting cross-layer terms inside one MCKP did not meaningfully improve over the self-only baseline, motivating a mechanism outside the MCKP.Source: §1 · p. 2
  • CASA uses a Kronecker-factored activation-aware loss for Stage 1, then performs greedy bit-swap search using calibration cross-entropy after re-quantization for Stage 2.The second stage evaluates non-additive interactions end-to-end rather than through a quadratic surrogate.Source: §1 · p. 2
  • CASA improves performance over scalar-based proxies across multiple models and quantization settings, with gains growing alongside the distortion bound.The cross-layer refinement adds further gains, particularly in the ultra-low-bit regime.Source: §1 · pp. 2–3

2 Preliminaries

Prior bit-allocation work frames mixed-precision assignment as constrained optimization and uses varied search and solver strategies. Sensitivity methods commonly reduce importance to a scalar, while cross-layer awareness has appeared in related quantization settings but remains unexplored for LLM bit allocation.

2.1 Bit-allocation as constrained optimization problem

  • Mixed-precision bit allocation is formulated as an MCKP under resource constraints, alongside genetic algorithms, mathematical solvers, reinforcement learning, and differentiable search.Constraints can concern model size or computational complexity.Source: §2.1 · p. 3

2.2 Sensitivity Measurement

  • HAWQ, OMPQ, Q-Palette, and HIGGS use sensitivity surrogates, but scalar reductions cannot represent calibration-dependent activation distributions or activation outliers [Dong et al., 2019, 2020, Yao et al., 2021, Ma et al., 2023, Lee and Song, 2025, Malinovskii et al., 2025].Recent quantization methods explicitly handle activation outliers [Lin et al., 2024, Xiao et al., 2023].Source: §2.2 · p. 3

2.3 Cross-layer Awareness for Quantization

3 Activation-Aware Quantization Error Proxy

The paper derives a per-layer quadratic proxy using input activation and output-sensitivity factors, then characterizes scalar proxies as coarser approximations. It argues that scalar distortion can prevent uniform approximation and that sensitivity misranking can produce arbitrarily large MCKP regret.

3.1 Derivation

  • Under a second-order loss approximation and layer-wise block-diagonal Hessian assumption, per-layer loss is expressed using input activation Gram matrix A and output-side Hessian B.A encodes channel correlations and magnitudes; B captures downstream sensitivity across output channels.Source: §3.1 · p. 4

3.2 Existing Proxies as Degenerate Cases

  • Scalar proxies collapse A and B to scaled identities, weighting every perturbation entry equally despite activation outliers and heterogeneous output-side sensitivity.The resulting loss of channel structure can misrank modules by quantization sensitivity.Source: §3.2 · p. 4

3.3 Theoretical Discussion

  • When either A or B is anisotropic, no scalar proxy uniformly approximates the activation-aware quadratic without the theorem's worst-case multiplicative distortion.The bound uses the condition numbers of A and B, defined by each factor's maximum-to-minimum eigenvalue ratio.Source: §3.3 · p. 4
  • Opposite rankings between true and proxy sensitivities can send the higher bit-width to the wrong module and produce regret that diverges as R grows.In the stated construction, regret equals (R −1)(dL −dH) when γ1 = R and γ2 = 1.Source: §3.3 · pp. 4–5

4 Cross-layer Activation-aware Sensitivity Allocation (CASA)

CASA separates activation-aware self-term allocation from cross-layer refinement: Stage 1 solves an MCKP, and Stage 2 searches for better allocations using cross-layer information and calibration loss.

4.1 Stage 1: MCKP Formulation

  • Stage 1 formulates per-module bit-width choices as an MCKP using activation-aware costs, subject to the target average bits-per-weight budget.The continuous relaxation under the high-rate distortion law admits a closed-form optimal allocation.Source: §4.1 · pp. 5–6

4.2 Stage 2: Local Search for considering Cross-Layer Interactions

  • The cross-layer approximation retains interactions between adjacent layers and same-module types, and exact cross-aware optimization is no worse than self-only optimization over the feasible set.The inequality is strict when the self-only solution does not minimize the quadratic cross-aware objective.Source: §4.2 · p. 6
  • Stage 2 refines the Stage-1 allocation with feasible bit-width swaps, screening by estimated cross-term improvement and evaluating candidates using calibration loss.It accepts the candidate with the largest loss reduction and stops if no candidate improves on the incumbent.Source: §4.2 · p. 7

5 Experiments

Across the reported experiments, CASA and its activation-aware ablation outperform Q-Palette, while cross-term-aware refinement adds gains that are largest at ultra-low bit budgets. The results also connect allocation patterns and empirical gains to module sensitivity and distortion.

Setup

  • The experiments cover five Llama and Qwen models, with WikiText-2 perplexity and average accuracy on common-sense reasoning benchmarks reported.The setup uses C4 calibration data and candidate bit-widths from 2 to 8.Source: §5 · p. 7

Results and Discussion

  • CASA and CASAself consistently achieve better performance than the scalar-based proxy Q-Palette.Table 2 reports WikiText-2 perplexity and six-task average accuracy across mixed-precision methods and bit budgets.Source: §5 · p. 7, Table 2 · p. 8
  • For Llama-3-8B at 2.5 BPW, CASA assigns 5–8 bits to early-layer V projections, including 8 bits on layer 2, while Q-Palette assigns 5 bits on layers 2–4.The reported input Gram condition numbers for V projections in layers 0–4 range from 2.4×10^11 to 1.5×10^15.Source: §5 · p. 7
  • CASA allocates more bits to modules with output-sensitive dimensions, reflecting substantial variation in output-side Hessian sensitivity.Figure 2 visualizes the allocation for Llama-3-8B at 2.5 BPW; blue denotes lower precision and white higher precision.Source: §5 · p. 8, Figure 2 · p. 9
  • Cross-term-aware local search consistently improves over CASAself, with the largest gains at 2.25–2.75 BPW.Using self-term rather than cross-term ranking improves over CASAself but remains below full CASA.Source: §5 · p. 8
  • At 3.0 BPW, CASA’s zero-shot accuracy gain over Q-Palette grows substantially for models with higher distortion.Table 3 reports mean distortion bounds and CASAself accuracy gain over Q-Palette at this budget.Source: §5 · p. 8, Table 3 · p. 9

6 Conclusion

CASA combines an activation-aware quadratic proxy with cross-layer local search, and its reported gains correlate with the per-model distortion bound. The conclusion limits the cross-term approximation to adjacent layers and notes that proxy rankings can diverge from downstream perplexity rankings in extreme regimes.

6 Conclusion

  • Across 7B–14B LLMs, the activation-aware proxy consistently outperforms scalar-based proxies, and cross-layer refinement adds gains particularly in the ultra-low-bit regime.The empirical improvement correlates with the per-model distortion bound, consistent with Theorem 3.1.Source: §6 · p. 9
  • The cross-term approximation models adjacent layer pairs only, and its activation-aware quadratic proxy may rank allocations differently from true downstream loss in extreme regimes.Longer-range or higher-order interactions are not captured because their incorporation would make calibration and optimization combinatorially intractable.Source: §6 · p. 9
  • The paper identifies both potential benefits for on-device and edge deployment under resource constraints and the broader risk of lowering barriers to LLM misuse.It states that this misuse concern is common to LLM efficiency research and is not amplified by this contribution.Source: §6 · p. 9

A Overall Algorithm

  • The overall algorithm collects module statistics, solves a self-term MCKP, then screens feasible swaps and evaluates selected candidates by calibration loss.The search stops when no candidate improves on the incumbent calibration loss.Source: A Overall Algorithm · p. 12

B Per-Module Distortion Bound

The scalar-proxy distortion bound depends on the condition numbers of the input- and output-side matrices and can vary sharply across modules. Reported heatmaps show especially large bounds in early layers and substantial within-layer variation.

B Per-Module Distortion Bound

  • Theorem 3.1 bounds scalar-proxy distortion by √(κ(A)κ(B)), where κ(A) and κ(B) are the input-side Gram and output-side curvature condition numbers.When the bound differs across modules, a scalar proxy can misrank their relative sensitivities.Source: B Per-Module Distortion Bound · p. 12
  • In Qwen3-8B, the distortion bound ranges from approximately 10^1.8 for o_proj in layer 2 to 10^13.2 for gate_proj in layer 1.This is a spread of over 11 orders of magnitude.Source: B Per-Module Distortion Bound · p. 12
  • Early layers 0–3 have the highest reported distortion across models, reaching 10^10 for Llama-3-8B Q projections and 10^13 for Qwen3-8B gate projections.Figure 3 maps log10 distortion across layer-module pairs, with brighter colors indicating larger worst-case distortion.Source: B Per-Module Distortion Bound · p. 12, Figure 3 · p. 13
  • Within a Qwen3 layer, the bound can vary by up to 8 orders of magnitude across modules, so scalar proxies have no uniform ranking guarantee.The reported modules include κ values around 10^10 alongside values around 10^2.Source: B Per-Module Distortion Bound · p. 13

C Ablation Study: Single-Stage MCKP Combining Self- and Cross-Terms

The ablation tests whether adding adjacent-layer cross terms directly to the MCKP improves on a self-term objective; neither tested coefficient choice yields a meaningful gain.

C Ablation Study: Single-Stage MCKP Combining Self- and Cross-Terms

  • These experiments motivate handling cross-layer effects outside the MCKP through CASA’s two-stage design rather than folding them into one knapsack objective.The tested single-stage formulation does not improve meaningfully over the self-term-only baseline.Source: C Ablation Study: Single-Stage MCKP… · p. 14
  • The single-stage objective adds adjacent-layer cross terms using either geometric-mean bounds or raw absolute cross-layer Fisher diagonals, then linearizes the bilinear products with McCormick constraints.The raw_abs option requires storing adjacent-layer activations; geo_abs reuses Stage-1 self-term diagonals.Source: C Ablation Study: Single-Stage MCKP… · p. 14
  • Neither geo_abs nor raw_abs meaningfully improves over self-term-only MCKP: geo_abs collapses to near-uniform allocations, while raw_abs only matches or marginally improves despite a larger optimization problem.The table compares Wikitext-2 perplexity and 4-task average accuracy across the single-stage objectives.Source: C Ablation Study: Single-Stage MCKP… · p. 14, Table 4 · p. 15

D Ablation Study: Local Search without Cross-Layer Awareness

The ablation finds that cross-term-aware local search improves on self-only search, especially at ultra-low bit widths. Additional figures visualize allocations for four models at 2.5 BPW.

D Ablation Study: Local Search without Cross-Layer Awareness

  • Cross-term coupling in the search objective is particularly beneficial for bit allocation in the low-bit regime, beyond the benefit of local search alone.The comparison reports that full CASA further improves over self-only local search in 13 of 15 configurations.Source: D Ablation Study: Local Search without… · p. 16
  • At 2.25 BPW on Llama-3-8B, CASA lowers perplexity from 28.15 (+LSself) to 22.41, a 20% relative reduction.Across 15 configurations, full CASA improves over +LSself in 13, while +LSself improves over CASAself in 14.Source: Table 5 · p. 16, D Ablation Study: Local Search without… · p. 16

E Visualization of bit-allocation

F Additional Experimental Results

Appendix F identifies Table 6 as additional results across mixed-precision methods and bit budgets. Appendix G supplies notation for its additional theoretical treatment of the activation-aware self-cost.

F Additional Experimental Results

G Additional Theoretical Results

  • The additional theoretical material introduces notation for the activation-aware self-cost by module and candidate bit width, suppressing module superscripts when discussing one module.The supplied passage describes the notation but does not provide the cost expression itself.Source: G Additional Theoretical Results · p. 19

G.1 A Lower Bound Showing Why Scalar Proxies Can Fail

This section establishes limits on scalar proxies and characterizes a high-rate activation-aware rule for continuous bit allocation. It shows that the bounded allocation problem has a unique clipped water-filling optimum whose threshold satisfies the budget.

G.1 A Lower Bound Showing Why Scalar Proxies Can Fail

  • The sharp scalar-proxy distortion is √(κ(A)κ(B)) for positive-definite factors, so anisotropy in either factor prevents uniform approximation.The bound compares the activation-aware quadratic with the best scalar rescaling and is attainable in the worst case.Source: §G.1 · p. 19
  • If either curvature factor is singular, every scalar proxy has infinite distortion for some perturbation.The corollary covers singularity in either the input-side factor A or output-side factor B.Source: §G.1 · p. 21
  • With opposite true and proxy sensitivity rankings, the proxy assigns the higher bit width to the wrong module and can incur arbitrarily large regret.For γ1 = R and γ2 = 1, regret is (R − 1)(dL − dH), which diverges as R grows.Source: §G.1 · p. 21

G.2 High-rate Quantization and the Optimal Bit-allocation Rule

  • Under the high-rate second-moment model, expected activation-aware distortion is summarized by a nonnegative module sensitivity Γ_i, computable from entrywise error second moments when the factors are diagonal.In the diagonal case, the formula does not require off-diagonal error covariances to vanish.Source: §G.2 · pp. 23, 25
  • The continuous bit-allocation relaxation has a unique global minimizer, assigning larger widths to modules with larger activation-aware sensitivity.The rule applies when each module has positive effective sensitivity and positive weight count.Source: §G.2 · p. 25
  • The bounded bit-allocation rule clips each coordinate to its permitted interval, with a shared threshold chosen to meet the budget.This is the clipped water-filling form for the bounded allocation problem.Source: §G.2 · p. 28, Then the problem · pp. 29–30

Then the problem

  • Strict convexity and a nonempty, closed, bounded feasible set ensure a unique global minimizer, and the budget constraint is active at that optimum.The proof derives the coordinate-wise clipped minimizer and shows that a threshold can be chosen to satisfy the active budget.Source: Then the problem · pp. 28–30

G.3 Why cross-layer extensions are theoretically beneficial

Exact cross-aware optimization over the feasible bit assignments cannot worsen the quadratic surrogate relative to self-only optimization, and can be strictly better. The omitted cross terms can create a substantial gap, motivating conservative cross-layer objectives.

G.3 Why cross-layer extensions are theoretically beneficial

  • Exact cross-aware optimization never has higher quadratic cost than self-only optimization, with strict improvement when the self-only assignment is not a full-cost minimizer.Both methods optimize over the same feasible set.Source: §G.3 · p. 30
  • The self-only gap is bounded by the difference in omitted cross terms, which is at most twice their maximum absolute value over feasible assignments.This bounds the quadratic-objective gap using the cross-term contribution.Source: §G.3 · p. 31
  • Self-costs can tie across every feasible assignment while a self-only choice is arbitrarily poor relative to the cross-aware optimum.The example gives a positive-definite two-module Hessian and an additive gap of 4ρR.Source: §G.3 · pp. 31–32
  • A conservative cross-layer objective upper-bounds the quadratic surrogate when pairwise cross terms satisfy the stated bounds; positive-semidefinite pair blocks permit ρ_ij = 1.The bound applies to every feasible assignment.Source: §G.3 · p. 33
  • Minimizing the robust cross-aware objective improves its conservative upper bound relative to the self-only assignment.The comparison follows because both assignments are feasible and the cross-aware one minimizes that bound.Source: §G.3 · p. 33

Related papers

Loading 2609.25916v1…