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.

  • Method

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

  • Results

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

  • Takeaways & Limitations

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

  • Takeaways & Limitations

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

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.
  • The compressed attention can be integrated into Transformer encoder-decoder layers through tensor splitting, concatenation, and end-to-end training.
  • The paper proves scaled dot-product attention can be linearly represented using orthonormal basis vectors.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • Block-Term Tensor Decomposition: 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.
  • Multi-head Attention: Transformer processes queries, keys, and values as matrices Q, K, and V in scaled dot-product attention.
  • Multi-head Attention: Multi-head attention concatenates multiple attention heads and applies an output projection.The formulation is MultiHeadAttention(Q, K, V ) = Concat(head1, . . . , headk)W O.
  • Multi-head Attention: Multiple parameter groups in multi-head attention produce a large number of redundant parameters.

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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.

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.
  • 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.
  • Language Modeling: 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.
  • Language Modeling: 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.
  • Language Modeling: 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.
  • 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.
  • 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.

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.
  • Compared with existing Transformer-based methods, the model achieves higher compression ratios and better experimental results, particularly in language modeling.
  • The authors plan to optimize the framework and apply it to other NLP tasks.

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.
  • Theorem 3.1: The attention output is represented as Attention(Q, K, V) = (e1, . . . , en)M, where M is an n×d coefficient matrix.
  • Theorem 3.1: Q, K, and V are expressed using the shared basis vectors with coefficient matrices α, β, and ξ.
  • Theorem 3.1: The input to softmax becomes a product of the coefficient matrices for Q and K.
  • 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.
  • Corollary 1: Figure 3 depicts the 3-order tensor output and the summation of tensor slices formed by splitting along index j.

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.

  • Multi-head attention uses distinct weight matrices, whereas Multi-linear attention uses three weight matrices and h weight vectors.
  • Figure 4 compares the parameter structures of Multi-linear attention and multi-head attention.
  • The compression ratio is computed from the parameter counts of the two attention mechanisms.
  • 8 is the approximate compression ratio when h = 8 and d = 512.
  • The method combines parameter sharing with low-rank approximation, with rank R smaller than the attention-matrix dimension d.
  • The compression ratio increases when the core-tensor rank R is smaller, and experiments set R to 18, below dmodel.

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.
  • Figure 5 presents structural information about the integration process.
  • The figure is intended to show how the proposed attention fits within the Transformer framework.

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.
  • Model configuration: Figure 5 depicts incorporating Multi-linear attention into a partial Transformer structure with shared parameters across single-block attentions.

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.
  • Translation setup: The experiment uses the same setup as Transformer and trains on one NVIDIA P40 GPU.

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.
  • Experimental comparison: FLOPs measure the number of floating-point operations.
  • 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.
  • Matched-size comparison: 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.
  • 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.
  • Implementation: Single-block attention constructs a core tensor and combines it with queries, keys, and values using an einsum operation.
  • Implementation: 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.
  • Implementation: The Multi-linear attention implementation exposes core count, tensor rank, projection dimensions, and dropout as configurable parameters.
Loading 1906.09777v3…