AltUp: Efficient Transformers
Posted by Xin Wang, Software Engineer, and Nishanth Dikkala, Research Scientist, Google Research
Contemporary deep learning models have achieved great success in various domains, such as natural language and computer vision. Transformer neural networks (transformers) have become a popular deep learning architecture and are widely used in natural language processing. They are also being applied to other domains like computer vision, robotics, and autonomous driving. Additionally, transformers serve as the foundation for state-of-the-art language models.
However, scaling up transformer networks often leads to increased compute cost and inference latency. This raises the question of whether we can benefit from larger models without incurring such burdensome costs. In our research paper, “Alternating Updates for Efficient Transformers”, accepted as a Spotlight at NeurIPS 2023, we introduce AltUp, a method that allows us to take advantage of increased token representation without increasing computation cost. AltUp is easy to implement, applicable to any transformer architecture, and requires minimal hyperparameter tuning. For example, by using a variant of AltUp on a 770M parameter T5-Large model, we can significantly improve model quality with just ~100 additional parameters.
Background
To understand how AltUp achieves this, let’s delve into how transformers work. Transformers divide the input into a sequence of tokens, and each token is mapped to an embedding vector called the token embedding. The dimension of this vector is known as the token representation dimension. The transformer then operates on these token embeddings using computation modules called layers, which are determined by the token representation dimension. Prior works have focused on efficiently scaling up the network parameters, but there is a research gap when it comes to scaling up the token representation dimension itself.
Method
AltUp addresses this gap by partitioning the widened representation vector into blocks and processing only one block at each layer. It utilizes a prediction-correction mechanism to infer the outputs of the other blocks. This allows AltUp to keep the model dimension and computation cost roughly constant while taking advantage of an increased token dimension. AltUp achieves this by activating different blocks of the representation vector at each layer. The activated block is passed through a transformer layer to compute its output, while a lightweight predictor computes a weighted combination of all input blocks. The predicted values, along with the computed value of the activated block, are used by a lightweight corrector to update the predictions. AltUp can be applied to an arbitrary number of blocks and can be combined with other techniques like MoE for additional performance gains.
Evaluation
We evaluated AltUp on T5 models for various benchmark language tasks. Models augmented with AltUp were consistently faster than extrapolated dense models at the same accuracy level. For example, a T5 Large model with AltUp achieved a 27% speedup on the GLUE benchmark, a 39% speedup on the SuperGLUE benchmark, an 87% speedup on the SQuAD benchmark, and a 29% speedup on the Trivia-QA benchmark. AltUp consistently led to models with better predictive performance compared to baseline models, while maintaining the same speed across different model sizes and benchmarks.
Extensions: Recycled AltUp
In scenarios where the vocabulary size is very large, AltUp may introduce additional computation and parameter costs. To address this, we propose an extension called Recycled-AltUp. Instead of widening the initial token embeddings, Recycled-AltUp replicates the embeddings multiple times to form a wider token representation. This allows us to benefit from a wider token dimension without adding additional parameters relative to the baseline transformer model.
