[论文解读] Language Modeling using LMUs: 10x Better Data Efficiency or Improved Scaling Compared to Transformers

Narsimha Chilkuri, Eric Hunsberger+3 more|arXiv (Cornell University)|Oct 5, 2021

Topic Modeling参考文献 14被引用 5

一句话总结

该论文提出了一种基于勒让德记忆单元（LMU）的语言模型，其数据效率比Transformer高出10倍，仅需1/10的训练标记数即可达到相同的损失。通过利用Legendre多项式投影实现O(n)内存和O(n ln n)计算复杂度，该模型在扩展性方面优于Transformer和LSTM，尤其在结合全局自注意力机制时表现更优。

ABSTRACT

Recent studies have demonstrated that the performance of transformers on the task of language modeling obeys a power-law relationship with model size over six orders of magnitude. While transformers exhibit impressive scaling, their performance hinges on processing large amounts of data, and their computational and memory requirements grow quadratically with sequence length. Motivated by these considerations, we construct a Legendre Memory Unit based model that introduces a general prior for sequence processing and exhibits an $O(n)$ and $O(n \ln n)$ (or better) dependency for memory and computation respectively. Over three orders of magnitude, we show that our new architecture attains the same accuracy as transformers with 10x fewer tokens. We also show that for the same amount of training our model improves the loss over transformers about as much as transformers improve over LSTMs. Additionally, we demonstrate that adding global self-attention complements our architecture and the augmented model improves performance even further.

研究动机与目标

• 为解决Transformer在语言建模中对高数据量和高计算资源的需求，特别是其随序列长度增长的O(n²)内存和计算复杂度。
• 探究非参数化、内存高效的架构（如勒让德记忆单元，LMU）是否能在扩展性与数据效率方面优于Transformer。
• 评估LMU状态中的隐式自注意力机制是否能与Transformer中的显式自注意力机制相媲美或超越其性能。
• 研究在LMU中引入全局自注意力是否能进一步提升性能与长上下文建模能力。
• 验证LMU模型在整个模型规模范围内是否表现出幂律扩展行为，并将其扩展指数与Transformer和LSTM进行比较。

提出的方法

• 该模型采用勒让德记忆单元（LMU），一种非参数化的线性时不变（LTI）系统，将输入标记的滑动窗口投影到q个勒让德多项式上，实现高效的时序表征。
• LMU状态通过固定矩阵递推更新：$\mathbf{m}_t = \bar{\mathbf{A}} \mathbf{m}_{t-1} + \bar{\mathbf{B}} x_t$，其中$\bar{\mathbf{A}}$和$\bar{\mathbf{B}}$为从LTI系统导出的冻结矩阵。
• LTI更新通过卷积并行计算：$\mathbf{m}_t = \sum_{j=1}^{t} \bar{\mathbf{A}}^{t-j} \bar{\mathbf{B}} x_j$，实现O(n)内存和O(n ln n)计算复杂度。
• 通过在每个时间步对LMU隐藏状态应用全局注意力机制，实现“隐式自注意力”，而非在时间步之间进行。
• 引入一种结合全局自注意力的变体，以提升长上下文建模能力，融合LMU的效率与注意力机制的表征能力。
• 模型在WebText2数据集上进行训练，并通过拟合模型规模（非嵌入参数量）与交叉熵损失之间的幂律关系来评估扩展规律。

实验结果

研究问题

• RQ1 LMU模型是否在模型规模与语言建模损失之间表现出类似Transformer和LSTM的幂律扩展关系？
• RQ2 由于数据效率提升，LMU模型是否能在显著更少的训练数据下实现与Transformer相当或更优的性能？
• RQ3 LMU模型中的隐式自注意力机制是否能有效捕捉长距离依赖关系，而无需显式的跨时间步注意力？
• RQ4 与基础LMU和Transformer相比，向LMU架构中添加全局自注意力对性能和扩展行为有何影响？
• RQ5 LMU模型的O(n)内存和O(n ln n)计算复杂度在多大程度上使其扩展性优于Transformer的O(n²)复杂度？

主要发现

• LMU模型在仅使用1/10训练标记数的情况下，即可达到与Transformer相同的交叉熵损失，证明其数据效率高出10倍。
• 在相同训练预算下，LMU模型在损失上优于Transformer的幅度，与Transformer相比LSTM的提升幅度相当。
• LMU模型表现出与模型规模相关的幂律扩展关系：$\text{LMU}(N) = \left(\frac{N}{1.95 \cdot 10^{14}}\right)^{-0.072}$，其扩展指数接近Transformer。
• 加入全局自注意力的LMU模型扩展关系为$\text{LMU}_G(N) = \left(\frac{N}{3.80 \cdot 10^{14}}\right)^{-0.069}$，表明其扩展性能略优于基础LMU。
• 在LMU模型中加入全局注意力后，其性能仍持续优于在10倍更多标记上训练的Transformer，证实其具有更优的数据效率与性能。
• 每标记损失曲线显示，基础LMU模型在约100个标记处即趋于平缓，表明其长上下文建模能力有限；而全局注意力变体则能持续提升长上下文性能，与Transformer相似。