💻 computer science

CLHA: Cross-Layer Hessian Aggregation for Quantizing Weight-Shared Mixture-of-Experts Transformers

This paper introduces CLHA, a cross-layer Hessian aggregation method for post-training quantization of weight-shared Mixture-of-Experts transformers that minimizes total reconstruction error by combining Hessian statistics across shared layers, significantly outperforming existing per-layer calibration approaches while also addressing critical implementation pitfalls in Hessian estimation.

Original authors: Kunal Dhanda

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Kunal Dhanda

Original paper licensed under CC BY 4.0 (https://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

The Big Picture: Shrinking a Giant Library

Imagine a massive library (a Large Language Model) that is too big to fit on a small shelf. To make it fit, we need to "compress" the books. In the world of AI, this compression is called quantization. It's like taking a high-definition photo and shrinking it to a thumbnail. If you do it poorly, the picture becomes blurry and unrecognizable.

This paper tackles a specific type of library called a Mixture-of-Experts (MoE) model. These libraries are special because they don't read every book for every question; they have "experts" (specialized sections) that only open when needed.

Recently, engineers invented a clever trick called Cross-Layer Weight Sharing (CLWS). Imagine two different floors in the library (Layer A and Layer B) that use the exact same set of reference books for their experts. This saves a huge amount of space because you don't need to buy two copies of the same books.

The Problem:
When people tried to shrink (quantize) these libraries, they made a mistake. They treated the shared books as if they only belonged to the first floor. They looked at the questions asked on Floor A, decided how to shrink the books, and applied that shrinkage to the books on Floor B as well.

But here's the catch: The people asking questions on Floor A are different from those on Floor B. The "vibe" of the questions changes as you go up the stairs. By only looking at Floor A, the shrinkage plan was wrong for Floor B, causing the library to lose its ability to answer questions correctly.

The Solution: CLHA (The "Double-Check" System)

The authors propose a new method called CLHA (Cross-Layer Hessian Aggregation).

The Analogy: The Tailor and the Twin Suits
Imagine a tailor making suits for a pair of identical twins.

The Old Way (PLIC): The tailor measures Twin A, cuts the fabric, and assumes Twin B is exactly the same size. But Twin B has slightly different posture. The suit fits Twin A perfectly but is tight on Twin B.
The CLHA Way: The tailor measures both twins. They take the measurements from Twin A and Twin B, average them out (giving more weight to whoever stands in the room more often), and create a "super-suit" pattern that fits both twins perfectly.

In technical terms, the paper says you shouldn't just look at the data from one layer. You must combine the "sensitivity" (mathematically called the Hessian) of both layers. This creates a combined map that tells you exactly how to shrink the shared weights so they work well for both floors of the library.

The Special Upgrade: CLHA-MP

The paper also introduces a "Mixed Precision" version called CLHA-MP.

The Analogy: The VIP Lounge
In these libraries, there are two types of experts:

Routed Experts: They only open their doors for specific, rare questions.
Shared Experts: They are "always open" and handle the most common, everyday questions.

Because the "Shared Experts" are used so much more often, they are more sensitive to errors. If you shrink them too much, the whole library breaks.

The Fix: CLHA-MP gives these "always-open" experts a little extra space (one extra bit of precision). It's like giving the VIP lounge a slightly wider door while keeping the regular doors narrow. This tiny extra space makes a huge difference in how well the library works.

The Hidden Trap: The "Ghost" Hook

The authors also found a sneaky bug in the software tools (PyTorch) used to build these models.
The Analogy: Imagine two people trying to write in the same notebook at the same time. If they both use the same pen without a label, their notes get mixed up, and you can't tell who wrote what.
In the code, when two layers share the same expert, the software's "forward hooks" (tools that track activity) would mix the data from both layers together, corrupting the measurements. The authors fixed this by creating a "layer-aware" system that labels the notes so the tailor knows exactly which twin is which.

What Did They Prove?

They tested this on a small model trained on a dataset called WikiText-2.

At 2-bit compression (very small size): The old method (PLIC) made the model very confused (high "perplexity").
The CLHA method: Made the model significantly smarter (4.3% to 5.4% better).
The CLHA-MP method: Made it even smarter (18.6% better) by giving the "always-open" experts that extra bit of space.

They also ran a test where they made the difference between the two floors (the layers) more extreme. The old method got worse and worse, but the new method stayed stable, proving that their "double-check" system is essential when the two floors are very different from each other.

Summary

The Issue: Sharing weights between layers saves space, but standard compression tools ignore the second layer, causing errors.
The Fix: Combine the data from both layers to create a better compression plan (CLHA).
The Bonus: Give the most-used experts a tiny bit more precision (CLHA-MP).
The Result: Much smarter, smaller models that don't lose their "brainpower" when shrunk down.

Technical Summary: CLHA – Cross-Layer Hessian Aggregation for Quantizing Weight-Shared Mixture-of-Experts Transformers

1. Problem Statement

Recent efficient Mixture-of-Experts (MoE) transformer architectures, such as those in Megrez2, employ Cross-Layer Weight Sharing (CLWS). In this design, adjacent transformer layers (e.g., $\ell$ and $\ell+1$ ) reuse the same routed expert weight tensors to halve the parameter footprint while maintaining the same number of activated parameters per token.

However, existing Post-Training Quantization (PTQ) methods, including GPTQ-style approaches, treat these shared weights as if they belong to a single layer. They calibrate the quantization using activation statistics from only one "primary" layer. This creates a fundamental mismatch:

The residual stream representation evolves between layers, meaning the token distributions activating the shared expert at layer $\ell_a$ differ from those at layer $\ell_b$ .
Existing methods ignore the activation statistics of the secondary layer, leading to suboptimal reconstruction error minimization, particularly in low-bit regimes (e.g., 2-bit).

2. Methodology

2.1 Theoretical Foundation: Cross-Layer Hessian Aggregation

The paper formalizes the objective for quantizing a weight $W$ shared across layers $\ell_1, \dots, \ell_n$ . The correct objective is to minimize the total reconstruction error across all layers using the weight:
$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$
where $H_{\ell_i}$ is the Hessian (activation covariance) of layer $\ell_i$ , and $\alpha_i$ represents the routing frequency (proportion of tokens) for that layer.

Proposition 1 establishes that minimizing this combined error is mathematically equivalent to applying GPTQ with a combined Hessian:
$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$
This contrasts with Per-Layer Independent Calibration (PLIC), which minimizes error using only a single layer's Hessian ( $H_{\ell_a}$ ).

2.2 CLHA Algorithm

The authors propose CLHA (Cross-Layer Hessian Aggregation), a three-pass calibration algorithm:

Collection: For each calibration sequence, accumulate per-layer Hessians and token counts for every active expert.
Aggregation: For each shared expert pair $(\ell_a, \ell_b)$ $(ℓ_{a}, ℓ_{b})$ , compute a combined Hessian. The paper evaluates two weighting strategies:
- Routing-Frequency Weighting ( $\alpha = n_a / (n_a + n_b)$ ): Weights Hessians by the actual token count received by each layer.
- Equal Weighting ( $\alpha = 0.5$ ): Acts as a regularizer, useful when calibration data is limited and secondary-layer Hessian estimates are noisy.
Quantization: Apply standard GPTQ column-sequential updates using the aggregated $H_{comb}$ .

2.3 CLHA-MP (Mixed Precision)

Recognizing that always-active shared experts receive all $N$ tokens (unlike routed experts which receive a fraction $K/E$ ), the paper notes their quantization sensitivity is $E/K$ times higher. CLHA-MP allocates $b+1$ bits to these shared experts while keeping routed experts at $b$ bits, incurring minimal storage overhead (10–20% of expert parameters).

2.4 Implementation Correction: Hook Ambiguity

The paper identifies a critical engineering pitfall: standard PyTorch forward hooks attached to a shared ExpertFFN module fire for any layer calling that module. This causes activations from different layers to be mixed into a single Hessian accumulator, corrupting per-layer attribution.

Solution: The authors propose a layer-aware callback mechanism passed through the MoELayer.forward() signature. This explicitly provides layer indices to the callback, ensuring separate Hessians are accumulated for $H_a$ and $H_b$ even when the underlying weight module is shared.

3. Key Contributions

Theoretical Proposition: Formalized that the optimal Hessian for shared weights is the routing-frequency-weighted sum of per-layer Hessians.
CLHA Algorithm: A practical implementation of combined Hessian calibration within the GPTQ framework.
CLHA-MP: A mixed-precision extension assigning an extra bit to highly sensitive shared experts.
Implementation Fix: Identified and resolved the "hook ambiguity" issue that silently corrupts calibration in shared-module architectures.
Empirical Validation: Demonstrated that gains from CLHA correlate with cross-layer routing divergence.

4. Experimental Results

Experiments were conducted on a 34.5M-parameter CLWS-MoE trained from scratch on WikiText-2.

Performance at 2-bit: This is the critical regime where miscalibration compounds.
- CLHA-rt (routing-weighted) reduced Perplexity (PPL) by 4.3% compared to PLIC.
- CLHA-eq (equal-weighted) reduced PPL by 5.4%.
- CLHA-MP achieved an 18.6% improvement over PLIC.
- At 4-bit and 3-bit, differences were within noise margins (<1%).
Calibration vs. Test Error: Interestingly, PLIC achieved lower Calibration Hessian Error (CHE) on the training set because it optimizes directly for that distribution. However, CLHA generalized better to the test set, indicating that the combined Hessian captures the true distribution of the shared weight better than a single-layer view.
Routing Sharpness Experiment: By artificially sharpening router logits (increasing cross-layer divergence), the authors observed:
- PLIC performance degraded by 5% as divergence increased.
- CLHA remained stable (<1% degradation).
- The performance gap between CLHA and PLIC grew from +4.3% (natural routing) to +8.0% (high divergence), confirming the theoretical prediction that CLHA gains scale with cross-layer divergence.

5. Significance and Claims

The paper claims to provide the first PTQ analysis specifically for cross-layer weight sharing in MoE models. Its significance lies in:

Correcting the Objective: Demonstrating that treating shared weights as single-layer entities is mathematically suboptimal and that Hessian aggregation is the necessary correction.
Scalability of Gains: Showing that as models become larger and routing patterns more specialized (increasing cross-layer divergence), the benefits of CLHA will likely grow, making it essential for future production CLWS-MoE models (e.g., Megrez2, DeepSeek, Qwen).
Practical Necessity: Highlighting that naive application of existing PTQ tools to CLWS architectures is fundamentally broken due to hook ambiguity, and providing the necessary fix for any future calibration method in this domain.

The authors remain modest regarding absolute performance, noting their model is small (34.5M) and trained for limited steps, but emphasize that the relative improvements over standard calibration are robust and theoretically grounded.