Mixture-of-Depths: Dynamically allocating compute in transformer-based language models

Understand

Transformer-based language models spread FLOPs uniformly across input sequences.

• In this work we demonstrate that transformers can instead learn to dynamically allocate FLOPs (or compute) to specific positions in a sequence, optimising the allocation along the sequence for different layers across the model depth.
• Our method enforces a total compute budget by capping the number of tokens ($k$) that can participate in the self-attention and MLP computations at a given layer.
• The tokens to be processed are determined by the network using a top-$k$ routing mechanism.