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

• URL: http://arxiv.org/abs/2404.02258v1
• Date: Tue, 2 Apr 2024 19:28:11 GMT
• Title: Mixture-of-Depths: Dynamically allocating compute in transformer-based language models
• Authors: David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, Adam Santoro,
• Abstract summary: Transformer-based language models spread FLOPs uniformly across input sequences. We show that transformers can learn to dynamically allocate FLOPs to specific positions in a sequence.
• Score: 8.774705201394916
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Algorithmic Capabilities of Random Transformers [49.73113518329544]
  We investigate what functions can be learned by randomly transformers in which only the embedding layers are optimized. We find that these random transformers can perform a wide range of meaningful algorithmic tasks. Our results indicate that some algorithmic capabilities are present in transformers even before these models are trained.
  arXiv  Detail & Related papers  (2024-10-06T06:04:23Z)
• IT$^3$: Idempotent Test-Time Training [95.78053599609044]
  This paper introduces Idempotent Test-Time Training (IT$3$), a novel approach to addressing the challenge of distribution shift. IT$3$ is based on the universal property of idempotence. We demonstrate the versatility of our approach across various tasks, including corrupted image classification.
  arXiv  Detail & Related papers  (2024-10-05T15:39:51Z)
• Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers [54.20763128054692]
  We study how a two-attention-layer transformer is trained to perform ICL on $n$-gram Markov chain data. We prove that the gradient flow with respect to a cross-entropy ICL loss converges to a limiting model.
  arXiv  Detail & Related papers  (2024-09-09T18:10:26Z)
• Adaptive Computation Modules: Granular Conditional Computation For Efficient Inference [13.000030080938078]
  computational cost of transformer models makes them inefficient in low-latency or low-power applications. We introduce the Adaptive Computation Module (ACM), a generic module that dynamically adapts its computational load to match the estimated difficulty of the input on a per-token basis. Our evaluation of transformer models in computer vision and speech recognition demonstrates that substituting layers with ACMs significantly reduces inference costs without degrading the downstream accuracy for a wide interval of user-defined budgets.
  arXiv  Detail & Related papers  (2023-12-15T20:39:43Z)
• Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition [62.83832841523525]
  We propose a fast and accurate parallel transformer, termed Paraformer. It accurately predicts the number of output tokens and extract hidden variables. It can attain comparable performance to the state-of-the-art AR transformer, with more than 10x speedup.
  arXiv  Detail & Related papers  (2022-06-16T17:24:14Z)
• You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling [38.34914626128062]
  We show that a Bernoulli sampling attention mechanism based on Locality Sensitive Hashing (LSH) decreases the quadratic complexity of such models to linear. We evaluate our algorithm on the GLUE benchmark with standard 512 sequence length where we see favorable performance relative to a standard pretrained Transformer.
  arXiv  Detail & Related papers  (2021-11-18T14:24:34Z)
• Classification and Feature Transformation with Fuzzy Cognitive Maps [0.3299672391663526]
  Fuzzy Cognitive Maps (FCMs) are considered a soft computing technique combining elements of fuzzy logic and recurrent neural networks. In this work we propose an FCM based classifier with a fully connected map structure. Weights were learned with a gradient algorithm and logloss or cross-entropy were used as the cost function.
  arXiv  Detail & Related papers  (2021-03-08T22:26:24Z)
• Self Normalizing Flows [65.73510214694987]
  We propose a flexible framework for training normalizing flows by replacing expensive terms in the gradient by learned approximate inverses at each layer. This reduces the computational complexity of each layer's exact update from $mathcalO(D3)$ to $mathcalO(D2)$. We show experimentally that such models are remarkably stable and optimize to similar data likelihood values as their exact gradient counterparts.
  arXiv  Detail & Related papers  (2020-11-14T09:51:51Z)
• Dynamic Tensor Rematerialization [11.204761128308542]
  Checkpointing enables the training of deep learning models under restricted memory budgets. Current checkpointing techniques statically plan these recomputations offline and assume static graphs. We demonstrate that a simple online algorithm can achieve comparable performance by introducing Dynamic Rematerialization (DTR)
  arXiv  Detail & Related papers  (2020-06-17T02:49:59Z)
• $O(n)$ Connections are Expressive Enough: Universal Approximability of Sparse Transformers [71.31712741938837]
  We show that sparse Transformers with only $O(n)$ connections per attention layer can approximate the same function class as the dense model with $n2$ connections. We also present experiments comparing different patterns/levels of sparsity on standard NLP tasks.
  arXiv  Detail & Related papers  (2020-06-08T18:30:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.