[论文解读] TensorCoder: Dimension-Wise Attention via Tensor Representation for Natural Language Modeling

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

Topic Modeling参考文献 26被引用 4

一句话总结

TensorCoder 提出了一种逐维注意力机制，通过张量运算重新定义查询-键交互，将自注意力复杂度从 O(N²d) 降低至 O(Nd²)，在 PTB 掩码语言建模任务上实现 13% 的性能提升，注意力层 FLOPs 减少 3.5 倍，优于标准 Transformer，同时在机器翻译任务中保持竞争力，计算量显著降低。

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.

研究动机与目标

• 解决标准 Transformer 自注意力机制的二次方计算复杂度问题，该问题限制了其在长序列上的可扩展性。
• 设计一种逐维注意力机制，在保持逐 token 注意力输出效用的同时降低复杂度。
• 通过张量运算和卷积操作，将逐维注意力集成到 Transformer 架构中以实现特征提取。
• 通过掩码的逐维注意力机制在解码器中保持自回归性质。
• 通过减少 FLOPs 和内存使用，实现预训练模型在低资源环境中的高效部署。

提出的方法

• 将标准的查询-键点积注意力（QKᵀ）替换为使用 QᵀK 的逐维注意力机制，将复杂度从 O(N²d) 降低至 O(Nd²)。
• 利用 Khatri-Rao（KR）积将注意力表示为三阶张量，以支持隐式和显式输出表示。
• 在张量表示上应用卷积操作，以提取并整合多维交互特征。
• 构建一种掩码的逐维注意力机制，限制信息流仅来自先前位置，从而保持自回归生成特性。
• 将逐维注意力机制集成到 Transformer 的编码器和解码器层中，替代标准的多头注意力机制。
• 通过基于张量的计算与表示，统一编码器与解码器结构中注意力的设计与实现。

实验结果

研究问题

• RQ1 能否通过逐维注意力机制在保持其表征能力的同时降低自注意力的二次方复杂度？
• RQ2 如何设计逐维注意力机制，使其输出形式与现有 Transformer 架构兼容？
• RQ3 基于张量的表示与操作能否有效捕捉序列建模中的跨维度交互？
• RQ4 所提出的掩码逐维注意力机制是否能在自回归解码中保持自回归生成特性？
• RQ5 逐维注意力的降低复杂度在长序列和低资源自然语言处理任务中，能在多大程度上提升效率与性能？

主要发现

• 在 PTB 掩码语言建模任务中，TensorCoder 相较于标准 Transformer 实现了 13% 的性能提升，推理时间缩短至三分之一。
• 在 WMT16 En-De 翻译任务中，TensorCoder 达到了与标准模型相当的 BLEU 分数（33.1 vs. 33.5），但注意力层 FLOPs 减少 3.5 倍。
• 在 IWSLT16 任务中，TensorCoder 将 FLOPs 降低至标准 Transformer 的三分之一，展现出显著的效率优势。
• 该逐维注意力机制将计算复杂度从 O(N²d) 降低至 O(Nd²)，实现了序列长度的线性扩展。
• 该模型在保持各任务竞争力的同时，显著降低了内存与计算需求，适用于边缘设备部署。
• Khatri-Rao 积与张量卷积的结合，有效实现了注意力机制内的特征提取与融合，支持高性能与低成本的平衡。