PaperlayerView exact revision on arXiv

Source-linked AI summary

A Tensorized Transformer for Language Modeling

Xindian Ma, Peng Zhang, Shuai Zhang, Nan Duan, Yuexian Hou, Dawei Song, Ming Zhou

arXiv:1906.09777v3cs.CLcs.LG

TL;DR

Transformer’s parameter-heavy multi-head attention limits deployment in resource-limited settings. The paper introduces Multi-linear attention using Block-Term Tensor Decomposition, low-rank approximation, and parameter sharing, and reports compressed models with improved results across evaluated tasks, especially language modeling.

  • Problem

    Transformer’s multi-head attention uses many parameters, motivating compression for resource-limited training and deployment.

    Source: §1 · pp. 1–2

  • Method

    Multi-linear attention uses Block-Term Tensor Decomposition to combine low-rank approximation and parameter sharing in a self-attention model.

    Source: §1 · p. 2

  • Results

    The method reduces parameters while achieving better experimental results than existing Transformer-based methods, particularly in language modeling.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    Multi-linear attention provides a tensorized alternative to multi-head attention for NLP models operating under limited-resource constraints.

    Source: §6 · p. 9, §5.4 · p. 9

    The model may suffer from overfitting when the number of cores is large in language modeling.

    Source: §6 · p. 9

Abstract

from arXiv · show

Latest development of neural models has connected the encoder and decoder through a self-attention mechanism. In particular, Transformer, which is solely based on self-attention, has led to breakthroughs in Natural Language Processing (NLP) tasks. However, the multi-head attention mechanism, as a key component of Transformer, limits the effective deployment of the model to a resource-limited setting. In this paper, based on the ideas of tensor decomposition and parameters sharing, we propose a novel self-attention model (namely Multi-linear attention) with Block-Term Tensor Decomposition (BTD). We test and verify the proposed attention method on three language modeling tasks (i.e., PTB, WikiText-103 and One-billion) and a neural machine translation task (i.e., WMT-2016 English-German). Multi-linear attention can not only largely compress the model parameters but also obtain performance improvements, compared with a number of language modeling approaches, such as Transformer, Transformer-XL, and Transformer with tensor train decomposition.

1 Introduction

Transformer’s parameter-heavy multi-head attention complicates deployment in resource-limited settings. The paper proposes Multi-linear attention, combining tensor-based low-rank approximation and parameter sharing to compress attention while retaining or improving performance.

  • Transformer’s multi-head attention contains many parameters, creating a compression challenge for resource-limited training and deployment.Source: §1 · pp. 1–2
  • The compressed attention can be integrated into Transformer encoder-decoder layers through tensor splitting, concatenation, and end-to-end training.Source: §1 · p. 2
  • The paper proves scaled dot-product attention can be linearly represented using orthonormal basis vectors.Source: §1 · p. 2
  • Multi-linear attention combines parameter sharing with low-rank approximation through Block-Term Tensor Decomposition.The method shares Q, K, and V factor matrices across multiple tensor blocks.Source: §1 · p. 2
  • Shared Q, K, and V factor matrices connect the three components and are intended to capture sufficient attention information.The paper also proves that Multi-linear attention can reconstruct the original Transformer attention.Source: §1 · p. 2
  • Experiments cover language modeling and neural machine translation, with higher compression ratios and promising results, especially for language modeling.The reported language-modeling results include One-Billion experiments where standard multi-head attention is compressed at higher ratios.Source: §1 · p. 2

2 Preliminaries

The preliminaries introduce tensors, Block-Term Tensor Decomposition, and Transformer’s scaled dot-product and multi-head attention. These concepts establish the tensor representation and parameter structure targeted by the proposed compression method.

Tensor

  • A tensor is treated as a multi-array; vectors and matrices are respectively first-order and second-order tensors.Source: §2.1 · p. 3

Block-Term Tensor Decomposition

  • Block-Term Tensor Decomposition combines CP and Tucker decomposition to represent a high-order tensor as P block terms.Each block contains a core tensor and factor matrices connected by tensor products.Source: §2.1 · p. 3
  • For the paper’s three-order tensors, P denotes the CP rank and R1, R2, and R3 denote the Tucker ranks.The paper assumes equal Tucker ranks: R=R1=R2=R3.Source: §2.1 · p. 3, Figure 1 · p. 3

Multi-head Attention

  • Transformer processes queries, keys, and values as matrices Q, K, and V in scaled dot-product attention.Source: §2.2 · p. 3
  • Multi-head attention concatenates multiple attention heads and applies an output projection.The formulation is MultiHeadAttention(Q, K, V ) = Concat(head1, . . . , headk)W O.Source: §2.2 · p. 3
  • Multiple parameter groups in multi-head attention produce a large number of redundant parameters.Source: §2.2 · p. 3

3 Tensorized Transformer

This section develops Tensorized Transformer attention by first expressing self-attention through Tucker decomposition, then extending it to shared multi-block representations using Block-Term decomposition. The resulting Multi-linear attention integrates into Transformer while reducing attention parameters and changing the computational complexity profile.

