PaperlayerView exact revision on arXiv

Source-linked AI summary

ARCHead: Activation-Metric Residual Correction for Large Language Model Output Heads

Şuayp Talha Kocabay, Talha Rüzgar Akkuş, Kamer Ali Yuksel

arXiv:2608.02703v1cs.CLcs.LG

TL;DR

Quantized LLM deployments often retain a large BF16/FP16 output head, while naive low-bit quantization can distort logits. ARCHead compresses this head with a quantized low-rank core, INT4 residuals, and activation-metric correction, achieving 3.7–3.9× persistent head-storage reduction with a practical quality–storage trade-off.

  • Problem

    Quantized LLM deployments often retain a large dense BF16/FP16 LM-head, whose direct effect on logits makes ordinary low-bit quantization difficult.

  • Method

    ARCHead combines quantized low-rank factors, group-wise INT4 residuals, and an activation-derived metric low-rank correction without storing a dense BF16 head.

  • Results

    3.7–3.9× persistent head-storage reduction is measured across tested heads while ARCHead avoids the large quality loss of storage-matched naive INT4.

  • Takeaways & Limitations

    ARCHead provides a practical quality–storage trade-off for dense output heads and can compose with AWQ or bitsandbytes without materially changing generation throughput.

  • Takeaways & Limitations

    ARCHead is specialized for dense output heads, and broader architectures, languages, context lengths, and calibration domains remain to be tested.

Abstract

from arXiv · show

Weight-only quantization substantially reduces the storage of large language model (LLM) transformer blocks, but practical backends often retain the final language-modeling head (LM-head) in BF16 or FP16. Quantizing this projection naively can strongly perturb the vocabulary-logit distribution. We present ARCHead, a packed LM-head compressor that combines a quantized low-rank core, group-wise INT4 residuals, and a low-rank correction fitted in an activation-derived metric. ARCHead stores no dense BF16 head and reduces persistent LM-head storage by 3.7-3.9x. On Qwen3-8B-Base, it uses 25.6% of BF16 head storage while attaining 1.007 relative perplexity; storage-matched naive INT4 yields 1.14-1.16. Replacing the BF16 head left by AWQ or bitsandbytes adds only 0.006-0.007 cross-entropy, with less than 2% throughput change in our measurements. ARCHead therefore complements block quantizers by compressing the large output projection they can leave untouched. Code is available at https://github.com/suayptalha/archead.

1 Introduction

ARCHead addresses the storage and quality problems of LM-heads that remain in BF16 or FP16 after transformer-block quantization. It combines quantized low-rank factors, group-wise INT4 residuals, and an activation-derived low-rank correction without storing a dense BF16 copy.

  • Motivation: PTQ implementations often quantize transformer-block weights while retaining the final LM-head in BF16 or FP16, leaving a large dense projection.The LM-head maps the final hidden state directly to a logit for every vocabulary item.
  • Motivation: 1.18 and 1.28 GB: Qwen3-8B-Base and Gemma-4-E4B BF16 output projections can become among the largest remaining dense tensors after four-bit block quantization.AWQ and bitsandbytes NF4 retained Qwen3-8B-Base lm_head as a BF16 Linear layer of shape 151,936 × 4,096.
  • Problem: Naive low-bit head quantization minimizes weight-space error despite anisotropic hidden-state activity, so errors on frequently activated directions can perturb many vocabulary logits.The relevant objective is the expected output error induced by the hidden-state distribution.
  • Method: ARCHead uses quantized low-rank factors, a group-wise INT4 residual, and a low-rank correction fitted after an activation-derived metric transform.The module is a drop-in output head and stores no dense BF16 copy of the original matrix.
  • Results: 25–27%: ARCHead’s measured state-dictionary footprint of the corresponding BF16 head across five model families.For a fixed quantized core, the correction is the best rank-r approximation in the activation-derived metric before factor quantization.
  • Results: ARCHead avoids the large quality loss of storage-matched naive INT4, generalizes across three output heads, and compresses BF16 heads left by AWQ and bitsandbytes at a small additional loss.The evaluation covers logit fidelity, downstream accuracy, calibration sensitivity, construction time, generation throughput, persistent storage, and runtime peak memory.

