Activation aware preconditioning improves tensor compression of large AI models
Activation-Aware Weight Tensorization: A Calibration-Time Preconditioner for Tensor-Network LLM Compression
Machine LearningArtificial Intelligence
Summary
Large language models use huge matrices to handle language tasks, but these can be compressed using tensor techniques to save space. The usual methods care only about matching the stored numbers, not how these changes affect the model’s output with real inputs. The authors introduce a simple pre-step that adjusts the weights based on how the model actually behaves, improving how well the compressed model works without retraining. This approach helps keep the model’s results closer to the original, even after strong compression, and works across different large models and tasks.
What this means in practice
- •For machine learning engineers: Enhance post-training compression of large language models to reduce size while preserving output quality without retraining.
- •For cloud ai service teams: Deploy large language models more efficiently by integrating activation-aware tensor compression to save memory and speed up inference.
- •For software developers for ai assistants: Improve responsiveness and reduce resource consumption of deployed language models using activation-aware weight tensorization compression.$Commercial implications: Enables smaller, faster large language model components in AI assistant apps sold to enterprises and consumers.
Authors
Alessandro Beatini, Marco Maronese, Emanuele Rodolà
Abstract
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.