3.1 Single-block Attention by Tucker Decomposition

  • Single-block attention represents Transformer self-attention as a linear combination of basis vectors using Tucker decomposition.The construction uses Q, K, and V as factor matrices and initializes a trainable core tensor G.Source: §3 · p. 3, §3.1 · p. 4

3.2 Multi-linear Attention

  • Multi-linear attention uses Block-Term tensor decomposition and shares factor matrices across multiple blocks to compress the multi-head mechanism.Each block uses a diagonal core tensor, while SplitConcat concatenates matrices after tensor splitting.Source: §3 · p. 4, §3.2 · p. 5

3.1 Single-block Attention by Tucker Decomposition

  • The Single-block attention output is a 3-order tensor that can reconstruct scaled dot-product attention by summing over its second index when N equals d.The paper instead uses tensor splitting followed by concatenation to obtain a new representation.Source: §3.1 · p. 5

3.2 Multi-linear Attention

  • The compressed attention can be incorporated into Transformer and trained end-to-end through tensor averaging, splitting, and concatenation.These operations produce the input to the next Transformer layer and preserve integration with encoder and decoder frameworks.Source: §3.2 · p. 5

3.3 Compression

  • 8 times fewer parameters can be achieved in the attention layer when h=8 and d=512.The stated ratio is based on the attention-layer compression formula; whole-Transformer compression is lower because other layers remain.Source: §3.3 · p. 6

3.4 Complexity

  • Multi-linear attention has complexity O(N^3), compared with O(N^2d) for Transformer self-attention.The paper states that its minimum number of sequential operations for different layers is approximately equal to Transformer self-attention.Source: §3.3 · p. 6

4 Related Work

Related work frames Transformer-based language models as effective but resource-intensive, motivating compression methods that reduce parameters and computation. This paper focuses on combining parameter sharing with low-rank approximation to compress Transformer multi-head attention.

Transformer Language Models

  • Transformer and its variants achieve strong language-modeling results but use many parameters, complicating training and deployment with limited resources.Transformer also has a fixed-length context in language modeling, while Transformer-XL addresses this through segment-level recurrence and positional encoding.Source: §4 · p. 6

Model Compression

  • Existing compression approaches include parameter pruning and sharing, low-rank approximation, knowledge transfer, transferred convolutional filters, and tensor decomposition.Tensor decomposition has been used to create different neural-network language-model structures.Source: §4 · p. 6

This Work

  • The paper combines parameter sharing with low-rank approximation to reduce parameters and computation complexity in Transformer multi-head attention.Its stated focus is specifically compression of the multi-head attention mechanism.Source: §4 · p. 6

5 Experiments

Experiments replace Transformer’s multi-head attention with Multi-linear attention and evaluate the resulting Tensorized Transformer on language modeling and translation. Across the reported tasks, the model reduces parameters while maintaining or improving benchmark performance.

Experimental Setup

  • The experiments replace standard multi-head attention with Multi-linear attention in Transformer-based language modeling and translation models.Evaluation covers PTB, WikiText-103, One-Billion Word, and WMT-2016 English-German.Source: §5 · p. 7, §5.1 · p. 7, §5.3 · p. 8

Language Modeling

  • On PTB, Tensorized Transformer achieves comparable SoTA results with two core tensors while nearly halving parameters relative to Transformer and Transformer-XL.Only the multi-head attention component is replaced; other model structures remain unchanged.Source: §5.2 · p. 7
  • On WikiText-103, Tensorized Transformer obtains 18.9 PPL versus the previous SoTA 18.3, demonstrating comparable language-modeling performance.WikiText-103 tests long-term dependency modeling.Source: §5.2 · p. 7
  • The One-Billion comparison reports a single-model SoTA improvement from 21.8 to 19.5 PPL.The benchmark primarily tests short-term dependency modeling because its sentences are shuffled.Source: §5.2 · p. 8
  • Across the three language-modeling datasets, the model reports lower PPL than other models with substantially fewer parameters.The comparison includes Transformer-XL with tensor-train embedding compression and Sparse Transformer.Source: §5.2 · p. 8

Neural Machine Translation

  • On WMT-2016 English-German, core-1 and core-2 achieve BLEU scores of 34.10 and 34.91, compared with 34.5 for the basic Transformer, using fewer parameters.The encoder attention layers are replaced with Multi-linear attention and decoding uses beam size 5 with length penalty α=0.6.Source: §5.3 · p. 8

Discussion

  • The discussion attributes possible gains to concatenating split tensor matrices to capture more information than summation and to reduced-parameter overfitting alleviation.The authors report overfitting when the number of core tensors exceeds two and with relatively large word-embedding dimensions.Source: §5.4 · p. 9

6 Conclusion and Further Work

The paper presents Multi-linear attention as a Tensorized Transformer encoder layer that compresses multi-head attention through Block-Term tensor decomposition, low-rank approximation, and parameter sharing. The authors report higher compression ratios and better experimental results, while identifying overfitting with many cores as future work.

  • The Tensorized Transformer combines 3-order tensors, low-rank approximation, and parameter sharing through Block-Term tensor decomposition.Source: §6 · p. 9
  • Compared with existing Transformer-based methods, the model achieves higher compression ratios and better experimental results, particularly in language modeling.Source: §6 · p. 9
  • The authors plan to optimize the framework and apply it to other NLP tasks.Source: §6 · p. 9

