《TRANSFORMER - VQ: LINEAR - TIME TRANSFORMERS VIA VECTOR QUANTIZATION》 作者是Lucas D. Lingle，该论文介绍了一种名为Transformer - VQ的解码器Transformer，它能够以线性时间计算基于softmax的密集自注意力。

  1. 研究背景：

    • 标准Transformer的自注意力机制具有二次时间复杂度，限制了其在处理长序列时的应用。
    • 虽然已有多种高效Transformer被提出，但仍有改进空间。

  2. 相关工作：

    • 介绍了向量量化（VQ）的相关概念和定理，包括向量量化器、代码簿以及基于向量量化的表示学习。
    • 讨论了与Transformer - VQ相关的工作，如分层注意力、可核化注意力、压缩注意力、门控序列模型以及与VQ、k - means和其他相关的工作。

  3. Transformer - VQ：

    • 二次时间公式：定义了向量量化自注意力（VQAttn）的函数，引入了门控注意力单元（GAU）的设计，使用了LayerNorm和STVQ等技术。
    • 线性时间编码器注意力：通过定理证明，在没有因果掩码的情况下，注意力权重可以进行因式分解，从而简化计算。
    • 线性时间解码器注意力：给出了在有因果掩码和滑动窗口偏置的情况下，计算产品WV的块级递归公式，使得Transformer - VQ能够在线性时间内计算解码器注意力。
    • 学习算法：定义了Transformer - VQ的训练损失，包括交叉熵损失和承诺损失，并介绍了训练更新的策略。

  4. 实验：

    • 初步研究：包括代码簿大小的消融实验、压缩缓存的消融实验以及延迟和吞吐量的测量。结果表明，较大的代码簿尺寸可以提高预测质量，但会增加训练时间；压缩缓存对于提高模型质量至关重要；Transformer - VQ在延迟和吞吐量方面具有明显优势。
    • 定量结果：在Enwik8、PG - 19和ImageNet64等数据集上进行实验，结果表明Transformer - VQ在学习长序列依赖方面具有很强的能力，在多个数据集上取得了具有竞争力的结果。

  5. 结论：

    • Transformer - VQ是一种在线性时间内计算自注意力的Transformer解码器，通过向量量化的键和新颖的缓存机制实现了高效的注意力计算。
    • 大规模实验表明，Transformer - VQ是一种高效灵活的自回归模型，在PG - 19和ImageNet64等数据集上取得了接近或达到最先进水平的结果。
    • 未来的工作方向包括研究正式的缩放定律、构建更大的模型以及移植到更低级的框架中。

实验细节：

  • 数据集：Enwik8（字节级语言建模数据集）、PG - 19（开放词汇语言建模数据集）、ImageNet64（图像数据集）。
  • 模型参数：不同数据集使用不同的模型参数，如在Enwik8上训练的Transformer - VQ有190M参数，在PG - 19上训练的有1.3B参数。
  • 超参数：不同数据集有不同的超参数设置，如序列长度、回传长度、块长度、模型维度、键维度、值维度、代码簿数量、GAU数量、正弦波丢弃率、残差丢弃率、层丢弃率、权重衰减、优化器等。
  • 训练细节：使用TPU v3进行训练，权重和令牌嵌入的初始化遵循特定方法，使用预归一化的LayerNorm，激活函数为SiLU，使用AdamW或Adafactor优化器，学习率进行线性预热和余弦衰减，代码簿承诺系数和EMA率也有特定设置。

相关代码和链接：

  • 代码实现：使用Jax和Flax库实现，相关伪代码在附录E中提供。
  • 开源链接：代码开源在https://github.com/transformer - vq/transformer_vq。