[论文解读] Random-LTD: Random and Layerwise Token Dropping Brings Efficient Training for Large-scale Transformers

Zhewei Yao, Xiaoxia Wu+5 more|arXiv (Cornell University)|Nov 17, 2022

Topic Modeling被引用 4

一句话总结

Random-LTD 提出了一种随机且逐层的标记丢弃方法，可在所有中间 Transformer 层中随机选择性地跳过计算，无需重要性评分或特殊标记处理。该方法在 GPT-3 1.3B 上实现了高达 33.3% 的理论计算量减少和 25.6% 的实际训练时间节省，同时保持了相近的零样本性能，且引入了一种新型的 LayerToken 学习率调度策略，以提升收敛性并减少调参负担。

ABSTRACT

Large-scale transformer models have become the de-facto architectures for various machine learning applications, e.g., CV and NLP. However, those large models also introduce prohibitive training costs. To mitigate this issue, we propose a novel random and layerwise token dropping method (random-LTD), which skips the computation of a subset of the input tokens at all middle layers. Particularly, random-LTD achieves considerable speedups and comparable accuracy as the standard training baseline. Compared to other token dropping methods, random-LTD does not require (1) any importance score-based metrics, (2) any special token treatment (e.g., [CLS]), and (3) many layers in full sequence length training except the first and the last layers. Besides, a new LayerToken learning rate schedule is proposed for pretraining problems that resolve the heavy tuning requirement for our proposed training mechanism. Finally, we demonstrate that random-LTD can be applied to broader applications, including GPT and BERT pretraining as well as ViT and GPT finetuning tasks. Our results show that random-LTD can save about 33.3% theoretical compute cost and 25.6% wall-clock training time while achieving similar zero-shot evaluations on GPT-31.3B as compared to baseline.

研究动机与目标

• 降低 NLP 和视觉应用中大规模 Transformer 模型训练的高昂成本。
• 消除在标记丢弃方法中对重要性评分或特殊标记处理（如 [CLS]）的依赖。
• 实现对所有中间层的标记丢弃，而不仅限于部分层，以最大化效率。
• 通过一种新型的 LayerToken 基础学习率调度策略，减少超参数调优负担。
• 在包括 BERT、GPT 和 ViT 预训练与微调在内的多样化任务中验证该方法。

提出的方法

• Random-LTD 在每个中间 Transformer 层中随机丢弃标记，不使用重要性评分，也不保留特殊标记。
• 每一层独立地执行标记丢弃，保留层间注意力依赖关系，避免因序列长度导致的瓶颈。
• 引入单调的序列长度增长调度策略，以减少梯度噪声，并缩小自回归模型中训练与推理之间的差距。
• 提出一种新的 LayerToken 学习率调度策略，根据每层累积的标记消耗量动态调整学习率，从而提升收敛性。
• 该方法在 BERT、GPT 和 ViT 架构中统一应用于预训练和微调。
• 该方法表现出正则化效果，与 Dropout 结合或替代 Dropout 时可提升性能。

实验结果

研究问题

• RQ1 随机且逐层的标记丢弃是否能在不依赖重要性评分或特殊标记处理的前提下，实现显著的训练效率提升？
• RQ2 与基线训练相比，Random-LTD 在计算成本、实际训练时间及模型准确率方面，在多样化架构中的表现如何？
• RQ3 基于 LayerToken 的学习率调度是否能有效减少调参工作量，并提升 Random-LTD 训练中的收敛性？
• RQ4 Random-LTD 是否表现出正则化效应，从而增强预训练和微调中模型的泛化能力？
• RQ5 Random-LTD 是否可泛化至不同 Transformer 架构和任务，包括因果语言模型和视觉 Transformer？

主要发现

• Random-LTD 在 GPT-3 1.3B 上实现了高达 33.3% 的理论计算量减少和最高 25.6% 的实际训练时间节省，同时保持了相近的零样本评估性能。
• 在使用 LayerToken 学习率调度时，该方法的验证损失与基线相当，且在 Random-LTD 训练中优于标准学习率调度。
• 在 BERT Large 预训练中，Random-LTD 在不使用 Dropout 的情况下，RACE-m 数据集上的准确率比标准基线高出 >1%，表明其具有正则化效应。
• 在 GPT-3 1.3B 上，使用 2880B LayerTokens 的 Random-LTD 可达到与使用 4320B LayerTokens 的基线相当的性能，证明了其高效的资源利用能力。
• LayerToken 学习率调度显著提升了收敛性，尤其是在训练后期，通过根据每层标记消耗量动态调整学习率实现。
• Random-LTD 在多样化任务中均表现出一致的性能提升，包括 BERT 和 GPT 的预训练以及 ViT 的微调，证实了其广泛的适用性。