Reproduce “Jump to Conclusions: Short-Cutting Transformers With Linear Transformations” on OpenResearch

About this paper

Alexander Yom Din, Taelin Karidi, Leshem Choshen, Mor Geva

Researchers from Hebrew University of Jerusalem and Tel Aviv University introduce mat, a method that uses learned linear transformations to accurately approximate future hidden representations of Transformer layers. This approach enhances model interpretability and enables more efficient early exiting, saving approximately 7.9% more layers for GPT-2 and 5.4% more for BERT during inference compared to the common identity mapping.

Key insights

• Simple linear transformations (mat) are significantly more effective than identity mapping at approximating future hidden representations and their associated predictions across Transformer layers.
• A notable portion of a Transformer's final prediction information is already present in its early layers, and mat can accurately extract this.
• Individual Transformer sub-modules, especially the attention mechanism, can be effectively approximated using linear mappings, potentially enabling fine-grained architectural optimizations.

Summary by alphaXiv.