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

Published at IEEE International Joint Conference on Neural Networks (IJCNN) 2025

Introduction & Motivation — Why compression?

Fig.1 The number of parameters in a deep neural network grows exponentially in modern AI systems.

As shown in Fig.1, since 1950s, the number of parameters in modern AI systems has exploded drastically as we scale the size of neural networks, with GPT-4 having over 1 trillion parameters. Storing, processing and moving this much data has become incredibly expensive computationally and environmentally. This indicates that brute-force, data-driven engineering is becoming increasingly unsustainable. If we want to build AI systems that run efficiently and affordably, we need smarter, more efficient and sustainable solutions. Therefore, model compression comes into play. It is not just a techical optimization, but is becoming a fundamental requirement for scaling AI is a practical and responsible way.

How to compress?

Fig.2 Post-training compression techniques

Fig.2 illustrates four post-training compression techniques, which are quantization, pruning, distillation, and low-rank approximation, respectively. Quantization involves using lower precision (bits) to represent the model weights, and pruning aims to remove (set to zero) some small or less significant weights. Distillation works by training a much smaller model that mimics the teacher model’s performance, based on logits (teacher output), features (hidden representation), etc. Finally, low-rank approximation removes less significant singular values and corresponding eigenvectors of the weight matrix, and this will be the focus of this work.

One of the papers that utilizes low-rank structure within a single weight matrix to compress the LLM is LASER (Layer-Selective Rank Reduction). The authors selectively compress 1 matrix in the feed-forward network (FFN) block in 1 layer of the transformer using singular value decomposition (SVD). Surprisingly, by carefully choosing which layer and matrix to compress, the performance on natural language understanding (NLU) tasks even improves, while achieving parameter compression.

Inspired by LASER which exploits the intra-matrix redundancies, we now want to explore inter-matrix redundancies in the attention block of the transformer using tensor decomposition. But, before getting into the technical details, let’s review some preliminary concepts in tensor algebra.

Tensor preliminaries

Fig.3 Tensor introduction

A tensor is a high dimensional array. A scalar is a 0-dimensional (0D) tensor. Stacking several scalars together forms a vector, which is a 1-dimensional (1D) tensor. Likewise, a matrix is a 2-dimensional (2D) tensor. There are many cases in our real lives which are naturally higher-order tensor. For example, a colorful image is naturally a 3D tensor, with the 1st dimension representing height, 2nd dimension representing width, and the 3rd dimension denoting color channels (E.g. RGB).

Tucker Decomposition (TD)

Fig.4 Tucker decomposition

TD decomposes a higher-order tensor into factor matrices and a much smaller tensor core. This can be considered as a generalization of SVD into a higher dimension, drastically reducing the space complexity from exponential to linear w.r.t. the dimension of each tensor mode. TD is just one form of tensor decomposition, and it illustrates how we can effectively bypass the Curse of Dimensionality using tensor algebra.

TensorLLM

Now, with the preliminary knowledge of tensor algebra, we can smartly reduce both inter-matrix and intra-matrix redundancies in the attention block of tranformers. Our design intuition is as follows:

1. Attention heads within the same layer capture the same level of patterns.
2. Different attention heads within the same layer learn different specialised knowledge.

Fig.5 TensorLLM architecture

In the vanilla transformer architecture, multi-head attention weight matrices are naturally 3D tensors of shape (embedding_dim, num_heads, head_dim). However, for the computational efficiency (e.g. GPU acceleration of matrix multiplication), those attention weights are stored in 2D matrices by folding the last 2 modes, yielding matrices of shape (embedding_dim, num_headshead_dim*).

In TensorLLM, we begin by splitting (unfolding) attention matrices, into 3D tensors, i.e. (embedding_dim, num_headshead_dim) → (embedding_dim, num_heads, head_dim). 4 such 3D tensors are then stacked together to form a 4D tensor of shape (embedding_dim, head_dim, 4, num_heads). Finally, we utilize Tucker decomposition with shared factor matrices to reduce redundancies and noises in this 4D tensor. Note that the last mode (num_heads*) is not decomposed, so that Tuker factor matrices are shared across all heads. This decomposition is shown mathematically as:

Eq.1 Tucker decomposition with shared factor matrices

But, why do we use shared factor matrices, and how does this align with the design intuition?

Based on domain knowlege of attention mechanism and our intuition mentioned above, we conjecture that the weights of the attention block in a single transformer layer contain reasoning-related information in a subspace which is shared across multiple attention heads. This common higher dimensional subspace can be characterized by the span of the Tucker factor matrices (Intuition 1). Within this shared subspace, each attention head is allowed to have its own learnt intrinsic knowlege. This specialized knowlege is represented by the Tucker core tensor, which encodes the variability information in the subspace designated by the factor matrices (Intuition 2).

Notably, our proposed TensorLLM is a plug-and-play module for attention block compression. It can be used together with other compression methods applied to FFN block, such as LASER.

Results

Tab.1 Performance comparison of TensorLLM-compressed model and the original pretrained model.

Table 1 compares the performance of pretrained and TensorLLM-compressed models on 4 different natural language understanding datasets. TensorLLM-compressed model outperforms the original model in terms of accuracy across 3 models and 4 datasets, while achieving compression ratios of up to ~250 times in the attention block.

To demonstrate that TensorLLM, when used in conjunction with other compression methods applied to FFN block, further improves the reasoning capabilities, we evaluate the following 3 cases:

1. LASER applied to 1 matrix in the FFN block
2. LASER applied to all matrices in the FFN and MHA blocks
3. TensorLLM (our method) applied to the MHA block; LASER applied to matrices in the FFN block

Tab.2 Performance of stand-alone and hybrid methods for LLM compression and reasoning.

Table 2 shows the stand-alone and hybrid methods in the 3 cases mentioned above. Clearly, TensorLLM (our method) applied to attention block, together with LASER applied to FFN block, yields the best overall performance, outperforming the configuration where LASER is applied to both attention and FFN blocks.

Finally, to verify the effectiveness of stacking 4 attention matrices together before compression, we tried compressing 4 matrices alone, and the results are listed in Table 3. It is obvious that compressing the stacked tensor gives better performance than compressing individual matrix.

Tab.3 The impact of compressing attention matrices separately and together.

Conclusion & Future work

In this paper, we proposed TensorLLM, a novel framework that compresses the LLM through tensor decomposition (a special variant of the Tucker decomposition). This approach denoises the pretrained weights, and gives better reasoning capabilities, which leads to better performance on downstream tasks. Our findings highlight the power of tensor structure in efficiently representing information, and their potential in LLM compression. In the future, we aim to extend TensorLLM to larger and more recent chain-of-though (CoT) reasoning models, as well as tasks beyond natural language understanding.

⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐

If you are interested in more details, feel free to read our paper and code!

📄https://arxiv.org/abs/2501.15674

🖥️https://github.com/guyuxuan9/TensorLLM

References

Sharma, P., Ash, J.T. and Misra, D., The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction. In The Twelfth International Conference on Learning Representations.

Cichocki, A., Mandic, D., De Lathauwer, L., Zhou, G., Zhao, Q., Caiafa, C. and Phan, H.A., 2015. Tensor decompositions for signal processing applications: From two-way to multiway component analysis. IEEE signal processing magazine, 32(2), pp.145–163.