上下文就绪的Transformer

近日，一篇题为“上下文就绪Transformer”的论文被提交至arXiv，并提出了一种创新的循环神经网络架构。该架构由D层Transformer模块和一个校正网络组成。与传统Transformer不同，该架构在每个token进入模块之前，通过校正网络将之前位置的模块输出（即缓存的过去上下文摘要）与当前token嵌入结合，从而实现预上下文化。这使得token在进入模块时不再是原始嵌入，而是已经携带了上下文信息。这种设计灵感来源于递归神经网络，但通过并行训练解决了递归难以并行化的问题。

在训练过程中，研究者将校正过程在完整序列上展开K次，每一步并行处理所有位置，使得训练高效且可并行化。此外，任何预训练的Transformer模型都可以通过添加一个零初始化的校正前馈网络并进行微调，转换为上下文就绪模型。这意味着现有的预训练模型可以轻松升级，而无需从头训练。

实验表明，该架构在各种配置下均表现出色。在A100 GPU上，一个D=5的模型击败了12层的标准Transformer，生成速度提升1.7倍。更令人印象深刻的是，当K=10时，单层模型（D=1）以2.6倍的推理速度优势超越了6层Transformer，且序贯推理的困惑度与并行K=10的版本仅相差0.01。这表明该架构不仅在性能上有所提升，还在推理效率上实现了显著突破。

该架构特别受益于宽表示和长上下文。在指针追踪任务中，使用BPTT训练的D=1模型成功解决了所有10个组合层次，而标准Transformer则表现出阶梯式的深度依赖，即需要更多层次才能处理更深的组合。这一突破表明，通过预上下文化和校正机制，较浅的网络也能达到甚至超越深层网络的性能，为构建更高效的Transformer模型开辟了新方向。

相关论文已被NeurIPS接收，共22页，作者为Mahesh Godavarti。