PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuxuan Gu, Wuyang Zhou, Giorgos Iacovides, Danilo Mandic

arXiv:2501.15674v2cs.CLcs.LG

TL;DR

Existing denoising methods primarily benefit FFN weights, while comparable gains from applying them to MHA have not been observed, [16]. TensorLLM applies multi-head tensorisation and Tucker decomposition to impose a shared higher-dimensional subspace across attention heads, and reports improved reasoning across model scales and encoder-only and decoder-only architectures.

  • Problem

    Existing denoising approaches achieve their strongest gains on FFN weights, while similar gains have not been observed when denoising MHA weights, [16].

    Source: §I · p. 1

  • Method

    The framework denoises and compresses MHA weights using multi-head tensorisation and Tucker decomposition to impose a shared higher-dimensional low-rank structure across attention heads.

    Source: §I · p. 2

  • Results

    The method improves reasoning in both uncompressed LLMs and models already denoised in FFN layers, spanning millions-to-billions of parameters and encoder-only to decoder-only architectures.

    Source: §I · p. 3

  • Takeaways & Limitations

    The method can be combined with FFN-weight denoising methods for further reasoning gains, without additional data, training, or fine-tuning.

    Source: §I · p. 3, §V · p. 8

    The best hyperparameter settings vary across datasets, and finding unified, generalisable settings remains future work.

    Source: §V · p. 8

Abstract

from arXiv · show

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.

I. INTRODUCTION

Prior denoising approaches focused on FFN weights and did not show similar gains when applied to MHA; TensorLLM addresses this by denoising MHA weights through shared structure across attention heads. Across four benchmarks and encoder-only and decoder-only models, the method improves reasoning without additional data, training, or fine-tuning.

I. INTRODUCTION

  • Post-training compression seeks to make overparameterized LLMs more practical while avoiding large compromises in inference performance and generative capabilities.The introduction frames this as a goal for resource-constrained environments.Source: §I · p. 1
  • Prior denoising methods showed their strongest gains on FFN weights, not MHA, despite MHA being central to transformer architectures.The authors attribute this gap to prior methods not leveraging the coherence of multiple attention heads within a layer.Source: §I · p. 1
  • The approach tests whether enforcing a shared higher-dimensional subspace across attention-head weights can improve LLM reasoning.This question builds on the conjecture that a layer’s MHA weights contain reasoning-related information in a subspace shared across heads.Source: §I · p. 2
  • TensorLLM tensorises multiple attention heads and applies a Tucker variant with shared factor matrices to denoise MHA weights in a common higher-dimensional subspace.Each head can encode different information within that shared subspace.Source: §I · p. 2
  • The framework can be combined with existing methods that denoise FFN layers to achieve further gains in reasoning.The paper evaluates models and datasets spanning multiple millions to billions of parameters and encoder-only to decoder-only architectures.Source: §I · pp. 2–3
  • Across four benchmarks and encoder-only and decoder-only models, the method improves reasoning in uncompressed LLMs and alongside FFN denoising, without additional data, training, or fine-tuning.The evaluated models range from multiple millions to billions of parameters.Source: §I · p. 3

II. BACKGROUND

The background introduces the paper’s mathematical notation and tensor preliminaries. It defines tensors and describes how stacking lower-dimensional arrays forms higher-dimensional tensors.

A. Mathematical Notations

B. Tensor Preliminaries

  • A tensor is a multidimensional array of order N, with mode n having size I_n.The passage defines the order as the number of modes and indexes modes from 1 to N.Source: B. Tensor Preliminaries · p. 3
  • Stacking scalars, vectors, and matrices illustrates how lower-dimensional tensors form vectors, matrices, and 3D tensors, respectively.The passage generalizes this iterative stacking process to higher-dimensional tensors.Source: B. Tensor Preliminaries · p. 3
  • A mode-n product combines a tensor A with a matrix B to produce a tensor C whose n-th mode has size J_n instead of I_n.The supplied text gives the output dimensions as I_1 through I_(n-1), J_n, and I_(n+1) through I_N.Source: B. Tensor Preliminaries · p. 3

C. Singular Value Decomposition (SVD)

SVD approximates a matrix with rank-1 factors and denoises it by discarding factors associated with the smallest singular values. Tucker decomposition extends this approach to tensors, using factor matrices and a smaller core tensor whose multilinear ranks can be set for compression.

C. Singular Value Decomposition (SVD)

