@inproceedings{
panigrahi2023trainable,
title={Trainable Transformer in Transformer},
author={Abhishek Panigrahi and Sadhika Malladi and Mengzhou Xia and Sanjeev Arora},
booktitle={R0-FoMo:Robustness of Few-shot and Zero-shot Learning in Large Foundation Models},
year={2023},
url={https://openreview.net/forum?id=7mEOK0EnbY}
}
Recent works attribute the capability of in-context learning (ICL) in large pre-trained language models to implicitly simulating and fine-tuning an internal model (e.g., linear or 2-layer MLP) during inference. However, such constructions require large memory overhead, which makes simulation of more sophisticated internal models intractable. In this work, we propose a new efficient construction, Transformer in Transformer (in short, TINT), that allows a transformer to simulate and fine-tune more complex models during inference (e.g., pre-trained language models). In particular, we introduce innovative approximation techniques that allow a TINT model with less than 2 billion parameters to simulate and fine-tune a 125 million parameter transformer model within a single forward pass. TINT accommodates many common transformer variants and its design ideas also improve the efficiency of past instantiations of simple models inside transformers. We conduct end-to-end experiments to validate the internal fine-tuning procedure of TINT on various language modeling and downstream tasks. For example, even with a limited one-step budget, we observe TINT for a OPT-125M model improves performance by 4 − 16% absolute on average compared to OPT-125M. These findings suggest that large pre-trained language models are capable of performing intricate subroutines. To facilitate further work, a modular and extensible codebase for TINT will be open-sourced.
Trainable Transformer in Transformer
Abhishek Panigrahi Sadhika Malladi Mengzhou Xia Sanjeev Arora