2 Related Work

Prior LLM quantization methods use second-order information, activation-aware protection, parameter optimization, distillation, outlier isolation, and related strategies. ARCHead builds on low-rank compression by approximating a transformed quantization residual in an LM-head-input-derived metric, while evaluating deployable storage from registered tensors and separating persistent from temporary memory.

  • LLM quantization: LLM quantization methods include GPTQ’s second-order information, AWQ’s activation-salient channel protection, and SmoothQuant’s activation-weight difficulty balancing.Related approaches also optimize quantization parameters, use layer-wise distillation, or isolate outliers.
  • Low-rank and residual correction: ARCHead differs from low-rank compression and parameter-efficient updates by approximating the quantization residual after transforming it with a metric estimated from LM-head inputs.This objective allocates correction capacity according to directions affecting observed logits.
  • Realized storage: Deployable compression requires measuring bytes from the packed head’s registered tensors rather than relying on intended bit width alone.Persistent storage is reported separately from temporary forward-pass memory.
  • Realized storage: The packed head discards the original BF16 matrix after construction and retains a quantized core plus an activation-metric low-rank correction.The decomposition is illustrated as a packed quantized core W_d and an activation-metric low-rank correction.

3 Method

ARCHead constructs a quantized low-rank LM-head core, corrects its residual with a low-rank approximation fitted in a damped activation-derived metric, and stores the resulting factors and residual in packed form. The method discards the dense BF16 head, while its correction optimality is conditional on the fixed core and before factor quantization.

  • Activation-derived metric: p = 0.75 is selected on calibration data for Qwen3-8B-Base, controlling how strongly dominant activation directions are emphasized.At p = 1/2, the induced objective is the damped form of the empirical logit-error objective.
  • Conditional optimality: For fixed Wd and invertible Tp, the unquantized rank-rr correction factors minimize the transformed residual objective by the Eckart–Young–Mirsky theorem.This guarantee applies before correction-factor quantization and is conditional on the fixed core.
  • Packed implementation: The packed module registers only the quantized residual, core and correction factors, scales, and metadata; the original dense BF16 V × D head is discarded.ARCHead can replace the output embedding after block quantization and is complementary to AWQ, bitsandbytes, or another block backend.
  • Memory accounting: Persistent output-head storage is evaluated separately from forward-pass peak memory, which includes activations, logits, temporary buffers, allocator behavior, workspaces, and the KV cache.The measurement protocol avoids unsupported claims about backend-dependent peak-memory reductions.
  • Construction: ARCHead estimates a damped activation covariance from calibration activations, quantizes a low-rank core, and applies group-wise low-bit quantization to the remaining residual.The inputs include W, H, core and correction ranks, group size, metric power p, and damping λ.
  • Residual correction: ARCHead eigendecomposes the covariance, transforms the core residual, truncates it to correction rank rr, maps the right factor back, and quantizes both correction factors with INT8.The transformed residual is eE = ETp, with Tp = QΛpQ⊤ and its inverse used for mapping.

4 Experimental Setup

The experiments evaluate ARCHead across multiple language-model heads using WikiText-103 calibration and held-out-token metrics. Comparisons include BF16, INT8, INT4, SVD8+INT4, and GPTQ-style heads under specified Qwen settings and matched calibration protocols.

  • Models: Five models are evaluated: Qwen3-8B-Base, Gemma-4-E4B, VibeThinker-3B, Mistral-7Bv0.3, and LFM2.5-8B-A1B.Head quality uses the first three models, while packed-storage validation additionally includes the latter two.
  • Data and metrics: 16,384 held-out WikiText-103 test tokens support cross-entropy, perplexity, relative PPL, ∆CE, logit MSE, KL divergence, prediction agreement, and persistent-head-byte measurements.Calibration activations come from the WikiText-103 training split.
  • Reproducibility protocol: GPTQ-head comparisons use fixed, calibration-disjoint activations, 2,048–16,384 calibration tokens, and three matched seeds: 0, 1, and 2.The broader ARCHead sensitivity sweep uses 4,096–65,536 calibration tokens.
  • Baselines and implementation: Head-only baselines comprise BF16, row-wise INT8, group INT4, SVD8+INT4, and GPTQ-style INT4, while transformer blocks remain BF16 unless AWQ or bitsandbytes NF4 is named.This setup isolates output-head compression while allowing explicit full-model hybrid comparisons.
  • Baselines and implementation: On Qwen, ARCHead uses rc = 10, rr = 6, groups of 64, p = 0.75, and ridge 10−3.Primary deployment measurements use one NVIDIA RTX Pro 6000.