D. Tucker decomposition

  • Tucker decomposition generalizes SVD by representing a tensor with factor matrices and a smaller core tensor.The factor matrices characterize projections onto subspaces, while the multilinear ranks are treated as hyperparameters.Source: D. Tucker decomposition · p. 3
  • Setting Tucker multilinear ranks much smaller than the original dimensions reduces the parameter count.Efficiently identifying the optimal ranks remains an active area of research.Source: D. Tucker decomposition · p. 3

E. Transformer Architecture

Transformer layers combine multi-head attention (MHA) and feed-forward network (FFN) blocks, with layer normalization and residual connections between them. The MHA weights can be represented as tensor networks to compare LASER, TRAWL, and the paper’s shared-factor Tucker decomposition.

E. Transformer Architecture

  • Each transformer layer contains an MHA block and an FFN block, with layer normalization and residual connections between them.The described architectures include encoder-only and decoder-only models.Source: E. Transformer Architecture · p. 3
  • The figure compares LASER’s single-matrix decomposition, TRAWL’s 3D-tensor decomposition, and Tucker decomposition with shared factor matrices.The diagrams represent tensors as circles and connect indices that are summed over.Source: E. Transformer Architecture · p. 4, Figure 2 · p. 4
  • MHA concatenates the outputs of its attention heads and projects them back into the original latent space.The attention heads independently learn distinct input representations.Source: E. Transformer Architecture · p. 4

III. PROPOSED METHODOLOGY

The method tensorises MHA weights and applies Tucker decomposition with shared factor matrices across attention heads. This places the head weights in a common subspace, and the framework is reported to improve reasoning while compressing MHA blocks.

III. PROPOSED METHODOLOGY

  • The method tensorises each attention head’s MHA weights and applies Tucker decomposition with shared factor matrices.The factor matrices are common across the decompositions for the heads in one transformer layer.Source: §III · p. 4
  • Shared factor matrices place the different attention-head weights in a common subspace.The framework uses a shared higher-dimensional low-rank structure to denoise attention weights.Source: §III · p. 4
  • The framework is reported to improve LLM reasoning while compressing the MHA blocks.Source: §III · p. 4

A. Multi-head Tensorisation

Multi-head tensorisation converts the MHA’s global 2D projection matrices into per-head 3D tensors, then stacks them into a 4D tensor. This prepares the weights for Tucker decomposition.

A. Multi-head Tensorisation

  • The method folds the global 2D query, key, value, and output projection matrices into higher-dimensional formats.It first splits each global matrix into local submatrices belonging to individual attention heads.Source: A. Multi-head Tensorisation · pp. 4–5
  • For each head, the four 2D submatrices are tensorised into a 3D tensor.The process is repeated across all heads.Source: A. Multi-head Tensorisation · p. 5
  • Stacking the per-head tensors produces a 4D tensor containing all MHA weights in a transformer layer.The higher-dimensional representation prepares the weights for Tucker decomposition.Source: A. Multi-head Tensorisation · p. 5

B. Tucker Decomposition with Shared Factor Matrices

The method tensorizes each attention head’s weights and applies Tucker decomposition with shared factor matrices, placing head weights in a common subspace while retaining head-specific cores. The resulting low-rank approximation structurally denoises and compresses MHA weights.

B. Tucker Decomposition with Shared Factor Matrices

  • Tucker decomposition represents a higher-dimensional tensor with a smaller core and factor matrices that designate a shared subspace for attention weights.The core represents variability within the subspace designated by the factors.Source: B. Tucker Decomposition with Shared… · p. 5
  • Each attention head’s 3D weight tensor is decomposed using the same three factor matrices, while retaining its own Tucker core.The shared factors define the common subspace; each head’s core contains its variability information.Source: B. Tucker Decomposition with Shared… · p. 5
  • The shared-factor formulation is a special Tucker decomposition of the tensor containing all attention weights in one layer.Its 4D core has a separate slice for each attention head.Source: B. Tucker Decomposition with Shared… · p. 5
  • Given multilinear ranks, denoising approximates the original attention-head weight tensors by minimizing their decomposition error.The approximation is applied to the tensorized MHA weights.Source: B. Tucker Decomposition with Shared… · p. 5

IV. EXPERIMENTS

Experiments evaluate the framework on four reasoning datasets and three LLMs, including encoder-only and decoder-only architectures. They use layer-selective application and compare a hybrid with LASER against alternative compression settings.

