Summary of Mixture-of-depths: Dynamically Allocating Compute in Transformer-based Language Models, by David Raposo et al.
Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
by David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro
First submitted to arxiv on: 2 Apr 2024
Categories
- Main: Machine Learning (cs.LG)
- Secondary: Computation and Language (cs.CL)
GrooveSquid.com Paper Summaries
GrooveSquid.com’s goal is to make artificial intelligence research accessible by summarizing AI papers in simpler terms. Each summary below covers the same AI paper, written at different levels of difficulty. The medium difficulty and low difficulty versions are original summaries written by GrooveSquid.com, while the high difficulty version is the paper’s original abstract. Feel free to learn from the version that suits you best!
| Summary difficulty | Written by | Summary |
|---|---|---|
| High | Paper authors | High Difficulty Summary Read the original abstract here |
| Medium | GrooveSquid.com (original content) | Medium Difficulty Summary Transformer-based language models typically distribute floating-point operations (FLOPs) uniformly across input sequences. This work shows that transformers can learn to dynamically allocate FLOPs to specific positions in a sequence, optimizing allocation for different layers throughout the model depth. The proposed method enforces a total compute budget by capping the number of tokens (k) participating in self-attention and MLP computations at each layer, using a top-k routing mechanism. By defining k beforehand, this approach uses a static computation graph with known tensor sizes, unlike other conditional computation techniques. However, since token identities are dynamic, this method can expend FLOPs non-uniformly across time and model depth dimensions. The resulting models efficiently learn to dynamically allocate compute, matching baseline performance for equivalent FLOPS and wall-clock times while requiring a fraction of the FLOPs per forward pass and being up to 50% faster during post-training sampling. |
| Low | GrooveSquid.com (original content) | Low Difficulty Summary This research paper is about improving how computers work with language. Normally, these computers spread their calculations evenly across what they’re working on. But this new method lets them focus on different parts of the text, depending on the task at hand. This makes them more efficient and faster than before. The key idea is to limit how many “pieces” (or tokens) the computer can work with at one time, so it learns to prioritize the most important ones. This results in computers that are just as good as before but use less energy and take less time to get things done. |
Keywords
* Artificial intelligence * Self attention * Token * Transformer
