[论文解读] Adaptive Computation with Elastic Input Sequence

Fuzhao Xue, Valerii Likhosherstov+4 more|arXiv (Cornell University)|Jan 30, 2023

Advanced Neural Network Applications被引用 5

一句话总结

AdaTape 提出了一种新颖的自适应计算框架，通过使用动态可学习的存储带（tape）为 Transformer 模型生成弹性输入序列，实现输入内容和序列长度的自适应。该方法在计算量匹配的条件下，在图像识别任务上实现了最先进性能，优于标准 Transformer 和自适应 Transformer，其核心在于引入了一种新型自适应存储带读取（ATR）机制，可从存储带库中选择信息量丰富的 token。

ABSTRACT

Humans have the ability to adapt the type of information they use, the procedure they employ, and the amount of time they spend when solving problems. However, most standard neural networks have a fixed function type and computation budget regardless of the sample's nature or difficulty. Adaptivity is a powerful paradigm as it not only imbues practitioners with flexibility pertaining to the downstream usage of these models but can also serve as a powerful inductive bias for solving certain challenging classes of problems. In this work, we introduce a new approach called AdaTape, which allows for dynamic computation in neural networks through adaptive tape tokens. AdaTape utilizes an elastic input sequence by equipping an architecture with a dynamic read-and-write tape. Specifically, we adaptively generate input sequences using tape tokens obtained from a tape bank which can be either trainable or derived from input data. We examine the challenges and requirements to obtain dynamic sequence content and length, and propose the Adaptive Tape Reading (ATR) algorithm to achieve both goals. Through extensive experiments on image recognition tasks, we show that AdaTape can achieve better performance while maintaining the computational cost. To facilitate further research, we have released code at https://github.com/google-research/scenic.

研究动机与目标

• 为解决标准 Transformer 固定计算预算的局限性，使其无法根据输入复杂度自适应调整。
• 探索一种新的自适应计算范式，通过动态调整输入序列长度和内容，而非模型深度或稀疏路由。
• 提供一种灵活的单模型解决方案，可在推理阶段无需微调即可动态扩展计算量。
• 证明自适应输入序列可作为解决高难度问题（如向量奇偶性任务）的强大归纳偏置。

提出的方法

• AdaTape 引入一个可学习或由输入生成的 token 存储带库，作为每个输入样本的动态记忆。
• 采用自适应存储带读取（ATR）算法，根据输入复杂度递归选择可变长度的存储带 token 序列。
• 将选定的存储带 token 追加到输入序列之前，再输入标准 Transformer 编码器。
• ATR 使用停止机制动态决定读取多少存储带 token，从而实现每个样本的可变计算量。
• 存储带库可为可训练参数，也可由输入特征（如 patch embeddings）生成，从而实现内容自适应。
• 该方法通过避免递归应用模型或引入参数稀疏性，保持计算效率，转而通过修改输入序列结构实现优化。

实验结果

研究问题

• RQ1 在不增加 FLOPs 的前提下，自适应输入序列长度和内容是否能提升模型性能？
• RQ2 所提出的 AdaTape 机制是否能使 Transformer 解决标准 Transformer 无法处理的任务（如向量奇偶性任务）？
• RQ3 在计算资源受限条件下，AdaTape 与现有自适应方法相比，在推理速度和准确率方面表现如何？
• RQ4 AdaTape 是否能在不同硬件和批量大小下，保持高性能的同时实现推理阶段的动态计算扩展？

主要发现

• 在 FLOPs 相当的前提下，AdaTape 在 ImageNet-10-shot 准确率上优于标准 ViT 和 DeiT 模型，当移除内容自适应性后，相比非自适应基线模型仍取得 2.3 个百分点的性能提升。
• 在具有挑战性的向量奇偶性任务中，AdaTape 能成功解决标准 Transformer 失败的问题，验证了其在高难度任务中的归纳偏置有效性。
• AdaTape-S/16 在推理速度上优于 Tiny 级别的基线模型，尽管其参数容量更高，表明动态序列长度带来了效率增益。
• 消融实验表明，移除内容自适应性会导致 ImageNet 10-shot 准确率显著下降 2.3 个百分点，证实其在性能中的关键作用。
• 更大的 AdaTape 模型显示出更高的序列长度方差，表明其在大规模下对计算成本具有更好的控制能力与更强的自适应性。
• 可视化结果表明，AdaTape 优先选择图像中心区域、信息量更丰富的 patch，与人类对显著性区域的直觉一致。