论文地址

[1911.02150] Fast Transformer Decoding: One Write-Head is All You Need

论文总结

《Fast Transformer Decoding: One Write - Head is All You Need》总结

这篇论文由Noam Shazeer撰写，提出了一种名为多查询注意力（Multi - Query Attention）的Transformer变体，旨在提高增量推理的速度，同时减少质量下降。

  1. 引言

    • Transformer模型中，多头注意力层在训练时并行性好，但增量推理速度慢，受限于内存带宽。
    • 本文提出多查询注意力，减少键和值张量的大小，降低内存带宽要求，提高推理速度。

  2. 背景：神经注意力

    • 点积注意力：通过查询向量与键向量的点积计算权重，对值向量进行加权求和得到输出向量。
    • 多头注意力：并行使用多个注意力层，查询向量、键向量和值向量通过不同的线性投影得到。
    • 多头注意力（批量）：批量处理多个查询，提高效率，可防止信息回流。
    • 多头注意力（增量）：在自回归语言模型中，由于输出的自注意力层会影响下一个位置的生成，导致无法并行计算。

  3. 多查询注意力

    • 与多头注意力类似，但不同的头共享一组键和值。
    • 性能分析：在增量设置中，内存访问与算术运算的比例降低，理论上能大幅提高增量生成的性能。

  4. 实验与结果

    • 实验设置：在WMT 2014英德翻译任务和亿词语言建模基准上进行实验，与基线模型进行对比。
    • 模型质量：多查询注意力模型略逊于基线模型，但比减少头数或键值维度的模型要好。
    • 速度：训练和推理速度在TPUv2上评估，多查询模型的编码器和解码器速度均比基线模型快。

  5. 结论

    • 多查询注意力降低了内存带宽要求，使基于注意力的序列模型在推理性能关键的应用中更广泛地采用。