Authors:
(1) David Raposo, Google DeepMind and with equal contribution;
(2) Sam Ritter, Google DeepMind;
(3) Blake Richards, Google DeepMind and McGill University & Mila;
(4) Timothy Lillicrap, Google DeepMind;
(5) Peter Conway Humphreys, Google DeepMind;
(6) Adam Santoro, Google DeepMind and with equal contribution. Editor's note: this is part 2 of 5 of a study detailing a way to make transformer-based language models more efficient by dynamically allocating computational resources. Read the rest below. Table of Links Introduction
Background
Implementing Mixture-of-Depths Transformers


3.1. Defining a compute budget


3.2. Routing around transformer blocks


3.3. Routing schemes


3.4. Routing implementation


3.5. Sampling and 3.6. Training methods



Results

4.1. Training, isoFLOP comparisons
4.2. Auto-regressive Evaluation and 4.3. Mixture-of-Depths-and-Experts (MoDE)


Discussion and References 2. Background The transformer architecture has become the workhorse of a revolution in practical artificial intelligence, bringing unprecedented capabilities at the cost of expensive training runs and serving procedures. This has spurred tremendous interest in making transformer architectures more efficient (Gupta and Agrawal, 2021; Tay et al., 2020). One of the promising approaches is conditional computation, whereby learned mechanisms determine when and how to expend computation. This terminology was introduced by Bengio (2013), and the concept was explored further over the next several years (Bengio et al., 2016, 2013; Cho and Bengio, 2014; Graves, 2016; Jernite et al., 2017; Wang et al., 2017). A wide variety of recent work has developed conditional computation methods for transformers. Some of this work focuses on "early exiting", that is, learning to decide when to end computation on a given token, allowing the token to skip any remaining transformer layers after the exit decision is made (Elbayad et al., 2019; Liu et al., 2021; Schuster et al., 2022). In MoD, unlike in early-exit methods, a token can skip middle layers, then be updated via self-attention with tokens that that have gone through all the middle layers. We speculate that this might be a useful property. Other work has developed methods for iterating transformer layers with shared weights for an adaptive number of steps (Dehghani et al., 2018; Simoulin and Crabbé, 2021). Bolya et al. (2023) developed a method for choosing tokens to merge when running inference on a trained vision transformer which notably requires no learning. Lei et al. (2023) make use of conditional computation in a fine tuning setting by building on adapter approaches (He et al., 2021) to learn to skip blocks of frozen pre-trained weights in favor of running only a small fine-tuned adapter. CoLT5 (Ainslie et al., 2023) uses conditional routing to select whether a given token will pass through a heavy or light pathway for each feedforward layer. Further, they use the same routing mechanism to select whether a token will attend to all other tokens or to a select few, as in Guo et al. (2022). Like MoD, CoLT5 uses soft top-k for making routing decisions. However, CoLT5 focuses on a encoder-decoder setting, and thus does need to contend with the problem of efficient sequential decoding given the non-causal nature of the top-k operation. In contrast, our current work with MoD focuses on the decoder-only setting, and so we propose a predictive router to enable efficient inference for conditional computation in transformers. One successful formulation of conditional computation is the the "mixture-of-experts" layer (MoE) as introduced by Shazeer et al. (2017). Developed initially in the context of LSTMs, later work showed compelling empirical results for MoE with transformers (Fedus et al., 2022; Lepikhin et al., 2020; Zoph et al., 2022). Unlike other conditional computation approaches that try to conserve or expend additional compute, MoE transformers use conditional logic to route tokens to one of many expert MLPs while keeping total compute expenditure constant. Our mixture-of-depths method can be thought of as using the routing logic from MoE transformers, but rather than having multiple experts, MoD deploys a single expert which can be dynamically skipped. This paper is available on arxiv under CC BY 4.0 DEED license. Authors: (1) David Raposo, Google DeepMind and with equal contribution; (2) Sam Ritter, Google DeepMind; (3) Blake Richards, Google DeepMind and McGill University & Mila; (4) Timothy Lillicrap, Google DeepMind; (5) Peter Conway Humphreys, Google DeepMind; (6) Adam Santoro, Google DeepMind and with equal contribution. Authors: Authors: (1) David Raposo, Google DeepMind and with equal contribution; (2) Sam Ritter, Google DeepMind; (3) Blake Richards, Google DeepMind and McGill University & Mila; (4) Timothy Lillicrap, Google DeepMind; (5) Peter Conway Humphreys, Google DeepMind; (6) Adam Santoro, Google DeepMind and with equal contribution. Editor's note: this is part 2 of 5 of a study detailing a way to make transformer-based language models more efficient by dynamically allocating computational resources. Read the rest below. Editor's note: this is part 2 of 5 of a study detailing a way to make transformer-based language models more efficient by dynamically allocating computational resources. Read the rest below. Editor's note: this is part 2 of 5 of a study detailing a way to make transformer-based language models more efficient by dynamically allocating computational resources. Read the rest below. Editor's note: this is part 2 of 5 of a study detailing a way to make transformer-based language models more efficient by dynamically allocating computational resources. Read the rest below. Table of Links Introduction Background Implementing Mixture-of-Depths Transformers


3.1. Defining a compute budget


3.2. Routing around transformer blocks


3.3. Routing schemes


3.4. Routing implementation


3.5. Sampling and 3.6. Training methods Results

4.1. Training, isoFLOP comparisons
4.2. Auto-regressive Evaluation and 4.3. Mixture-of-Depths-and-Experts (MoDE) Discussion and References Introduction Introduction Background Background Implementing Mixture-of-Depths Transformers 3.1. Defining a compute budget


