[Paper Review] TensorCoder: Dimension-Wise Attention via Tensor Representation for Natural Language Modeling

Shuai Zhang, Peng Zhang+4 more|arXiv (Cornell University)|Jul 28, 2020

Topic Modeling26 references4 citations

TL;DR

TensorCoder proposes a dimension-wise attention mechanism that reduces self-attention complexity from O(N²d) to O(Nd²) by redefining query-key interaction via tensor operations, achieving 13% better performance on PTB masked language modeling and 3.5× fewer FLOPs than standard Transformer, while maintaining competitive results on machine translation with reduced computation.

ABSTRACT

Transformer has been widely-used in many Natural Language Processing (NLP) tasks and the scaled dot-product attention between tokens is a core module of Transformer. This attention is a token-wise design and its complexity is quadratic to the length of sequence, limiting its application potential for long sequence tasks. In this paper, we propose a dimension-wise attention mechanism based on which a novel language modeling approach (namely TensorCoder) can be developed. The dimension-wise attention can reduce the attention complexity from the original $O(N^2d)$ to $O(Nd^2)$, where $N$ is the length of the sequence and $d$ is the dimensionality of head. We verify TensorCoder on two tasks including masked language modeling and neural machine translation. Compared with the original Transformer, TensorCoder not only greatly reduces the calculation of the original model but also obtains improved performance on masked language modeling task (in PTB dataset) and comparable performance on machine translation tasks.

Motivation & Objective

• Address the quadratic computational complexity of standard Transformer self-attention, which limits scalability to long sequences.
• Develop a dimension-wise attention mechanism that maintains the utility of token-wise attention output while reducing complexity.
• Integrate dimension-wise attention into the Transformer architecture using tensor operations and convolution for feature extraction.
• Preserve auto-regressive property in decoders via a masked dimension-wise attention mechanism.
• Enable efficient deployment of pre-trained models in low-resource environments through reduced FLOPs and memory usage.

Proposed method

• Replace standard query-key dot product attention (QKᵀ) with a dimension-wise attention mechanism using QᵀK, reducing complexity from O(N²d) to O(Nd²).
• Represent attention as a third-order tensor using the Khatri-Rao (KR) product to enable implicit and explicit output representations.
• Apply convolutional operations on the tensor representation to extract and integrate multi-dimensional interactive features.
• Construct a masked dimension-wise attention mechanism that restricts information flow to previous positions, preserving auto-regressive generation.
• Integrate the dimension-wise attention into both encoder and decoder layers of the Transformer, replacing standard multi-head attention.
• Use tensor-based computation and representation to unify the design and implementation of attention across encoder and decoder structures.

Experimental results

Research questions

• RQ1 Can a dimension-wise attention mechanism reduce the quadratic complexity of self-attention while preserving its representational power?
• RQ2 How can dimension-wise attention be designed to produce output forms compatible with existing Transformer architectures?
• RQ3 Can tensor-based representations and operations effectively capture cross-dimensional interactions in sequence modeling?
• RQ4 Does the proposed masked dimension-wise attention maintain auto-regressive generation properties in autoregressive decoding?
• RQ5 To what extent can the reduced complexity of dimension-wise attention improve efficiency and performance on long-sequence and low-resource NLP tasks?

Key findings

• On the PTB masked language modeling task, TensorCoder achieved a 13% performance improvement over the standard Transformer, with inference time reduced to one-third.
• In WMT16 En-De translation, TensorCoder achieved comparable BLEU scores (33.1 vs. 33.5) but with 3.5× fewer FLOPs in the attention layers.
• On IWSLT16, TensorCoder reduced FLOPs to one-third of the standard Transformer’s, demonstrating strong efficiency gains.
• The dimension-wise attention mechanism reduced computational complexity from O(N²d) to O(Nd²), enabling linear scaling with sequence length.
• The model maintained competitive performance across tasks while significantly reducing memory and computation demands, making it suitable for edge devices.
• The use of Khatri-Rao product and tensor convolution enabled effective feature extraction and integration within the attention mechanism, supporting high performance with lower cost.