让Transformer模型按需计算，节省20%~50%算力
ElastiFormer: Learned Redundancy Reduction in Transformer via Self-Distillation
- 用轻量路由模块动态选择每层处理的参数和输入
- 实现20%至50%算力节省，且支持图像、文本等多模态
- 无需重训练，适配各类预训练模型，适合资源受限场景
我们提出ElastiFormer，一种后训练技术，可将预训练Transformer模型转化为可变推理计算量的弹性版本。ElastiFormer引入极小的路由模块（额外可训练参数低至0.00006%），以输入相关方式动态选择每层处理的网络参数和输入标记。路由模块通过自蒸馏损失训练，使弹性模型输出与原模型尽可能一致。该方法不依赖模型模态，可应用于因果语言建模、图像建模及视觉-语言建模任务。实验显示，不同Transformer组件可节省20%至50%计算量，进一步结合低秩LoRA权重（秩为1）可降低更多开销。通过在ImageNet不同子集上训练路由模块，验证了ElastiFormer对训练域的鲁棒性。
原文摘要 · Abstract (English)
We introduce ElastiFormer, a post-training technique that adapts pretrained Transformer models into an elastic counterpart with variable inference time compute. ElastiFormer introduces small routing modules (as low as .00006% additional trainable parameters) to dynamically selects subsets of network parameters and input tokens to be processed by each layer of the pretrained network in an inputdependent manner. The routing modules are trained using self-distillation losses to minimize the differences between the output of the pretrained-model and their elastic counterparts. As ElastiFormer makes no assumption regarding the modality of the pretrained Transformer model, it can be readily applied to all modalities covering causal language modeling, image modeling as well as visual-language modeling tasks. We show that 20% to 50% compute saving could be achieved for different components of the transformer layer, which could be further reduced by adding very low rank LoRA weights (rank 1) trained via the same distillation objective. Finally, by comparing routing trained on different subsets of ImageNet, we show that ElastiFormer is robust against the training domain.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。