B Theorem 3.1

The theorem develops a basis-vector representation of self-attention and connects it to a tensorized Single-block attention representation. Under specified conditions, summing tensor slices reconstructs Scaled Dot-Product attention.

Theorem 3.1

  • Self-attention outputs can be represented as linear combinations of orthonormal basis vectors when Q, K, and V are represented in that basis.Source: B Theorem 3.1 · p. 12
  • The attention output is represented as Attention(Q, K, V) = (e1, . . . , en)M, where M is an n×d coefficient matrix.Source: B Theorem 3.1 · p. 12
  • Q, K, and V are expressed using the shared basis vectors with coefficient matrices α, β, and ξ.Source: B Theorem 3.1 · p. 12
  • The input to softmax becomes a product of the coefficient matrices for Q and K.Source: B Theorem 3.1 · p. 12

Corollary 1

  • Under the theorem's conditions and N = d, summing Single-block attention tensor slices over the second index reconstructs Scaled Dot-Product attention.Source: C Corollary 1 · pp. 12–13
  • Figure 3 depicts the 3-order tensor output and the summation of tensor slices formed by splitting along index j.Source: Figure 3 · p. 13

D Compression Ratio about Multi-Linear Attention

The compression analysis compares Multi-linear attention with multi-head attention by combining parameter sharing and low-rank approximation. With h = 8 and d = 512, the reported compression ratio approximates 8.

E Experiment

This section introduces the Transformer integration structure for Multi-linear attention. Figure 5 provides information about how the proposed attention is incorporated into Transformer.

  • The paper proposes incorporating Multi-linear attention into Transformer.Source: §E.1 · p. 14
  • Figure 5 presents structural information about the integration process.Source: §E.1 · p. 14
  • The figure is intended to show how the proposed attention fits within the Transformer framework.Source: §E.1 · p. 14

E.2 Experimental Details in Language Modeling

The language-modeling experiments replace multi-head attention with BTD-based Multi-linear attention while retaining the other Transformer layers. The setup uses shared projections and varies training and optimization details across the evaluated datasets.

Model configuration

  • Multi-linear attention uses BTD to construct a 3-order tensor representation from linearly projected queries, keys, and values.The projected outputs are shared h times, where h is the number of core tensors.Source: §E.2 · p. 14
  • Figure 5 depicts incorporating Multi-linear attention into a partial Transformer structure with shared parameters across single-block attentions.Source: Figure 5 · p. 15

E.3 Experiment Details in Neural Machine Translation

The neural machine translation experiment replaces encoder multi-head attention with Multi-linear attention while keeping decoder multi-head attention and using the Transformer setup.

Translation setup

  • The translation model replaces multi-head attention only in the encoder and retains multi-head attention in the decoder.This setup tests the effectiveness of encoding a sentence with Multi-linear attention.Source: §E.3 · p. 15
  • The experiment uses the same setup as Transformer and trains on one NVIDIA P40 GPU.Source: §E.3 · p. 15

E.4 Experimental comparison

The comparison evaluates perplexity and FLOPs on PTB and WikiText-103, including matched-parameter and matched-hyperparameter settings. Tensorized Transformer achieves better results than Transformer-XL while using fewer parameters or lower computational cost.

Experimental comparison

  • The PTB and WikiText-103 comparisons report perplexity and FLOPs for Tensorized Transformer and baseline models.Table 5 presents experimental comparisons on the two datasets.Source: §E.4 · p. 15, Table 5 · p. 16
  • FLOPs measure the number of floating-point operations.Source: §E.4 · p. 15

Matched-size comparison

  • At matched parameter sizes, Tensorized Transformer achieves better results and lower FLOPs than Transformer-XL.The comparison uses 12M parameters on PTB and 85.5M on WikiText-103 for Tensorized Transformer.Source: §E.4 · pp. 15–16
  • Tensorized Transformer can be enlarged by increasing layers and sequence length to match Transformer-XL parameter counts.The reported configurations include 24M on PTB and 257M on WikiText-103.Source: §E.4 · p. 16

Matched-configuration comparison

  • With the same hyperparameters, Tensorized Transformer performs better than Transformer-XL on PTB and WikiText-103.The experiments also vary dmodel, with larger dimensions tested on both datasets.Source: §E.4 · p. 16, Table 6 · p. 16

Implementation

  • Single-block attention constructs a core tensor and combines it with queries, keys, and values using an einsum operation.Source: F Partial Code · pp. 16–17
  • The implementation defines SingleBlockAttention as a component of Multi-linear attention and applies its output through projection, dropout, and residual layer normalization.For multiple cores, the implementation computes two attention outputs and averages them.Source: F Partial Code · pp. 16–17
  • The Multi-linear attention implementation exposes core count, tensor rank, projection dimensions, and dropout as configurable parameters.Source: F Partial Code · p. 17

Related papers

Loading 1906.09777v3…