5 Results

ARCHead compresses LM-head storage while preserving perplexity and dense-logit fidelity across architectures and block-quantization backends. Its activation-metric correction is decisive, with small throughput and downstream-evaluation effects.

  • Perplexity and architecture results: 1.007 relative PPL remains near the dense head, versus 1.151 for storage-matched Group INT4, which ordinary low-rank augmentation does not repair.The improvement does not come from the quantized core alone.
  • Perplexity and architecture results: 1.007, 1.010, and 1.027 relative PPL are obtained on Qwen, Gemma, and VibeThinker, respectively.The gain over naive INT4 is largest on Qwen and narrower on Gemma.
  • Backend replacement and storage: 25.6% of BF16 head storage remains after replacing the dense head, with extra CE of 0.006 for AWQ and 0.007 for NF4.Both backends retained a 1.18 GB BF16 head, and ARCHead operates orthogonally to transformer-block quantization.
  • Backend replacement and storage: 3.71–3.91× compression is measured across five heads in persistent state-dictionary storage, not total forward-pass peak VRAM.The latter also includes logits, activations, allocator state, backend workspaces, and the KV cache.
  • Comparison with GPTQ: 44.0%, 49.3%, 33.6%, and 40.6% reductions in GPTQ degradation occur from 2K through 16K tokens, while ARCHead constructs the head 2.36–2.58× faster.ARCHead has lower mean relative PPL and logit MSE at every tested budget, with lower variance despite one 8K seed exception.
  • Ablation and fidelity: 1.134 relative PPL from the quantized core alone falls to 1.007 after adding rank-6 activation-metric correction, requiring a 2.6-point increase in storage ratio.The correction is fitted to the core’s remaining, activation-weighted error.

6 Conclusion

ARCHead compresses the BF16 output projection left after transformer-block quantization by repairing compact-core errors with an activation-metric residual branch. Its packed module reduces persistent head storage by 3.7–3.9× and complements, rather than replaces, full-model quantizers.

  • Conclusion: 3.7–3.9× persistent head-storage reduction is achieved by ARCHead’s packed module.The reduction targets the large BF16 output projection that can remain after transformer-block quantization.
  • Conclusion: The activation-metric residual branch repairs structured error left by a compact quantized core.This mechanism supports ARCHead’s practical quality–storage trade-off across tested heads.
  • Conclusion: ARCHead is complementary to full-model quantizers rather than a replacement for them.It can be composed with AWQ or bitsandbytes without materially changing deployment behavior.

Limitations

ARCHead is specialized to LM output heads, with benefits dependent on vocabulary-projection characteristics and broader evaluation still limited. Its memory and optimality claims also have workload- and quantization-dependent qualifications, while downstream tasks provide only a sanity check.

  • ARCHead compresses output heads but not transformer MLP or attention weights, limiting its scope to LM-head storage.
  • On Gemma-4-E4B, naive INT4 is already slightly better in the displayed rounded relative-PPL result.
  • The cross-model study covers three quality evaluations and packed size is checked on five heads, limiting breadth across architectures, languages, context lengths, and calibration domains.
  • The principal memory claim concerns serialized and load-time parameter tensors, while peak GPU memory depends on workload and backend factors such as activations, workspaces, allocator behavior, and KV caches.
  • The downstream tasks are only a sanity check, and conditional optimality applies only to an unquantized rank-r correction for a fixed core.

Ethical Considerations … C Packed Buffers and Model-Conversion Lifecycle

