Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TensorLLM: Tensorising Multi-Head Attention for Enhanced Reasoning and Compression in LLMs

Yuxuan Gu, Wuyang Zhou, Giorgos Iacovides, Danilo Mandic

TL;DR

This work addresses the challenge of enhancing reasoning in LLMs while compressing their parameters by focusing on the multi-head attention (MHA) block. It introduces a novel framework that tensorises per-head MHA weights and applies a Tucker decomposition with shared factor matrices across heads, enabling denoising within a common high-dimensional subspace and achieving substantial compression without extra data or training. Empirical results across encoder-only and decoder-only models on four reasoning datasets show consistent improvements in accuracy and up to approximately 250x compression in the MHA parameters, and the method can be combined with existing FFN-denoising techniques for further gains. The approach provides a practical, training-free path to more efficient and capable LLM inference, with ablation studies validating the importance of joint tensorisation across all MHA weight matrices.

Abstract

The reasoning abilities of Large Language Models (LLMs) can be improved by structurally denoising their weights, yet existing techniques primarily focus on denoising the feed-forward network (FFN) of the transformer block, and can not efficiently utilise the Multi-head Attention (MHA) block, which is the core of transformer architectures. To address this issue, we propose a novel intuitive framework that, at its very core, performs MHA compression through a multi-head tensorisation process and the Tucker decomposition. This enables both higher-dimensional structured denoising and compression of the MHA weights, by enforcing a shared higher-dimensional subspace across the weights of the multiple attention heads. We demonstrate that this approach consistently enhances the reasoning capabilities of LLMs across multiple benchmark datasets, and for both encoder-only and decoder-only architectures, while achieving compression rates of up to $\sim 250$ times in the MHA weights, all without requiring any additional data, training, or fine-tuning. Furthermore, we show that the proposed method can be seamlessly combined with existing FFN-only-based denoising techniques to achieve further improvements in LLM reasoning performance.

TensorLLM: Tensorising Multi-Head Attention for Enhanced Reasoning and Compression in LLMs

TL;DR

This work addresses the challenge of enhancing reasoning in LLMs while compressing their parameters by focusing on the multi-head attention (MHA) block. It introduces a novel framework that tensorises per-head MHA weights and applies a Tucker decomposition with shared factor matrices across heads, enabling denoising within a common high-dimensional subspace and achieving substantial compression without extra data or training. Empirical results across encoder-only and decoder-only models on four reasoning datasets show consistent improvements in accuracy and up to approximately 250x compression in the MHA parameters, and the method can be combined with existing FFN-denoising techniques for further gains. The approach provides a practical, training-free path to more efficient and capable LLM inference, with ablation studies validating the importance of joint tensorisation across all MHA weight matrices.

Abstract

The reasoning abilities of Large Language Models (LLMs) can be improved by structurally denoising their weights, yet existing techniques primarily focus on denoising the feed-forward network (FFN) of the transformer block, and can not efficiently utilise the Multi-head Attention (MHA) block, which is the core of transformer architectures. To address this issue, we propose a novel intuitive framework that, at its very core, performs MHA compression through a multi-head tensorisation process and the Tucker decomposition. This enables both higher-dimensional structured denoising and compression of the MHA weights, by enforcing a shared higher-dimensional subspace across the weights of the multiple attention heads. We demonstrate that this approach consistently enhances the reasoning capabilities of LLMs across multiple benchmark datasets, and for both encoder-only and decoder-only architectures, while achieving compression rates of up to times in the MHA weights, all without requiring any additional data, training, or fine-tuning. Furthermore, we show that the proposed method can be seamlessly combined with existing FFN-only-based denoising techniques to achieve further improvements in LLM reasoning performance.
Paper Structure (24 sections, 14 equations, 2 figures, 4 tables)

This paper contains 24 sections, 14 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Mathematical Notations
  4. Tensor Preliminaries
  5. Definition
  6. Mode-$n$ product
  7. Singular Value Decomposition (SVD)
  8. Tucker decomposition
  9. Transformer Architecture
  10. Multi-Head Attention (MHA)
  11. Proposed Methodology
  12. Multi-head Tensorisation
  13. Tucker Decomposition with Shared Factor Matrices
  14. Experiments
  15. LLM Models
  16. ...and 9 more sections

Figures (2)

  • Figure 1: Structures of our proposed framework for the compression of the MHA block and the existing methods (such as LASER) which apply to the FFN block. Left: The three-step denoising process applied to the MHA block (our method): (1) split the weight matrices into multiple heads, (2) tensorise the matrices of each attention head into a 3D tensor, and (3) apply Tucker decomposition to this set of 3D tensors with a common set of factor matrices to perform denoising. Middle: A standard (vanilla) decoder-only or encoder-only transformer architecture. Right: Existing methods applied to the FFN block; an illustration of LASER.
  • Figure 2: Tensor network orus2014practical topologies of different decomposition methods applied to the MHA weights in a single transformer layer. Note that while both LASER and TRAWL reported performances when only applied to the FFN blocks, they could also be applied to the MHA weights. The symbol $d_{model}$ represents the embedding dimension, $h$ is the number of attention heads, while $d_v = \frac{d_{model}}{h}$ stands for the head dimension. ($\mathbf{a}, \mathbf{d}$) LASER sharma2023truth: decomposition of a single weight matrix into $\mathbf{U} \in \mathbb{R}^{d_{model} \times R}$, $\mathbf{V} \in \mathbb{R}^{h \cdot d_v \times R}$, and a diagonal matrix $\mathbf{\Sigma} \in \mathbb{R}^{R \times R}$. ($\mathbf{b},\mathbf{e}$) TRAWL luo2024trawl: decomposition of a 3D tensor into the factor matrices $\mathbf{A} \in \mathbb{R}^{4 \times R_3}, \mathbf{B} \in \mathbb{R}^{h \cdot d_v \times R_2}$, and $\mathbf{C} \in \mathbb{R}^{d_{model} \times R_1}$, along with a core $\mathcal{G} \in \mathbb{R}^{R_1 \times R_2 \times R_3}$. ($\mathbf{c}, \mathbf{f}$) Our method -- Tucker decomposition with shared factor matrices: decomposing a set of 3D tensors using Tucker decomposition while having a common set of shared factor matrices $\mathbf{U}^{(1)} \in \mathbb{R}^{d_{model} \times R_1}, \mathbf{U}^{(2)} \in \mathbb{R}^{d_{v} \times R_2} , \mathbf{U}^{(3)} \in \mathbb{R}^{4 \times R_3}$, with $\mathcal{G}_{all} \in \mathbb{R}^{R_1 \times R_2 \times R_3 \times h}$ as the core tensor.

Theorems & Definitions (3)

  • Remark 1
  • Remark 2
  • Remark 3