title: Mixture-of-Depths: Dynamically allocating compute in transformer-based
    language models
summary: 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. Since k is defined a priori, this simple
  procedure uses a static computation graph with known tensor sizes, unlike other
  conditional computation techniques. Nevertheless, since the identities of the
  k tokens are fluid, this method can expend FLOPs non-uniformly across the
  time and model depth dimensions. Thus, compute expenditure is entirely
  predictable in sum total, but dynamic and context-sensitive at the token-level.
  Not only do models trained in this way learn to dynamically allocate compute,
  they do so efficiently. These models match baseline performance for equivalent
  FLOPS and wall-clock times to train, but require a fraction of the FLOPs per
  forward pass, and can be upwards of 50\% faster to step during post-training
  sampling.
paper_authors._id: 660e2a251714de5d5bd49825
paper_authors.name: David Raposo
paper_authors._id: 660e2a251714de5d5bd49826
paper_authors.name: Sam Ritter
paper_authors._id: 660e2a251714de5d5bd49827
paper_authors.name: Blake Richards
paper_authors._id: 660e2a251714de5d5bd49828
paper_authors.name: Timothy Lillicrap
paper_authors._id: 660e2a251714de5d5bd49829
paper_authors.name: Peter Conway Humphreys
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2024-04-25T07:34:08.695Z
paper_authors._id: 638524061769b7c4b10ca8af
paper_authors.avatarUrl: /avatars/b2167b6777364b24ea3ad433655fa66d.svg
paper_authors.fullname: Peter Humphreys
paper_authors.type: user
paper_authors.user: peterhumphreys
paper_authors._id: 660e2a251714de5d5bd4982a
paper_authors.name: Adam Santoro
paper_title: Mixture-of-Depths: Dynamically allocating compute in transformer-based
    language models
paper_summary: 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. Since k is defined a priori, this simple
  procedure uses a static computation graph with known tensor sizes, unlike other
  conditional computation techniques. Nevertheless, since the identities of the
  k tokens are fluid, this method can expend FLOPs non-uniformly across the
  time and model depth dimensions. Thus, compute expenditure is entirely
  predictable in sum total, but dynamic and context-sensitive at the token-level.
  Not only do models trained in this way learn to dynamically allocate compute,
  they do so efficiently. These models match baseline performance for equivalent
  FLOPS and wall-clock times to train, but require a fraction of the FLOPs per
  forward pass, and can be upwards of 50\% faster to step during post-training
  sampling.
paper_ai_summary: Transformer-based models can dynamically allocate computational resources to specific positions in input sequences, leading to efficient performance with reduced FLOPs.
paper_ai_keywords: Transformer-based language models FLOPs self-attention MLP top-k routing mechanism model depth tensor sizes compute expenditure wall-clock times forward pass post-training sampling