[论文解读] Controlling Computation versus Quality for Neural Sequence Models

Ankur Bapna, Naveen Arivazhagan+1 more|arXiv (Cornell University)|Feb 17, 2020

Topic Modeling参考文献 33被引用 13

一句话总结

本文提出条件计算Transformer（CCT），一种通过可学习控制网络在推理时动态控制Transformer模型计算量的方法。通过多任务学习在不同计算预算下进行训练，CCT在完整预算下达到具有竞争力的性能，且在计算受限的预算下显著优于基线模型，实现了在输入复杂度和资源限制下高效、自适应的推理。

ABSTRACT

Most neural networks utilize the same amount of compute for every example independent of the inherent complexity of the input. Further, methods that adapt the amount of computation to the example focus on finding a fixed inference-time computational graph per example, ignoring any external computational budgets or varying inference time limitations. In this work, we utilize conditional computation to make neural sequence models (Transformer) more efficient and computation-aware during inference. We first modify the Transformer architecture, making each set of operations conditionally executable depending on the output of a learned control network. We then train this model in a multi-task setting, where each task corresponds to a particular computation budget. This allows us to train a single model that can be controlled to operate on different points of the computation-quality trade-off curve, depending on the available computation budget at inference time. We evaluate our approach on two tasks: (i) WMT English-French Translation and (ii) Unsupervised representation learning (BERT). Our experiments demonstrate that the proposed Conditional Computation Transformer (CCT) is competitive with vanilla Transformers when allowed to utilize its full computational budget, while improving significantly over computationally equivalent baselines when operating on smaller computational budgets.

研究动机与目标

• 为解决标准Transformer的僵化性，其对每个输入均使用固定的计算量，而不论输入复杂度或可用资源如何。
• 实现推理过程中对每个样本的计算量进行动态、实时调整，基于输入难度和可用计算预算。
• 训练单一模型，使其能在多个计算-质量权衡下高效运行，而无需为每个预算重新训练。
• 开发一种可微训练方案，弥合训练期间的软门控与推理时的离散决策之间的差距。
• 将条件计算扩展至解码器层之外，包括编码器和注意力子网络，提升适用范围。

提出的方法

• 为每个Transformer层引入控制网络，在训练期间输出0到1之间的连续值，使用Sigmoid激活的门控函数。
• 在训练期间使用随机软门控：每层的输出是完整计算和残差的加权和，由门控值控制。
• 在训练期间向门控网络的预激活值添加零均值高斯噪声，以鼓励离散行为并提升泛化能力。
• 在推理时，门控值以0.5为阈值进行二值化：若门控值 ≥ 0.5，则执行子网络；否则跳过。
• 使用多任务目标端到端训练整个模型，每个任务对应一个不同的目标计算预算。
• 将该框架应用于编码器和解码器子网络，包括自注意力中的键-值投影，实现对计算的细粒度控制。

实验结果

研究问题

• RQ1 我们能否训练一个单一的Transformer模型，使其能够根据输入复杂度和可用预算动态调整每输入的计算量？
• RQ2 通过可学习控制网络实现的条件计算，在不同计算预算下如何影响模型质量与效率？
• RQ3 我们能否使用可微软门控端到端训练此类模型，同时仍支持推理时的离散决策？
• RQ4 该模型是否学会在不同难度的输入上高效分配计算？
• RQ5 所提方法能否扩展至解码器层之外的其他组件，如注意力投影或编码器层？

主要发现

• 条件计算Transformer（CCT）在获得其完整计算预算时，性能与标准Transformer相当。
• 在较小计算预算下，CCT在WMT英语-法语翻译任务和BERT风格的无监督表征学习任务中，均显著优于计算量相当的基线模型。
• 控制网络输出显示，模型学会在更难的输入上激活更复杂的子网络，体现出智能的计算分配能力。
• 模型有效学会在较简单的样本上跳过不必要的计算，显著降低推理成本，且质量损失可忽略。
• 多任务训练目标使单一模型能泛化至多个计算预算，支持按需推理。
• 可视化结果表明，键-值子网络仅在对应翻译中结构复杂部分的标记上被选择性激活，表明计算控制具有实际意义。