IV. EXPERIMENTS

  • Experiments cover four reasoning datasets and three LLMs spanning encoder-only and decoder-only architectures.Table II compares RoBERTa, GPT-J, and LLaMA2 with and without the method across four datasets.Source: §IV · p. 5, Table II · p. 6
  • The framework is applied to one transformer layer at a time for layer-selective comparisons.This follows the layer-selective approach used in.Source: §IV · p. 5
  • The experiments assess combining the MHA method with LASER on FFN weights and compare it with LASER-only settings and separate Q, K, V, and output tensor denoising.The ablation compares the proposed approach with denoising separate query, key, value, and output weight matrices.Source: §IV · pp. 5–6

A. LLM Models

The evaluation uses RoBERTa 125M, GPT-J 6B, and LLaMA 2 7B. RoBERTa is encoder-only, while GPT-J and LLaMA 2 are decoder-only.

A. LLM Models

  • The proposed method was tested on three LLM models: RoBERTa 125M [32], GPT-J 6B, and LLaMA 2 7B [4].Source: A. LLM Models · p. 6
  • RoBERTa uses an encoder-only setup that predicts missing tokens, with five mask tokens appended to each question.The five <mask> tokens are added before the question is input to the model.Source: A. LLM Models · p. 6
  • GPT-J and LLaMA 2 are decoder-only models that generate tokens autoregressively from preceding tokens.They predict the next token based on the preceding tokens.Source: A. LLM Models · p. 6

B. Datasets

The evaluation uses four reasoning datasets and reports improved test accuracy across three models and all four datasets, with MHA compression up to 247.3 times. A hybrid of the proposed MHA method and LASER on FFN weights has the highest test accuracy in all but one model-dataset combination.

B. Datasets

  • HotPotQA questions come from its training and validation sets, and a generated answer counts as correct when it contains the answer.Inputs exceeding 15 tokens are discarded, and prompts use the form “<question> The answer is”.Source: B. Datasets · p. 6
  • FEVER evaluates fact-checking with binary true/false labels, using validation and test question-answer pairs after conflicting-label claims are filtered.A claim is classified as true when its probability exceeds the probability of false.Source: B. Datasets · p. 6
  • Bios Profession evaluates profession prediction from biographies, using validation examples for the same 10 professions as LASER.The model selects the profession with the highest probability among the 10 possibilities.Source: B. Datasets · p. 6
  • BigBench-WikidataQA evaluates question answering using structured Wikidata knowledge, with question-answer pairs from its train and validation sets.A generated answer counts as correct when it appears in the generated text.Source: B. Datasets · p. 6

C. Experimental Results

  • Test accuracy improved across all four datasets for all three models, with MHA-weight compression up to 247.3 times.Table II compares accuracy, loss, and MHA compression rate for models with and without the method; test accuracy uses the last 20% of each dataset.Source: C. Experimental Results · p. 6
  • The hybrid using the proposed method on MHA and LASER on FFN achieved the highest test accuracy except for LLaMA2 on BigBench-WikidataQA.The comparison includes LASER on one FFN matrix, LASER on both MHA and FFN, and the proposed MHA method with LASER on FFN.Source: C. Experimental Results · p. 7, Table III · p. 7

D. Ablation study

Jointly tensorizing the MHA weights outperformed compressing individual weight types across four reasoning datasets. The conclusion also reports benefits across encoder-only and decoder-only models, gains from combining the method with FFN denoising, and dataset-dependent best hyperparameters.

D. Ablation study

  • The proposed method achieved the best performance across all four reasoning datasets when compressing all MHA weights together.The comparison evaluates joint compression against compressing query, key, value, or output weights separately.Source: D. Ablation study · p. 7, Table IV · p. 7

V. CONCLUSION

  • The framework applies shared higher-dimensional low-rank denoising to each attention head’s weights using common Tucker factor matrices.The conclusion describes this as multi-head tensorisation with a special variant of the Tucker decomposition.Source: §V · p. 7
  • The method enhances reasoning and compresses parameters in encoder-only and decoder-only LLMs ranging from hundreds of millions to billions of parameters.Source: §V · p. 7
  • Combining the framework with existing reasoning methods, including FFN-weight denoising, yields further performance gains.Source: §V · p. 8
  • The best results require different hyperparameter settings for different datasets.The authors identify unified, generalisable settings as future work.Source: §V · p. 8
Loading 2501.15674v2…