The paper describes ARCHead’s metric-aware construction, packed representation, and conversion lifecycle while noting ethical risks inherited from source checkpoints. It distinguishes persistent serialized storage from construction and inference memory, and specifies that the dense head is discarded before serialization.

  • Ethical Considerations: ARCHead inherits the biases, unsafe behaviors, risks, licenses, and intended-use constraints of its source checkpoint, despite potentially lowering deployment costs.The work uses publicly released checkpoints and WikiText-103 without human-subject data collection or annotation.
  • A Extended Notation and Objective: The activation-derived metric weights directions by activation energy, unlike Frobenius reconstruction, which treats covariance as the identity.ARCHead uses a damped covariance and interpolates between unweighted residual approximation and empirical logit-MSE geometry through the metric transform.
  • A Extended Notation and Objective: p = 0.75 is selected for Qwen; damping stabilizes weakly observed directions and keeps the transform inverse well defined.Values above one half emphasize dominant activation directions more strongly.
  • A Extended Notation and Objective: The truncated SVD correction is optimal only for a fixed core, invertible metric transform, and unquantized correction, not for joint optimization of the complete packed module.The theorem characterizes the ideal correction stage before INT8 factor quantization.
  • B Expanded Construction Procedure: The construction estimates activation geometry, quantizes a low-rank core and group-wise INT4 residual, then fits and packs an activation-metric low-rank correction.The correction targets realized errors from both core approximation and factor quantization, and its factors are stored row-wise and group-wise in INT8.
  • B Expanded Construction Procedure: The serialized module registers only packed tensors: the original BF16 V × D head and construction temporaries are discarded, while inference may use temporary dequantized state.Factorized evaluation avoids registering a reconstructed dense head.
  • C Packed Buffers and Model-Conversion Lifecycle: During conversion, ARCHead is fitted with the dense head available, replaces the model’s output embedding, and releases the dense tensor before serialization and reload.Integration occurs after transformer-block quantization, targeting the remaining dense output projection without changing the block representation.
  • C Packed Buffers and Model-Conversion Lifecycle: 3.7–3.9× refers only to persistent registered-tensor storage, not construction working memory or forward working memory.Construction and inference can additionally include dense weights, covariance factors, SVD workspaces, logits, activations, KV cache, allocator state, and backend workspaces.

D Evaluation and Reproducibility Checklist · E Full Head-Only Results · F Cross-Model Head-Only Results

The evaluation isolates LM-head effects through head-only and hybrid protocols, fixed calibration and evaluation procedures, and explicit metric interpretations. Head-only results show ARCHead’s advantage over naive INT4 is model-dependent but substantial across the reported models.

  • D Evaluation and Reproducibility Checklist: Head-only experiments keep transformer blocks in BF16, changing only the LM-head representation to isolate head quantization effects.
  • D Evaluation and Reproducibility Checklist: Hybrid experiments quantize blocks with AWQ or bitsandbytes, while BF16-head and ARCHead rows use the same block backend.
  • D Evaluation and Reproducibility Checklist: Calibration uses WikiText-103 training hidden states, while evaluation uses 16,384 held-out test tokens that are not reused for fitting.The GPTQ-style head and ARCHead also share activation caches across 2K, 4K, 8K, and 16K budgets with seeds 0, 1, and 2.
  • D Evaluation and Reproducibility Checklist: The checklist distinguishes quality, dense-head fidelity, prediction agreement, and serialized-size metrics rather than collapsing them into one claim.Cross-entropy and relative perplexity measure language-model quality; logit MSE and KL divergence measure fidelity; top-k agreement measures leading-candidate preservation; persistent head ratio measures serialized size.
  • D Evaluation and Reproducibility Checklist: The downstream HellaSwag, TruthfulQA MC2, and WinoGrande suite is a sanity check for obvious regressions, not evidence of task-level improvement or comprehensive capability preservation.
  • E Full Head-Only Results: Table 14 reports comprehensive Qwen3-8B-Base head-only results with BF16 transformer blocks, where ARCHead matches naive Group INT4 storage while reducing cross-entropy and perplexity degradation.
  • F Cross-Model Head-Only Results: ARCHead’s relative advantage varies by model: it is strongest on Qwen3-8B, narrower but competitive on Gemma-4-E4B, and substantial despite larger loss on VibeThinker-3B.The reported pattern reflects model-dependent LM-head quantization sensitivity, including naive INT4 baseline strength and smaller hidden dimension.

