Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper introduces Activation-aware Weight Tensorization (AWT), a calibration-time preconditioner that scales weight matrices based on their activation distributions before tensor-network decomposition, aiming to minimize functional error rather than just Frobenius norm. AWT consistently improves tensor-network compression (TT/TTN) across Llama 3.1, Mistral, and Qwen2.5 models, closing a significant portion of the perplexity gap to dense baselines and showing gains in downstream tasks. The method's effectiveness is attributed to diagonal preconditioning acting as a robustness-modularity tradeoff, outperforming full-covariance oracles in functional fidelity.
A simple, training-free calibration wrapper boosts tensor-network LLM compression by 2-6x, closing up to 60% of the perplexity gap without altering the underlying decomposition solver.
Post-training tensor-network compression replaces Transformer linear layers with Tensor Train (TT) or Tree Tensor Network (TTN) operators, but standard decompositions minimize weight-space Frobenius error rather than functional error under the layer's activation distribution. We propose Activation-aware Weight Tensorization (AWT), a training-free calibration wrapper that preconditions each weight matrix with a diagonal activation-derived scale before an unchanged TT/TTN solver and deploys the result with only an input-side elementwise rescaling. Across Llama 3.1 8B, Ministral 8B, and Qwen2.5 7B, AWT consistently improves vanilla TT/TTN tensorization at 2-6 times compression: under single-operator replacement, AWT closes 12-35% of the WikiText perplexity gap to the dense baseline across the three model families and 2-6 times compression settings; while under multi-operator Llama suffix replacement it closes 27-60% across attention-group and all-seven-matrix settings. The gains also transfer to downstream HellaSwag and ARC-Challenge evaluations. We further show that diagonal preconditioning is a robustness-modularity tradeoff rather than a diagonal-covariance assumption: a dense full-covariance oracle wins its own weighted objective in 80/81 cases, yet diagonal AWT gives better held-out functional fidelity in 53/81 cases. Together, these results position AWT as a principled, modular preconditioner for improving functional fidelity in fixed TT/TTN compression pipelines without modifying the decomposition solver.