TensorCoder: Dimension-Wise Attention via Tensor Representation for Natural Language Modeling

1 Introduction

• Dimension-wise attention replaces QK^T with Q^T K, reducing attention complexity from O(N^2d) to O(Nd^2).The paper describes the resulting matrix as encoding global dependencies among dimensions.Source: §1 · pp. 1–2
• TensorCoder uses a Khatri-Rao tensor representation and convolution to extract explicit and implicit output features, and masks decoder attention to prevent leftward information flow.The mask ensures prediction at position i depends only on known outputs at positions less than i.Source: §1 · p. 2
• TensorCoder improves PTB performance by 13% with a time cost of 1/3 of multi-head token-wise attention, while WMT16 results are comparable with 1/3.5 the Transformer FLOPs.The paper also reports comparable performance on machine translation tasks.Source: §1 · p. 2

2.1 Tensor and Khatri-Rao Product

• A third-order tensor has three modes, and the Khatri-Rao product forms a matching-columnwise Kronecker product of two matrices.Its columns are pairwise Kronecker products of corresponding input columns.Source: §2.1 · p. 2

2.2 Attention

• Token-wise attention computes query-key dot products across tokens, applies softmax, and weights values; multi-head attention concatenates heads from different representation subspaces.The inputs Q, K, and V have shape N×d, with d the head dimensionality.Source: §2.2 · p. 3

3.1 Dimension-Wise Attention Matrix

• Dimension-wise attention computes S = Q^T K, relating embedding dimensions and using whole-sequence context rather than token-pair correlations.The paper relates S to a covariance matrix when each head's input is standardized and centralized.Source: §3.1 · pp. 3–4

3.2 KR Product Tensor Representation

• The Khatri-Rao product combines S and V into a third-order tensor, whose explicit and implicit representations provide alternative attention outputs.The paper describes explicit output as weighted coefficients under one basis vector and implicit output as weighted coefficients across all basis vectors.Source: §3.2 · p. 4

3.3 Masked Dimension-Wise Attention

• Masked dimension-wise attention uses an upper-triangular mask so each position's dimension-attention score depends only on known tokens at earlier positions.The masked tensor is combined with V to construct a third-order tensor.Source: §3.3 · pp. 5–6

3.4 Feature Extraction

• Sliding convolution filters the third-order tensor along sequence dimension N to produce output O for the next network layer.The filter is shared across slices, with different parameters for each column.Source: §3.4 · p. 6

3.5 Analysis of Complexity

• TensorCoder's encoder attention, KR tensor construction, and feature extraction each cost O(Nd^2), while decoder masked attention costs O(N^2d^2).The paper notes an encoder-only advantage over Transformer for tasks such as masked language modeling.Source: §3.5 · p. 6

• Multi-linear attention compresses parameters and improves performance, but cannot replace decoder attention because it lacks masking.The passage limits its application to the Transformer encoder.Source: §4 · p. 6
• Sparse Transformer reduces attention complexity from O(N^2) to O(N), while Reformer reduces it to O(N log N); both remain nonlinear in sequence length.The section characterizes both methods as achieving relatively good results.Source: §4 · p. 6
• Linear Transformer achieves O(N) time and memory complexity by discarding softmax and using the associative property of matrix multiplication.It uses feature representations for queries, keys, and values.Source: §4 · p. 6

5 Experiments

• The experiments replace Transformer multi-head attention with dimension-wise attention, comparing TensorCoder against Transformer and Tensorized Transformer, with matched attention-output counts in the single- and eight-output configurations.Source: §5 · p. 7

5.1 Evaluation on Masked Language Modeling

• For masked language modeling, 15% of tokens were selected for prediction, with 80% replaced by [MASK], 10% by a random word, and 10% kept unchanged; PTB sequences had length 100 and Wiki-103 sequences length 160.Models were evaluated by negative log-likelihood loss and FLOPs; the task used the encoder with other Transformer parts unchanged.Source: §5.1 · p. 7
• On PTB, both TensorCoder convolution methods had lower loss than Transformer and up to 3.05 times lower complexity; on Wiki-103, loss and FLOPs were much smaller than Tensorized Transformer.Source: §5.1 · pp. 7–8, Figure 3 · p. 8

5.2 Evaluation on Neural Machine Translation

• The translation evaluation used WMT16 and IWSLT16 English-to-German; FLOPs report encoder attention layers, with sequence length set to the training-set maximum.IWSLT16 used 197k training sentence pairs, and evaluation used beam size 5 and length penalty α=0.6.Source: §5.2 · p. 8
• In WMT16, TensorCoder-1conv and TensorCoder-8conv scored BLEU 30.7 and 33.1; Transformer used 3.5 times the FLOPs at similar performance, while IWSLT16 TensorCoder used one-third the FLOPs.Source: §5.2 · p. 8, Table 2 · p. 8, Table 3 · p. 8

• TensorCoder replaces token-wise attention with dimension-wise attention, reducing attention complexity to O(Nd^2) and making it linear rather than quadratic in sequence length.Source: §6 · p. 9

• Transformer token-wise attention is described as unaware of interactions between embedding dimensions and as quadratic in sequence length, limiting long-sequence and resource-limited use.Source: §7 · p. 9
• TensorCoder's linear sequence-length complexity is proposed as potentially useful for document-level NLP and resource-limited devices, including phones and tablets.Source: §7 · p. 9
• Future work will study how dimension-wise attention dimensionality relates to sequence length and how to set it for a given length.Source: §7 · p. 9