G Hybrid Quantizer + ARCHead Results … J Additional Result Visualizations

ARCHead compresses the dense LM-head retained by AWQ and bitsandbytes with minimal added cross-entropy, while activation-metric residual correction is essential for near-baseline perplexity. Measurements and visualizations indicate negligible throughput impact, a favorable storage–perplexity trade-off, and limited calibration sensitivity.

  • G Hybrid Quantizer + ARCHead Results: +0.006 to +0.007 cross-entropy is added when ARCHead replaces the dense BF16 head retained by AWQ or bitsandbytes.ARCHead is a drop-in head compressor rather than a method that improves the base perplexity of either backend.
  • G Hybrid Quantizer + ARCHead Results: Table 16 reports hybrid Qwen3-8B-Base results in which ARCHead compresses the remaining dense head with minimal extra CE.The hybrid setup targets heads left dense by the underlying backend.
  • H Ablation Details: 1.134 relative PPL for ARCHead core only falls to 1.007 after adding the activation-metric residual correction branch.The ablation supports optimizing the correction in covariance-weighted logit space.
  • H Ablation Details: Table 17’s expanded ablation details highlight the importance of the correction branch.The reported relative-PPL change shows that the quantized core alone is insufficient.
  • I Throughput Measurement Details: ARCHead does not bottleneck generation throughput, and the small positive difference in the BNB NF4 configuration is attributed to measurement noise.Throughput is measured independently from persistent memory.
  • J Additional Result Visualizations: ARCHead occupies a storage range similar to group INT4 while remaining close to the dense head in the storage–perplexity trade-off.The comparison is shown for the Qwen3-8B-Base head.
  • J Additional Result Visualizations: Below 2% throughput differences between BF16 and packed ARCHead heads are treated as negligible measurement variation.This visualization summarizes measured generation throughput with either head representation.
  • J Additional Result Visualizations: Relative PPL remains in a narrow band from 4K through 64K calibration tokens.The calibration-sensitivity visualization covers Qwen3-8B-Base.

K Failure Modes and Limitations · L Preliminary Future Work: Fused Triton Kernel Integration

ARCHead’s benefits depend on head density, model sensitivity, and whether prior LM-head compression already works well, while its evaluation and fused-kernel evidence remain limited. A preliminary fused Triton implementation avoids materializing the full FP16 head and shows differences within 2% on one GPU.

  • K Failure Modes and Limitations: ARCHead is most beneficial when the LM-head constitutes a large dense parameter block.
  • K Failure Modes and Limitations: ARCHead’s relative advantage diminishes when naive INT4 head quantization already performs exceptionally well, as observed to some degree with Gemma.
  • K Failure Modes and Limitations: Successful native LM-head quantization by a backend may leave ARCHead unable to provide additional quality improvements.
  • K Failure Modes and Limitations: ARCHead is developed for LM-head compression, not for compressing Feed-Forward Networks (FFN/MLP) in this paper.
  • K Failure Modes and Limitations: WikiText perplexity and logit fidelity are primary measures, while three downstream tasks provide only a limited sanity check and exhaustive capability and safety evaluation remains future work.
  • K Failure Modes and Limitations: The included fused-kernel result is preliminary, limited to one accelerator, and leaves broader hardware and production backend optimization as future work.
  • L Preliminary Future Work: Fused Triton Kernel Integration: A custom Triton kernel preliminarily fuses INT8 dequantization with low-rank residual evaluation, loading the quantized core and row-wise scales directly in one pass.
  • L Preliminary Future Work: Fused Triton Kernel Integration: Within 2%, the fused-kernel throughput differences on one NVIDIA RTX Pro 6000 are treated as noise, supporting no practical inference-generation bottleneck with AWQ or bitsandbytes NF4 backends.The fused pass avoids the prohibitive memory-bandwidth overhead of sequentially instantiating the full FP16 matrix before multiplication.
Loading 2608.02703v1…