ElastiFormer: Learned Redundancy Reduction in Transformer via Self-Distillation

TL;DR

ElastiFormer is a post-training method that dynamically reduces computation in pretrained Transformers by using tiny routing modules trained with self-distillation, applicable across various modalities and robust to domain shifts.

Contribution

It introduces a novel elastic Transformer framework with minimal additional parameters, enabling dynamic inference time adjustment through self-distillation training.

Findings

Achieves 20-50% compute savings in transformer components.

Can be combined with low-rank LoRA weights for further reduction.

Demonstrates robustness across different training domains.

Abstract

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.