ElastiFormer: Learned Redundancy Reduction in Transformer via Self-Distillation
This addresses the high computational cost of Transformer inference for users in NLP, vision, and multimodal domains, offering an incremental improvement through efficient adaptation.
The paper tackles the problem of reducing inference compute in pretrained Transformer models by introducing ElastiFormer, a post-training technique that uses self-distillation to dynamically select subsets of parameters and tokens, achieving 20% to 50% compute savings across various tasks.
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.