# ElastiFormer: Learned Redundancy Reduction in Transformer via Self-Distillation

arXiv: https://arxiv.org/abs/2411.15281 · HTML: https://mlsanity.com/paper/2411.15281/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 35 of its 86 references that are in the map.
- First posted on arXiv 22 Nov 2024.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)
- Knowledge distillation / [General knowledge distillation methods](https://mlsanity.com/distill/general-knowledge-distillation-methods/)

## Builds on (most cited first, 30 of 35)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [TinyBERT: Distilling BERT for Natural Language Understanding](https://mlsanity.com/paper/1909.10351/) (2019)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models](https://mlsanity.com/paper/2401.06066/) (2024)
- [Patient Knowledge Distillation for BERT Model Compression](https://mlsanity.com/paper/1908.09355/) (2019)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models](https://mlsanity.com/paper/2308.13137/) (2023)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Structured Pruning Learns Compact and Accurate Models](https://mlsanity.com/paper/2204.00408/) (2022)
- [MoEfication: Transformer Feed-forward Layers are Mixtures of Experts](https://mlsanity.com/paper/2110.01786/) (2021)
- [FQ-ViT: Post-Training Quantization for Fully Quantized Vision Transformer](https://mlsanity.com/paper/2111.13824/) (2021)
- [LLaMA-MoE: Building Mixture-of-Experts from LLaMA with Continual Pre-training](https://mlsanity.com/paper/2406.16554/) (2024)
- [A Practical Review of Mechanistic Interpretability for Transformer-Based Language Models](https://mlsanity.com/paper/2407.02646/) (2024)
- [Circuit Component Reuse Across Tasks in Transformer Language Models](https://mlsanity.com/paper/2310.08744/) (2023)
- [A Mechanistic Interpretation of Arithmetic Reasoning in Language Models using Causal Mediation Analysis](https://mlsanity.com/paper/2305.15054/) (2023)
- [What Matters in Transformers? Not All Attention is Needed](https://mlsanity.com/paper/2406.15786/) (2024)
- [Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods](https://mlsanity.com/paper/2402.02834/) (2024)
- [Unified Visual Transformer Compression](https://mlsanity.com/paper/2203.08243/) (2022)
- [Successor Heads: Recurring, Interpretable Attention Heads In The Wild](https://mlsanity.com/paper/2312.09230/) (2023)
- [Copy Suppression: Comprehensively Understanding an Attention Head](https://mlsanity.com/paper/2310.04625/) (2023)
- [ThinK: Thinner Key Cache by Query-Driven Pruning](https://mlsanity.com/paper/2407.21018/) (2024)
- [MoEBERT: from BERT to Mixture-of-Experts via Importance-Guided Adaptation](https://mlsanity.com/paper/2204.07675/) (2022)

## Built on by (most cited first, 1 of 1)

- [UNIFERENCE: A Discrete Event Simulation Framework for Developing Distributed AI Models](https://mlsanity.com/paper/2603.26469/) (2026)

## Abstract (arXiv)

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.