3.2. Routing around transformer blocks


3.3. Routing schemes


3.4. Routing implementation


3.5. Sampling and 3.6. Training methods Implementing Mixture-of-Depths Transformers 3.1. Defining a compute budget 3.2. Routing around transformer blocks 3.3. Routing schemes 3.4. Routing implementation 3.5. Sampling and 3.6. Training methods 3.1. Defining a compute budget 3.1. Defining a compute budget 3.2. Routing around transformer blocks 3.2. Routing around transformer blocks 3.3. Routing schemes 3.3. Routing schemes 3.4. Routing implementation 3.4. Routing implementation 3.5. Sampling and 3.6. Training methods 3.5. Sampling and 3.6. Training methods Results 4.1. Training, isoFLOP comparisons
4.2. Auto-regressive Evaluation and 4.3. Mixture-of-Depths-and-Experts (MoDE) Results 4.1. Training, isoFLOP comparisons 4.2. Auto-regressive Evaluation and 4.3. Mixture-of-Depths-and-Experts (MoDE) 4.1. Training, isoFLOP comparisons 4.2. Auto-regressive Evaluation and 4.3. Mixture-of-Depths-and-Experts (MoDE) Discussion and References Discussion and References 2. Background The transformer architecture has become the workhorse of a revolution in practical artificial intelligence, bringing unprecedented capabilities at the cost of expensive training runs and serving procedures. This has spurred tremendous interest in making transformer architectures more efficient (Gupta and Agrawal, 2021; Tay et al., 2020). One of the promising approaches is conditional computation, whereby learned mechanisms determine when and how to expend computation. This terminology was introduced by Bengio (2013), and the concept was explored further over the next several years (Bengio et al., 2016, 2013; Cho and Bengio, 2014; Graves, 2016; Jernite et al., 2017; Wang et al., 2017). A wide variety of recent work has developed conditional computation methods for transformers. Some of this work focuses on "early exiting", that is, learning to decide when to end computation on a given token, allowing the token to skip any remaining transformer layers after the exit decision is made (Elbayad et al., 2019; Liu et al., 2021; Schuster et al., 2022). In MoD, unlike in early-exit methods, a token can skip middle layers, then be updated via self-attention with tokens that that have gone through all the middle layers. We speculate that this might be a useful property. Other work has developed methods for iterating transformer layers with shared weights for an adaptive number of steps (Dehghani et al., 2018; Simoulin and Crabbé, 2021). Bolya et al. (2023) developed a method for choosing tokens to merge when running inference on a trained vision transformer which notably requires no learning. Lei et al. (2023) make use of conditional computation in a fine tuning setting by building on adapter approaches (He et al., 2021) to learn to skip blocks of frozen pre-trained weights in favor of running only a small fine-tuned adapter. CoLT5 (Ainslie et al., 2023) uses conditional routing to select whether a given token will pass through a heavy or light pathway for each feedforward layer. Further, they use the same routing mechanism to select whether a token will attend to all other tokens or to a select few, as in Guo et al. (2022). Like MoD, CoLT5 uses soft top-k for making routing decisions. However, CoLT5 focuses on a encoder-decoder setting, and thus does need to contend with the problem of efficient sequential decoding given the non-causal nature of the top-k operation. In contrast, our current work with MoD focuses on the decoder-only setting, and so we propose a predictive router to enable efficient inference for conditional computation in transformers. One successful formulation of conditional computation is the the "mixture-of-experts" layer (MoE) as introduced by Shazeer et al. (2017). Developed initially in the context of LSTMs, later work showed compelling empirical results for MoE with transformers (Fedus et al., 2022; Lepikhin et al., 2020; Zoph et al., 2022). Unlike other conditional computation approaches that try to conserve or expend additional compute, MoE transformers use conditional logic to route tokens to one of many expert MLPs while keeping total compute expenditure constant. Our mixture-of-depths method can be thought of as using the routing logic from MoE transformers, but rather than having multiple experts, MoD deploys a single expert which can be dynamically skipped. This paper is available on arxiv under CC BY 4.0 DEED license. This paper is available on arxiv under CC BY 4.0 DEED license. available on arxiv

Part of HackerNoon's growing list of open-source research papers, promoting free access to academic material.

Google

This Clever AI Hack Could Cut Processing Costs in Half

Google Researchers Develop New AI Tech That Doesn't Waste Brainpower on Useless Words

AI Models on HackerNoon

Too Long; Didn't Read

Make resilience your competitive advantage

New AI Method Lets Models Decide What to Think About

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

Untitled Story

Unlock Peak Mobile Performance: A Deep Dive into PowerInfer-2's Neuron-Aware Runtime

The Noonification: Use This 7-Step McKinsey Framework to Solve Any Problem (1/10/2023)

The Noonification: A Taxonomy of Inclusiveness (1/11/2024)

The Noonification: What is the InfiniteNature-Zero AI Model? (11/19/2022)

10 Ways AI Has Changed Our Lives

Unlock Peak Mobile Performance: A Deep Dive into PowerInfer-2's Neuron-Aware Runtime

The Noonification: Use This 7-Step McKinsey Framework to Solve Any Problem (1/10/2023)

The Noonification: A Taxonomy of Inclusiveness (1/11/2024)

The Noonification: What is the InfiniteNature-Zero AI Model? (11/19/2022)

10 Ways AI Has Changed Our Lives

Light-Mode

Classic

Newspaper

Dark-Mode

Neon Noir

Minty

HN StartUps