Original: P ARALLEL T RAINING OF D EEP N ETWORKS WITH L OCAL U PDATES
Revised: P ARALLEL T RAINING OF D EEP N ETWORKS WITH L OCAL U PDATES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep learning models trained on large data sets have been widely successful in both vision and language domains.
Revised: Deep learning models trained on large data sets have been widely successful in both vision and language domains.

Original: As state-of-the-art deep learning architectures have continued to grow in parameter count so have the compute budgets and times required to train them, increasing the need for compute-efﬁcient methods that parallelize training.
Revised: As state-of-the-art deep learning architectures have continued to grow in parameter count so have the compute budgets and times required to train them, increasing the need for compute-efﬁcient methods that parallelize training.

Original: Two common approaches to parallelize the training of deep networks have been data and model parallelism.
Revised: Two common approaches to parallelize the training of deep networks have been data and model parallelism.

Original: While useful, data and model parallelism suffer from diminishing returns in terms of compute efﬁciency for large batch sizes.
Revised: While useful, data and model parallelism suffer from diminishing returns in terms of compute efﬁciency for large batch sizes.

Original: In this paper, we investigate how to continue scaling compute efﬁciently beyond the point of diminishing returns for large batches through local parallelism , a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation.
Revised: In this paper, we investigate how to continue scaling compute efﬁciently beyond the point of diminishing returns for large batches through local parallelism , a framework which parallelizes training of individual layers in deep networks by replacing global backpropagation with truncated layer-wise backpropagation.

Original: Local parallelism enables fully asynchronous layer-wise parallelism with a low memory footprint, and requires little communication overhead compared with model parallelism.
Revised: Local parallelism enables fully asynchronous layer-wise parallelism with a low memory footprint, and requires little communication overhead compared with model parallelism.

Original: We show results in both vision and language domains across a diverse set of architectures, and ﬁnd that local parallelism is particularly effective in the high-compute regime.
Revised: We show results in both vision and language domains across a diverse set of architectures, and ﬁnd that local parallelism is particularly effective in the high-compute regime.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Backpropagation (Rumelhart et al., 1985) is by far the most common method used to train neural networks.
Revised: Backpropagation (Rumelhart et al., 1985) is by far the most common method used to train neural networks.

Original: Alternatives to backpropagation are typically used only when backpropagation is impractical due to a non-differentiable loss (Schulman et al., 2015), non-smooth loss landscape (Metz et al., 2019), or due to memory and/or compute requirements (Ororbia et al., 2020).
Revised: Alternatives to backpropagation are typically used only when backpropagation is impractical due to a non-differentiable loss (Schulman et al., 2015), non-smooth loss landscape (Metz et al., 2019), or due to memory and/or compute requirements (Ororbia et al., 2020).

Original: However, progress in deep learning is producing ever larger models in terms of parameter count and depth, in vision (H ´ enaff et al., 2019;
Revised: However, progress in deep learning is producing ever larger models in terms of parameter count and depth, in vision (H ´ enaff et al., 2019;

Original: Chen et al., 2020), language (Radford et al., 2019; Brown et al., 2020), and many other domains (Silver et al., 2017; Vinyals et al., 2019; Berner et al., 2019).
Revised: Chen et al., 2020), language (Radford et al., 2019; Brown et al., 2020), and many other domains (Silver et al., 2017; Vinyals et al., 2019; Berner et al., 2019).

Original: As model size increases, backpropagation incurs growing computational, memory, and synchronization overhead (Ben-Nun & Hoeﬂer, 2018).
Revised: As model size increases, backpropagation incurs growing computational, memory, and synchronization overhead (Ben-Nun & Hoeﬂer, 2018).

Original: This raises the question of whether there are more efﬁcient training strategies, even for models and losses that are considered well matched to training by backpropagation.
Revised: This raises the question of whether there are more efﬁcient training strategies, even for models and losses that are considered well matched to training by backpropagation.

Original: Much of the work on training large scale models focuses on designing compute infrastructure which makes backpropagation more efﬁcient, despite growing model size (Dean et al., 2012b; Chen et al., 2015; Sergeev & Balso, 2018).
Revised: Much of the work on training large scale models focuses on designing compute infrastructure which makes backpropagation more efﬁcient, despite growing model size (Dean et al., 2012b; Chen et al., 2015; Sergeev & Balso, 2018).

Original: One of the most common ways to achieve efﬁcient training of deep neural networks with backpropagation is to scale utilizing data parallelism (Zhang et al., 1989; Chen et al., 2016), training on bigger batch sizes spread across multiple devices.
Revised: One of the most common ways to achieve efﬁcient training of deep neural networks with backpropagation is to scale utilizing data parallelism (Zhang et al., 1989; Chen et al., 2016), training on bigger batch sizes spread across multiple devices.

Original: However, diminishing returns have been reported with this method for larger batch sizes, effectively wasting compute (Goyal et al., 2017; Masters & Luschi, 2018; Shallue et al., 2018;
Revised: However, diminishing returns have been reported with this method for larger batch sizes, effectively wasting compute (Goyal et al., 2017; Masters & Luschi, 2018; Shallue et al., 2018;

Original: McCandlish et al., 2018).
Revised: McCandlish et al., 2018).

Original: Training based on pipeline parallelism has also been introduced, but still requires large batches for efﬁcient training (Petrowski et al., 1993; Ben-Nun & Hoeﬂer, 2018; Huang et al., 2019).
Revised: Training based on pipeline parallelism has also been introduced, but still requires large batches for efﬁcient training (Petrowski et al., 1993; Ben-Nun & Hoeﬂer, 2018; Huang et al., 2019).

Original: Moreover, in addition to the limitation that in the forward pass each layer can only process the input data in sequence ( forward locking ), the use of backpropagation implies that the network parameters of each layer can only be updated in turn after completing the full forward pass ( backward locking ).
Revised: Moreover, in addition to the limitation that in the forward pass each layer can only process the input data in sequence ( forward locking ), the use of backpropagation implies that the network parameters of each layer can only be updated in turn after completing the full forward pass ( backward locking ).

Original: This backward locking results in increased memory overhead, and precludes efﬁcient parallel processing across layers (Jaderberg et al., 2017).
Revised: This backward locking results in increased memory overhead, and precludes efﬁcient parallel processing across layers (Jaderberg et al., 2017).

Original: The challenges of scaling compute infrastructure to support deep networks trained with backpropagation motivate the need for alternative approaches to training deep neural networks.
Revised: The challenges of scaling compute infrastructure to support deep networks trained with backpropagation motivate the need for alternative approaches to training deep neural networks.

Original: [Figure]
Revised: [Figure]

Original: While data, model, and pipeline parallelism are existing paradigms for parallelizing learning, we investigate another way of parallelizing learning through local layer-wise training shown in (d).
Revised: While data, model, and pipeline parallelism are existing paradigms for parallelizing learning, we investigate another way of parallelizing learning through local layer-wise training shown in (d).

Original: In this work, we explore how layer-wise local updates (Belilovsky et al., 2019a; L¨owe et al., 2019;
Revised: In this work, we explore how layer-wise local updates (Belilovsky et al., 2019a; L¨owe et al., 2019;

Original: Xiong et al., 2020) can help overcome these challenges and scale more efﬁciently with compute than backpropagation.
Revised: Xiong et al., 2020) can help overcome these challenges and scale more efﬁciently with compute than backpropagation.

Original: With local updates, each layer is updated before even completing a full forward pass through the network.
Revised: With local updates, each layer is updated before even completing a full forward pass through the network.

Original: This remedies the forward and backward locking problems which harm memory efﬁciency and update latency in standard backprop.
Revised: This remedies the forward and backward locking problems which harm memory efﬁciency and update latency in standard backprop.

Original: Layer-wise local updates are not proportional to gradients of the original loss, and are not even guaranteed to descend a loss function.
Revised: Layer-wise local updates are not proportional to gradients of the original loss, and are not even guaranteed to descend a loss function.

Original: Nevertheless, in practice they are effective at training neural networks.
Revised: Nevertheless, in practice they are effective at training neural networks.

Original: We refer to this approach of parallelizing compute, which is alternative and complementary to data and model parallelism, as local parallelism .
Revised: We refer to this approach of parallelizing compute, which is alternative and complementary to data and model parallelism, as local parallelism .

Original: Our investigation focuses on the trade-offs of using local update methods as opposed to global backpropagation.
Revised: Our investigation focuses on the trade-offs of using local update methods as opposed to global backpropagation.

Original: To summarize our contributions: (i)
Revised: To summarize our contributions: (i)

Original: We provide the ﬁrst large scale investigation into local update methods in both vision and language domains.
Revised: We provide the ﬁrst large scale investigation into local update methods in both vision and language domains.

Original: We ﬁnd training speedups (as measured by the reduction in required sequential compute steps) of up to 10 × on simple MLPs, and 2 × on Transformer architectures.
Revised: We ﬁnd training speedups (as measured by the reduction in required sequential compute steps) of up to 10 × on simple MLPs, and 2 × on Transformer architectures.

Original: These training speedups are the result of local training methods being able to leverage more parallel compute than backprop.
Revised: These training speedups are the result of local training methods being able to leverage more parallel compute than backprop.

Original: (ii) We provide insight into how local parallelism methods work, and experimentally compare the similarity of their gradient and features to those from backprop.
Revised: (ii) We provide insight into how local parallelism methods work, and experimentally compare the similarity of their gradient and features to those from backprop.

Original: (iii)
Revised: (iii)

Original: We demonstrate a prototype implementation of local parallelism for
Revised: We demonstrate a prototype implementation of local parallelism for

Original: ResNets, and show up to a 40% increase in sample throughput (number of training points per second) relative to backprop, due to higher hardware utilization.
Revised: ResNets, and show up to a 40% increase in sample throughput (number of training points per second) relative to backprop, due to higher hardware utilization.

Original: We believe that local parallelism will provide beneﬁts whenever there are diminishing returns from data parallelism, and avoid stale weights from pipelined model parallelism.
Revised: We believe that local parallelism will provide beneﬁts whenever there are diminishing returns from data parallelism, and avoid stale weights from pipelined model parallelism.

Original: Additionally, we have released code showing an example of local parallelism, available at hiddenurl .
Revised: Additionally, we have released code showing an example of local parallelism, available at hiddenurl .

Original:  R ELATED W ORK 2. P ARALLELIZATION IN D EEP L EARNING
Revised:  R ELATED W ORK 2. P ARALLELIZATION IN D EEP L EARNING

Original: Scaling large models has led to the development of a number of techniques to train deep models in a parallel fashion (Ben-Nun & Hoeﬂer, 2018), summarized in Figure 1.
Revised: Scaling large models has led to the development of a number of techniques to train deep models in a parallel fashion (Ben-Nun & Hoeﬂer, 2018), summarized in Figure 1.

Original: Data Parallelism : Data Parallelism (Zhang et al., 1989) is an attempt to speed up training of a model by splitting the data among multiple identical models and training each model on a shard of the data independently.
Revised: Data Parallelism : Data Parallelism (Zhang et al., 1989) is an attempt to speed up training of a model by splitting the data among multiple identical models and training each model on a shard of the data independently.

Original: Data parallelism is effectively training with larger minibatches (Kaplan et al., 2020).
Revised: Data parallelism is effectively training with larger minibatches (Kaplan et al., 2020).

Original: This creates issues around the consistency of a model which then needs to be synchronized (Deng et al., 2012; Dean et al., 2012a).
Revised: This creates issues around the consistency of a model which then needs to be synchronized (Deng et al., 2012; Dean et al., 2012a).

Original: There are two main ways to synchronize weights across model copies: (i) Synchronous optimization , where data parallel training synchronizes at the end of every minibatch (Das et al., 2016; Chen et al., 2016), with a communication overhead that increases with the number of devices; (ii) Asynchronous optimization that implements data parallel training with independent updates of local model parameters without global synchronization (Niu et al., 2011; Dean et al., 2012a) – this increases device utilization, but empirically gradients are computed on stale weights, which results in a poor sample efﬁciency and thus slower overall training time compared to synchronous optimization.
Revised: There are two main ways to synchronize weights across model copies: (i) Synchronous optimization , where data parallel training synchronizes at the end of every minibatch (Das et al., 2016; Chen et al., 2016), with a communication overhead that increases with the number of devices; (ii) Asynchronous optimization that implements data parallel training with independent updates of local model parameters without global synchronization (Niu et al., 2011; Dean et al., 2012a) – this increases device utilization, but empirically gradients are computed on stale weights, which results in a poor sample efﬁciency and thus slower overall training time compared to synchronous optimization.

Original: Model Parallelism : Model Parallelism is used when a model is too large to ﬁt in the memory of a single device and is instead spread over multiple processors (Krizhevsky et al., 2012; Shazeer et al., 2018; Harlap et al., 2018;
Revised: Model Parallelism : Model Parallelism is used when a model is too large to ﬁt in the memory of a single device and is instead spread over multiple processors (Krizhevsky et al., 2012; Shazeer et al., 2018; Harlap et al., 2018;

Original: Lepikhin et al., 2020).
Revised: Lepikhin et al., 2020).

Original: This is increasingly common as state of the art performance continues to improve with increasing model size (Brown et al., 2020).
Revised: This is increasingly common as state of the art performance continues to improve with increasing model size (Brown et al., 2020).

Original: Model parallelism unfortunately has a few downsides: (i) High communication costs – the total training time for larger networks can become dominated by communication costs (Simonyan & Zisserman, 2015), which in the worst case can grow quadratically with the number of devices, and can reach up to 85% of the total training time of a large model such as VGG-16 (Harlap et al., 2018; Simonyan & Zisserman, 2015); (ii) Device under-utilization – forward propagation and backward propagation are both synchronous operations, which can result in processor under-utilization in model-parallel systems.
Revised: Model parallelism unfortunately has a few downsides: (i) High communication costs – the total training time for larger networks can become dominated by communication costs (Simonyan & Zisserman, 2015), which in the worst case can grow quadratically with the number of devices, and can reach up to 85% of the total training time of a large model such as VGG-16 (Harlap et al., 2018; Simonyan & Zisserman, 2015); (ii) Device under-utilization – forward propagation and backward propagation are both synchronous operations, which can result in processor under-utilization in model-parallel systems.

Original: This problem becomes worse as we increase the number of layers (Ben-Nun & Hoeﬂer, 2018; Jia et al., 2014; Collobert et al., 2011; Abadi et al., 2016; Huang et al., 2018).
Revised: This problem becomes worse as we increase the number of layers (Ben-Nun & Hoeﬂer, 2018; Jia et al., 2014; Collobert et al., 2011; Abadi et al., 2016; Huang et al., 2018).

Original: Pipeline Parallelism : Due to the forward and backward locking, using multiple devices to process consecutive blocks of the deep model would make an inefﬁcient use of the hardware resources.
Revised: Pipeline Parallelism : Due to the forward and backward locking, using multiple devices to process consecutive blocks of the deep model would make an inefﬁcient use of the hardware resources.

Original: Pipelining (Harlap et al., 2018) concurrently passes multiple mini-batches to multiple layers on multiple devices.
Revised: Pipelining (Harlap et al., 2018) concurrently passes multiple mini-batches to multiple layers on multiple devices.

Original: This increases device utilization but can introduce staleness and consistency issues which lead to unstable training.
Revised: This increases device utilization but can introduce staleness and consistency issues which lead to unstable training.

Original: Harlap et al.
Revised: Harlap et al.

Original: (2018) alleviates the consistency issue by storing past versions of each layer.
Revised: (2018) alleviates the consistency issue by storing past versions of each layer.

Original: Huang et al.
Revised: Huang et al.

Original: (2019) addresses the staleness issue by pipelining microbatches and synchronously updating at the end of each minibatch.
Revised: (2019) addresses the staleness issue by pipelining microbatches and synchronously updating at the end of each minibatch.

Original: Guan et al.
Revised: Guan et al.

Original: (2019) builds on this work by introducing a weight prediction strategy and Yang et al.
Revised: (2019) builds on this work by introducing a weight prediction strategy and Yang et al.

Original: (2020) investigates to what extent the tradeoff between staleness/consistency and device utilization is necessary.
Revised: (2020) investigates to what extent the tradeoff between staleness/consistency and device utilization is necessary.

Original: Local updates on the other hand can keep device utilization high with both small and large batches and avoid the weight staleness problem.
Revised: Local updates on the other hand can keep device utilization high with both small and large batches and avoid the weight staleness problem.

Original: Local Learning Rules
Revised: Local Learning Rules

Original: : Local learning describes a family of methods that perform parameter updates based only on local information, where locality is deﬁned as dependence of neighboring neurons, layers, or groups of layers.
Revised: : Local learning describes a family of methods that perform parameter updates based only on local information, where locality is deﬁned as dependence of neighboring neurons, layers, or groups of layers.

Original: The earliest local method we are aware of is Hebbian Learning (Hebb, 1949) which has further been explored in BCM theory (Izhikevich & Desai, 2003; Coesmans et al., 2004), Oja’s rule (Oja, 1982), Generalized Hebbian Learning (Sanger, 1989), and meta-learned local learning rules (Bengio et al., 1990; 1992; Metz et al., 2018;
Revised: The earliest local method we are aware of is Hebbian Learning (Hebb, 1949) which has further been explored in BCM theory (Izhikevich & Desai, 2003; Coesmans et al., 2004), Oja’s rule (Oja, 1982), Generalized Hebbian Learning (Sanger, 1989), and meta-learned local learning rules (Bengio et al., 1990; 1992; Metz et al., 2018;

Original: Gu et al., 2019).
Revised: Gu et al., 2019).

Original: Architectures like Hopﬁeld Networks (Hopﬁeld, 1982) and Boltzmann Machines (Ackley et al., 1985) also employ a local update, and predate backprogation in deep learning.
Revised: Architectures like Hopﬁeld Networks (Hopﬁeld, 1982) and Boltzmann Machines (Ackley et al., 1985) also employ a local update, and predate backprogation in deep learning.

Original: Modern variants of local training methods have attempted to bridge the performance gap with backpropagation.
Revised: Modern variants of local training methods have attempted to bridge the performance gap with backpropagation.

Original: These include projection methods such as Hebbian learning rules for deep networks (Krotov & Hopﬁeld, 2019; Grinberg et al., 2019; Ryali et al., 2020), and local layer-wise learning with auxiliary losses (Belilovsky et al., 2019a;b).
Revised: These include projection methods such as Hebbian learning rules for deep networks (Krotov & Hopﬁeld, 2019; Grinberg et al., 2019; Ryali et al., 2020), and local layer-wise learning with auxiliary losses (Belilovsky et al., 2019a;b).

Original: Most similar to our work is decoupled greedy layer-wise learning (Belilovsky et al., 2019b; L ¨ owe et al., 2019), which trained auxiliary image classiﬁers greedily, and local contrastive learning (Xiong et al., 2020).
Revised: Most similar to our work is decoupled greedy layer-wise learning (Belilovsky et al., 2019b; L ¨ owe et al., 2019), which trained auxiliary image classiﬁers greedily, and local contrastive learning (Xiong et al., 2020).

Original: These methods mainly focus on matching the performance of backpropagation with respect to training epochs, whereas our work focuses on tradeoffs.
Revised: These methods mainly focus on matching the performance of backpropagation with respect to training epochs, whereas our work focuses on tradeoffs.

Original: Finally, while not local in the sense that parallelized layers still optimize for the global objective,
Revised: Finally, while not local in the sense that parallelized layers still optimize for the global objective,

Original: Huo et al.
Revised: Huo et al.

Original: (2018b) parallelize layers by caching gradients and using delayed gradient signals to overcome the backward locking problem and update decoupled layers in parallel.
Revised: (2018b) parallelize layers by caching gradients and using delayed gradient signals to overcome the backward locking problem and update decoupled layers in parallel.

Original:  L OCAL P ARALLELISM
Revised:  L OCAL P ARALLELISM

Original: Given a deep neural network, we divide the layers into a sequence of J blocks, which may contain one or more layers.
Revised: Given a deep neural network, we divide the layers into a sequence of J blocks, which may contain one or more layers.

Original: Each block is trained independently with an auxiliary objective, and receives the activations output by the previous block as input or, in the case of the ﬁrst block, the data from the sampled minibatch.
Revised: Each block is trained independently with an auxiliary objective, and receives the activations output by the previous block as input or, in the case of the ﬁrst block, the data from the sampled minibatch.

Original: We consider ﬁve variants to train this sequence of J blocks: backpropagation, greedy local parallelism, overlapping local parallelism, and chunked local parallelism, as shown in
Revised: We consider ﬁve variants to train this sequence of J blocks: backpropagation, greedy local parallelism, overlapping local parallelism, and chunked local parallelism, as shown in

Original: Figure 2.
Revised: Figure 2.

Original: We also include a baseline method of just training the last, or last two, layers.
Revised: We also include a baseline method of just training the last, or last two, layers.

Original: In all of the local methods, training occurs by attaching objective functions to the end of each block and back propagating the signal locally into the corresponding block or blocks.
Revised: In all of the local methods, training occurs by attaching objective functions to the end of each block and back propagating the signal locally into the corresponding block or blocks.

Original: In this work the auxiliary objective functions that we use take the same form as the global objective.
Revised: In this work the auxiliary objective functions that we use take the same form as the global objective.

Original: For example, to train a classiﬁer on CIFAR-10, we attach auxiliary linear classiﬁers to each local block.
Revised: For example, to train a classiﬁer on CIFAR-10, we attach auxiliary linear classiﬁers to each local block.

Original: See Belilovsky et al.
Revised: See Belilovsky et al.

Original: (2019b) for further discussion on the form of this objective.
Revised: (2019b) for further discussion on the form of this objective.

Original: [Figure]
Revised: [Figure]

Original: Backpropagation:
Revised: Backpropagation:

Original: In our notation, backpropagation groups all layers into one block and thus J = 1 .
Revised: In our notation, backpropagation groups all layers into one block and thus J = 1 .

Original: The parameters are updated with one instance of global error correction.
Revised: The parameters are updated with one instance of global error correction.

Original: While backpropagation ensures that all weights are updated according to the ﬁnal output loss, it also suffers from forward and backward locking (Jaderberg et al., 2017), an issue that local parallelized methods aim to resolve.
Revised: While backpropagation ensures that all weights are updated according to the ﬁnal output loss, it also suffers from forward and backward locking (Jaderberg et al., 2017), an issue that local parallelized methods aim to resolve.

Original: Greedy local parallelism: A straightforward approach to enable local training is to attach an auxiliary network to each local layer, which generates predictions from the activations of hidden layers.
Revised: Greedy local parallelism: A straightforward approach to enable local training is to attach an auxiliary network to each local layer, which generates predictions from the activations of hidden layers.

Original: After generating predictions, each local gradient is backpropagated to its respective local block, shown in Figure 2(b).
Revised: After generating predictions, each local gradient is backpropagated to its respective local block, shown in Figure 2(b).

Original: The activations are then passed as input to the next layer.
Revised: The activations are then passed as input to the next layer.

Original: We refer to this approach, introduced in (Belilovsky et al., 2019b), as greedy .
Revised: We refer to this approach, introduced in (Belilovsky et al., 2019b), as greedy .

Original: Greedy local parallelism is the most parallelizable of all the schemes we consider.
Revised: Greedy local parallelism is the most parallelizable of all the schemes we consider.

Original: However, a potential downside is that fully greedy updates force the layers to learn features that are only relevant to their local objective and preclude inter-layer communication, which may result in lower evaluation performance for the global objective, or worse generalization.
Revised: However, a potential downside is that fully greedy updates force the layers to learn features that are only relevant to their local objective and preclude inter-layer communication, which may result in lower evaluation performance for the global objective, or worse generalization.

Original: Overlapping local parallelism: One issue with the purely greedy approach is that features learned for any individual block may not be useful for subsequent blocks, since there is no inter-block propagation of gradient.
Revised: Overlapping local parallelism: One issue with the purely greedy approach is that features learned for any individual block may not be useful for subsequent blocks, since there is no inter-block propagation of gradient.

Original: For this reason, we consider overlapping local architectures where the ﬁrst layer of each block is also the last layer of the previous block, as shown in Figure 2(c), though overlapping of more layers is also possible.
Revised: For this reason, we consider overlapping local architectures where the ﬁrst layer of each block is also the last layer of the previous block, as shown in Figure 2(c), though overlapping of more layers is also possible.

Original: This redundancy enables inter-block propagation of gradient that is still local, since only neighboring blocks overlap.
Revised: This redundancy enables inter-block propagation of gradient that is still local, since only neighboring blocks overlap.

Original: However, this comes at the cost of running additional backward passes.
Revised: However, this comes at the cost of running additional backward passes.

Original: The overlapping architecture has appeared before in Xiong et al.
Revised: The overlapping architecture has appeared before in Xiong et al.

Original: but was used only for contrastive losses.
Revised: but was used only for contrastive losses.

Original: Ours is the ﬁrst work to investigate overlapping local architectures for standard prediction objectives in computer vision and language.
Revised: Ours is the ﬁrst work to investigate overlapping local architectures for standard prediction objectives in computer vision and language.

Original: Overlapping updates are parallelizable, but come with the additional complexity of keeping duplicates of the overlapping components and averaging updates for these layers.
Revised: Overlapping updates are parallelizable, but come with the additional complexity of keeping duplicates of the overlapping components and averaging updates for these layers.

Original: Chunked local parallelism : The greedy architecture is maximally parallel in the sense that it distributes one layer per block.
Revised: Chunked local parallelism : The greedy architecture is maximally parallel in the sense that it distributes one layer per block.

Original: However, it is also possible to have fewer parallel blocks by combining multiple layers into one.
Revised: However, it is also possible to have fewer parallel blocks by combining multiple layers into one.

Original: We refer to this architecture, shown in Figure 2(d), as chunked local parallelism.
Revised: We refer to this architecture, shown in Figure 2(d), as chunked local parallelism.

Original: This method trades off parallelizability and therefore throughput for an error signal that propagates through more consecutive layers.
Revised: This method trades off parallelizability and therefore throughput for an error signal that propagates through more consecutive layers.

Original: It differs from overlapping local parallelism by not needing to duplicate any layer.
Revised: It differs from overlapping local parallelism by not needing to duplicate any layer.

Original: While previous work has investigated the asymptotic performance of chunked parallelism (Belilovsky et al., 2019b), ours is the ﬁrst to consider the compute efﬁciency and parallelizability of local parallelism.
Revised: While previous work has investigated the asymptotic performance of chunked parallelism (Belilovsky et al., 2019b), ours is the ﬁrst to consider the compute efﬁciency and parallelizability of local parallelism.

Original: By stacking multiple layers per each parallelized block, chunked parallelism sits between fully parallelized methods, such as greedy and overlapping updates, and fully sequential methods like backpropagation.
Revised: By stacking multiple layers per each parallelized block, chunked parallelism sits between fully parallelized methods, such as greedy and overlapping updates, and fully sequential methods like backpropagation.

Original:  E FFICIENT T RAINING ON P ARETO F RONTIERS
Revised:  E FFICIENT T RAINING ON P ARETO F RONTIERS

Original: We explore the trade off between total computational cost and the amount of wallclock time needed to train a particular machine learning model to a target performance, similar to the analysis in Mc-
Revised: We explore the trade off between total computational cost and the amount of wallclock time needed to train a particular machine learning model to a target performance, similar to the analysis in Mc-

Original: [Figure]
Revised: [Figure]

Original: Candlish et al.
Revised: Candlish et al.

Original: We use ﬂoating point operations (FLOPs) as our unit of both cost and time, as they do not couple us to a particular choice of hardware.
Revised: We use ﬂoating point operations (FLOPs) as our unit of both cost and time, as they do not couple us to a particular choice of hardware.

Original: Cost is proportional to the total FLOPs used.
Revised: Cost is proportional to the total FLOPs used.

Original: We report time as the number of sequential FLOPs needed assuming we can run each example, and in the case of the local methods, each layer, in parallel.
Revised: We report time as the number of sequential FLOPs needed assuming we can run each example, and in the case of the local methods, each layer, in parallel.

Original: We refer the reader to Appendix A for detailed information on how total and sequential FLOPs are computed for each experiment.
Revised: We refer the reader to Appendix A for detailed information on how total and sequential FLOPs are computed for each experiment.

Original: We compare how backpropagation scales with compute across a variety of local methods: (i) greedy (Figure 2(b)), (ii) overlapping (Figure 2(c)), (iii) two and three chunk greedy (Figure 2(d)), where we split the network into two or three pieces that are trained in a greedy fashion, (iv) last layer & last two layers, a simple baseline where we only backpropagate through the last one or two layers and keep the rest of the network parameters ﬁxed.
Revised: We compare how backpropagation scales with compute across a variety of local methods: (i) greedy (Figure 2(b)), (ii) overlapping (Figure 2(c)), (iii) two and three chunk greedy (Figure 2(d)), where we split the network into two or three pieces that are trained in a greedy fashion, (iv) last layer & last two layers, a simple baseline where we only backpropagate through the last one or two layers and keep the rest of the network parameters ﬁxed.

Original: We apply these methods on a variety of architectures and data including a dense feed-forward network, a ResNet50 network (He et al., 2016) trained on ImageNet (Russakovsky et al., 2015), and a Transformer (Vaswani et al., 2017) model trained on LM1B (Chelba et al., 2013).
Revised: We apply these methods on a variety of architectures and data including a dense feed-forward network, a ResNet50 network (He et al., 2016) trained on ImageNet (Russakovsky et al., 2015), and a Transformer (Vaswani et al., 2017) model trained on LM1B (Chelba et al., 2013).

Original: In Appendix C, we provide results for additional feed-forward networks, a ResNet18 trained on ImageNet, and a larger Transformer, as well as further architecture details.
Revised: In Appendix C, we provide results for additional feed-forward networks, a ResNet18 trained on ImageNet, and a larger Transformer, as well as further architecture details.

Original: For each model and training method, we perform a large sweep over batch size as well as other optimization hyperparameters, and only display the best-performing runs on the Pareto optimal frontier.
Revised: For each model and training method, we perform a large sweep over batch size as well as other optimization hyperparameters, and only display the best-performing runs on the Pareto optimal frontier.

Original: See Appendix B for more detail.
Revised: See Appendix B for more detail.

Original: The resulting ﬁgures all follow the same general structure.
Revised: The resulting ﬁgures all follow the same general structure.

Original: Models train with low total cost when the amount of available compute is large.
Revised: Models train with low total cost when the amount of available compute is large.

Original: By increasing batch size, the amount of compute utilized per parallel process can be reduced efﬁciently until a critical batch size is reached, at which point further increasing the batch size results in diminishing returns in terms of compute efﬁciency, which is similar to results reported for backpropagation in (McCandlish et al., 2018).
Revised: By increasing batch size, the amount of compute utilized per parallel process can be reduced efﬁciently until a critical batch size is reached, at which point further increasing the batch size results in diminishing returns in terms of compute efﬁciency, which is similar to results reported for backpropagation in (McCandlish et al., 2018).

Original: We ﬁnd that, in most cases, local updates signiﬁcantly increase the training speed over deep networks in the highcompute regime, and therefore utilize less total compute than backpropagation.
Revised: We ﬁnd that, in most cases, local updates signiﬁcantly increase the training speed over deep networks in the highcompute regime, and therefore utilize less total compute than backpropagation.

Original: When applicable, we additionally show tables of the best achieved results across all parameters ignoring the time to reach these values.
Revised: When applicable, we additionally show tables of the best achieved results across all parameters ignoring the time to reach these values.

Original: In this setting, we ﬁnd that backpropagation usually achieves the best performance.
Revised: In this setting, we ﬁnd that backpropagation usually achieves the best performance.

Original: This is partially due to the fact that all of these models are trained for a ﬁxed number of examples, and partially due to the fact that backpropagation makes higher use of the capacity of a given model, which we further investigate in Section 5.
Revised: This is partially due to the fact that all of these models are trained for a ﬁxed number of examples, and partially due to the fact that backpropagation makes higher use of the capacity of a given model, which we further investigate in Section 5.

Original: 4. S YNTHETIC : MLP’ S O VER - FITTING TO CIFAR- 
Revised: 4. S YNTHETIC : MLP’ S O VER - FITTING TO CIFAR-

Original: As a proof of concept we ﬁrst demonstrate optimization performance on an eight layer MLP with 4096 hidden units, performing classiﬁcation on the CIFAR-10 dataset (Krizhevsky et al., 2009).
Revised: As a proof of concept we ﬁrst demonstrate optimization performance on an eight layer MLP with 4096 hidden units, performing classiﬁcation on the CIFAR-10 dataset (Krizhevsky et al., 2009).

Original: Hyperparameter and optimization details can be found in Appendix B.1.
Revised: Hyperparameter and optimization details can be found in Appendix B.1.

Original: From the resulting Pareto frontiers shown in Figure 3, we ﬁnd that in no circumstance is backpropagation the best method to use.
Revised: From the resulting Pareto frontiers shown in Figure 3, we ﬁnd that in no circumstance is backpropagation the best method to use.

Original: In the high compute regime, we ﬁnd that local methods enable training up to 10 × faster (e.g. in 0.001 cutoff).
Revised: In the high compute regime, we ﬁnd that local methods enable training up to 10 × faster (e.g. in 0.001 cutoff).

Original: 4. L ANGUAGE M ODELING : T RANSFORMERS ON LM1B
Revised: 4. L ANGUAGE M ODELING : T RANSFORMERS ON LM1B

Original: Next we explore a small (6M parameter) Transformer (Vaswani et al., 2017) trained on
Revised: Next we explore a small (6M parameter) Transformer (Vaswani et al., 2017) trained on

Original: LM1B (Chelba et al., 2013).
Revised: LM1B (Chelba et al., 2013).

Original: We build off of an implementation in Flax Developers (2020).
Revised: We build off of an implementation in Flax Developers (2020).

Original: Hyperparameters and optimization details can be found in Appendix B.2.
Revised: Hyperparameters and optimization details can be found in Appendix B.2.

Original: We ﬁnd that, for the higher cutoffs, many of the local methods vastly outperform backpropagation.
Revised: We ﬁnd that, for the higher cutoffs, many of the local methods vastly outperform backpropagation.

Original: For the lower cuttofs ( ≤ 4 . 0 )
Revised: For the lower cuttofs ( ≤ 4 . 0 )

Original: , we
Revised: , we

Original: [Figure]
Revised: [Figure]

Original: ﬁnd that while backpropagation is more efﬁcient in the high-time regime, local methods train signiﬁcantly faster in the high-compute regime, and can train 2 × faster than backpropagation .
Revised: ﬁnd that while backpropagation is more efﬁcient in the high-time regime, local methods train signiﬁcantly faster in the high-compute regime, and can train 2 × faster than backpropagation .

Original: These local methods do not reach as low of a minimum in the given training time however.
Revised: These local methods do not reach as low of a minimum in the given training time however.

Original: See Figure
Revised: See Figure

Original: [Figure]
Revised: [Figure]

Original: I MAGE C
Revised: I MAGE C

Original: : R ES N ET 50 ON
Revised: : R ES N ET 50 ON

Original: I MAGE N ET
Revised: I MAGE N ET

Original: Next we explore performance of optimization parallelism on a ResNet50 model trained on the ImageNet dataset (Russakovsky et al., 2015) (Figure 5).
Revised: Next we explore performance of optimization parallelism on a ResNet50 model trained on the ImageNet dataset (Russakovsky et al., 2015) (Figure 5).

Original: Hyperparameter and conﬁguration details can be found in Appendix C.1.
Revised: Hyperparameter and conﬁguration details can be found in Appendix C.1.

Original: We ﬁnd, as before, that for many cutoff values local parallelism shows gains over backpropagation in the high-compute regime.
Revised: We ﬁnd, as before, that for many cutoff values local parallelism shows gains over backpropagation in the high-compute regime.

Original: However, at the cutoff of 74% these gains shrink and the local methods are slightly less efﬁcient.
Revised: However, at the cutoff of 74% these gains shrink and the local methods are slightly less efﬁcient.

Original: We hypothesize this is in-part due to increased overﬁtting by the local methods.
Revised: We hypothesize this is in-part due to increased overﬁtting by the local methods.

Original: To see this we can observe that local methods are much more competitive when evaluated on training accuracy.
Revised: To see this we can observe that local methods are much more competitive when evaluated on training accuracy.

Original: This suggests that given more data these local methods will be competitive.
Revised: This suggests that given more data these local methods will be competitive.

Original:  P ROBING B EHAVIOR OF L OCAL U PDATES
Revised:  P ROBING B EHAVIOR OF L OCAL U PDATES

Original: In the previous section we showed that in some cases local parallelism can provide large speedups over backpropagation but suffers in terms of the best achievable performance.
Revised: In the previous section we showed that in some cases local parallelism can provide large speedups over backpropagation but suffers in terms of the best achievable performance.

Original: In this section we explore why and how these methods work, and discuss limitations.
Revised: In this section we explore why and how these methods work, and discuss limitations.

Original: Gradient Angles: Local parallelism does not follow the gradient of the underlying function.
Revised: Gradient Angles: Local parallelism does not follow the gradient of the underlying function.

Original: Instead it computes a local, greedy approximation.
Revised: Instead it computes a local, greedy approximation.

Original: To check the quality of this approximation we measure the angle between the true gradient, and the gradient computed with our greedy method (Figure 6a).
Revised: To check the quality of this approximation we measure the angle between the true gradient, and the gradient computed with our greedy method (Figure 6a).

Original: We ﬁnd positive angles which imply that these directions are still descent directions.
Revised: We ﬁnd positive angles which imply that these directions are still descent directions.

Original: As one moves further away from the end of the network these similarities shrink.
Revised: As one moves further away from the end of the network these similarities shrink.

Original: Larger Block Sizes Improve Generalization: As noted in Huo et al.
Revised: Larger Block Sizes Improve Generalization: As noted in Huo et al.

Original: (2018a;b) and Belilovsky et al.
Revised: (2018a;b) and Belilovsky et al.

Original: (2019b), using chunked local parallelism with more parallel blocks can decrease performance.
Revised: (2019b), using chunked local parallelism with more parallel blocks can decrease performance.

Original: Here we show that practically this reduction in performance seems to stem mainly from a worsening generalization gap, with train and test results shown for various chunk sizes in Figure 6.
Revised: Here we show that practically this reduction in performance seems to stem mainly from a worsening generalization gap, with train and test results shown for various chunk sizes in Figure 6.

Original: A chunk size of nine is simply backprop, and a chunksize of one is fully greedy.
Revised: A chunk size of nine is simply backprop, and a chunksize of one is fully greedy.

Original: Capacity: Ability to Fit Random Labels:
Revised: Capacity: Ability to Fit Random Labels:

Original: Throughout our work we ﬁnd that models trained with local updates don’t make as efﬁcient use of model capacity.
Revised: Throughout our work we ﬁnd that models trained with local updates don’t make as efﬁcient use of model capacity.

Original: This is not necessarily a problem, but represents a tradeoff.
Revised: This is not necessarily a problem, but represents a tradeoff.

Original: Researchers have found that increased model sizes can be used to train faster without leveraging the extra capacity to its fullest (Raffel et al., 2019; Kaplan et al., 2020).
Revised: Researchers have found that increased model sizes can be used to train faster without leveraging the extra capacity to its fullest (Raffel et al., 2019; Kaplan et al., 2020).

Original: Additionally, techniques like distillation can be used to reduce model size (Hinton et al., 2015).
Revised: Additionally, techniques like distillation can be used to reduce model size (Hinton et al., 2015).

Original: We demonstrate this capacity issue by ﬁtting random labels with a ResNet on CIFAR-10, shown in
Revised: We demonstrate this capacity issue by ﬁtting random labels with a ResNet on CIFAR-10, shown in

Original: Figure 6.
Revised: Figure 6.

Original: [Figure]
Revised: [Figure]

Original: Backprop
Revised: Backprop

Original: [Figure]
Revised: [Figure]

Original: Valid 0.
Revised: Valid 0.

Original: Train 0.
Revised: Train 0.

Original: for ResNet50 models trained on ImageNet.
Revised: for ResNet50 models trained on ImageNet.

Original: We show the cost/time to reach a certain cutoff measured on validation accuracy (top) and training accuracy (bottom).
Revised: We show the cost/time to reach a certain cutoff measured on validation accuracy (top) and training accuracy (bottom).

Original: With low cutoffs (50%, 60%, 70%), modest speedups can be obtained on validation performance.
Revised: With low cutoffs (50%, 60%, 70%), modest speedups can be obtained on validation performance.

Original: With higher cutoffs (74%) however backprop is optimal.
Revised: With higher cutoffs (74%) however backprop is optimal.

Original: In the subsequent table, we show the best accuracies reached for each method across all conﬁgurations.
Revised: In the subsequent table, we show the best accuracies reached for each method across all conﬁgurations.

Original: We ﬁnd that the least parallelized method, Two Chunk Greedy, is the only local method competitive with backprop on validation accuracy.
Revised: We ﬁnd that the least parallelized method, Two Chunk Greedy, is the only local method competitive with backprop on validation accuracy.

Original: [Figure]
Revised: [Figure]

Original: demonstrates higher capacity, in that it is able to memorize the dataset better than local greedy backpropagation.
Revised: demonstrates higher capacity, in that it is able to memorize the dataset better than local greedy backpropagation.

Original: Local Methods Learn Different Features: One way to show differences between local and nonlocal methods is to look at the features learned.
Revised: Local Methods Learn Different Features: One way to show differences between local and nonlocal methods is to look at the features learned.

Original: For each method we test we take the best performing model and visualize the ﬁrst layer features.
Revised: For each method we test we take the best performing model and visualize the ﬁrst layer features.

Original: The results are shown in Figure 7.
Revised: The results are shown in Figure 7.

Original: Qualitatively, we see similar ﬁrst layer features from Backprop and Two/Three Chunk local parallelism.
Revised: Qualitatively, we see similar ﬁrst layer features from Backprop and Two/Three Chunk local parallelism.

Original: The more greedy approaches (Overlap, Greedy) yield a different set of features with fewer edge detectors.
Revised: The more greedy approaches (Overlap, Greedy) yield a different set of features with fewer edge detectors.

Original: Finally, when training with only the last layers, the input layer is not updated, and the features are random.
Revised: Finally, when training with only the last layers, the input layer is not updated, and the features are random.

Original: [Figure]
Revised: [Figure]

Original: We ﬁnd the more global methods (Backprop, Two Chunk, Three Chunk) learn similar distributions over features.
Revised: We ﬁnd the more global methods (Backprop, Two Chunk, Three Chunk) learn similar distributions over features.

Original: However more greedy approaches (Overlap and Greedy) learn visually distinct, less edge-like, features.
Revised: However more greedy approaches (Overlap and Greedy) learn visually distinct, less edge-like, features.

Original: Finally the Last 2 ﬁlters are random, because the input layer is never updated.
Revised: Finally the Last 2 ﬁlters are random, because the input layer is never updated.

Original:  R EALIZED P ERFORMANCE G AINS
Revised:  R EALIZED P ERFORMANCE G AINS

Original: Here we show that performance gains of local parallelism can be realized on real hardware, and that they are similar to or better than pipelined backpropagation despite the increased computation needed for auxiliary losses.
Revised: Here we show that performance gains of local parallelism can be realized on real hardware, and that they are similar to or better than pipelined backpropagation despite the increased computation needed for auxiliary losses.

Original: We train ResNet34, ResNet50 and ResNet101 (He et al., 2016) on the ImageNet dataset (Deng et al., 2009), and compare throughput (images per second) between chunked local parallelism and synchronous pipelined backpropagation (Huang et al., 2019).
Revised: We train ResNet34, ResNet50 and ResNet101 (He et al., 2016) on the ImageNet dataset (Deng et al., 2009), and compare throughput (images per second) between chunked local parallelism and synchronous pipelined backpropagation (Huang et al., 2019).

Original: We implement the models in TensorFlow (Abadi et al., 2016) and train them across 4 or 8 Intelligence Processing Units (IPUs – see details in Appendix E).
Revised: We implement the models in TensorFlow (Abadi et al., 2016) and train them across 4 or 8 Intelligence Processing Units (IPUs – see details in Appendix E).

Original: Note that both local and pipeline conﬁgurations are not optimized for overall throughputLM: This feels weird becuase they are optimized for throughput – just only using a single smaller number of chips right? – they do not make use of data parallelism which could be applied identically in both cases.
Revised: Note that both local and pipeline conﬁgurations are not optimized for overall throughput – they do not make use of data parallelism which could be applied identically in both cases.

Original: We use activation recomputation in the case of pipelined backprop (see discussion in Appendix D.3).
Revised: We use activation recomputation in the case of pipelined backprop (see discussion in Appendix D.3).

Original: The results in Table 1 show that chunked local parallelism can achieve similar or greater throughput compared to pipelined backpropagation, for the same local batch size.
Revised: The results in Table 1 show that chunked local parallelism can achieve similar or greater throughput compared to pipelined backpropagation, for the same local batch size.

Original: This provides evidence that local parallelism can enable similar hardware efﬁciency without necessitating an increase of minibatch size.
Revised: This provides evidence that local parallelism can enable similar hardware efﬁciency without necessitating an increase of minibatch size.

Original: It is therefore amenable to a greater level of data parallelism before performance degradation due to a large global batch size.
Revised: It is therefore amenable to a greater level of data parallelism before performance degradation due to a large global batch size.

Original: The difference in throughput between backpropagation and local parallelism with the same local batch size is primarily due to the poor utilisation during the “ramp-up” and “ramp-down” phases of the pipelined backpropagation.
Revised: The difference in throughput between backpropagation and local parallelism with the same local batch size is primarily due to the poor utilisation during the “ramp-up” and “ramp-down” phases of the pipelined backpropagation.

Original: This can be mitigated by running the pipeline in the steady state for more stages (compare rows 4 and 5 of Table 1).
Revised: This can be mitigated by running the pipeline in the steady state for more stages (compare rows 4 and 5 of Table 1).

Original: However, this results in the accumulation of gradients from a larger number of local batches, thus costing a larger effective batch size.
Revised: However, this results in the accumulation of gradients from a larger number of local batches, thus costing a larger effective batch size.

Original: With greedy local parallelism, updates can be applied asynchronously and the pipeline can be run in steady state indeﬁnitely, after an initial ramp-up phase.
Revised: With greedy local parallelism, updates can be applied asynchronously and the pipeline can be run in steady state indeﬁnitely, after an initial ramp-up phase.

Original: Hardware utilization analysis and further discussion can be found in Appendix D.
Revised: Hardware utilization analysis and further discussion can be found in Appendix D.

Original: [Table]
Revised: [Table]

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work we demonstrated that local parallelism is a competitive alternative to backpropagation in the high-compute training regime, and explored design decisions and trade-offs inherent in training with local parallelism.
Revised: In this work we demonstrated that local parallelism is a competitive alternative to backpropagation in the high-compute training regime [Table] and explored design decisions and trade-offs inherent in training with local parallelism.

Original: We summarize some main takeaways from our work:
Revised: We summarize some main takeaways from our work:

Original: • Speed vs. Performance:
Revised: • Speed vs. Performance:

Original: Greedy local parallelism should be used if speed and computeefﬁciency are the primary objectives.
Revised: Greedy local parallelism should be used if speed and computeefﬁciency are the primary objectives.

Original: Chunked local parallelism should be used if performance is the primary objective.
Revised: Chunked local parallelism should be used if performance is the primary objective.

Original: • Gains in High-Compute Regime: Local parallelism can be useful to prolong computeefﬁcient scaling, and therefore faster training, with larger batch sizes once data parallelism begins to saturate.
Revised: • Gains in High-Compute Regime: Local parallelism can be useful to prolong computeefﬁcient scaling, and therefore faster training, with larger batch sizes once data parallelism begins to saturate.

Original: • Comprehensive Analysis: Local parallelism can be applied across multiple modalities (vision, language) and architectures (MLPs, ResNets, Transformers).
Revised: • Comprehensive Analysis: Local parallelism can be applied across multiple modalities (vision, language) and architectures (MLPs, ResNets, Transformers).

Original: We hope that local methods will enable new research into large models.
Revised: We hope that local methods will enable new research into large models.

Original: By lowering communication requirements – particularly latency requirements surrounding synchronization – we believe that local parallelism can be used to scale up and train more massive models in a more distributed fashion
Revised: By lowering communication requirements – particularly latency requirements surrounding synchronization – we believe that local parallelism can be used to scale up and train more massive models in a more distributed fashion

Original: A C ALCULATION OF T OTAL FLOP S AND S EQUENTIAL FLOP S
Revised: A C ALCULATION OF T OTAL FLOP S AND S EQUENTIAL FLOP S

Original: To construct the Pareto curves used in this work we need some estimate of compute time.
Revised: To construct the Pareto curves used in this work we need some estimate of compute time.

Original: Obtaining hardware independent measurements of compute cost and compute time is desirable, but in general impossible, as different hardware makes different trade offs for compute efﬁciency.
Revised: Obtaining hardware independent measurements of compute cost and compute time is desirable, but in general impossible, as different hardware makes different trade offs for compute efﬁciency.

Original: In this work we choose to use a theoretical estimate of compute costs based on ﬂoating point operation (FLOP) counting.
Revised: In this work we choose to use a theoretical estimate of compute costs based on ﬂoating point operation (FLOP) counting.

Original: In all three models, we divide the costs up into three measurements: FLOPs needed for a forward pass through a layer, ﬂops needed for the auxiliary loss computation, and a multiplier to compute the number of ﬂops for a backward pass.
Revised: In all three models, we divide the costs up into three measurements: FLOPs needed for a forward pass through a layer, ﬂops needed for the auxiliary loss computation, and a multiplier to compute the number of ﬂops for a backward pass.

Original: For simplicity, we average the compute costs across layers.
Revised: For simplicity, we average the compute costs across layers.

Original: While this is strictly not feasible in reality with a batch size of one per device, we can come close to approximating it by using more or less parallel hardware per layer.
Revised: While this is strictly not feasible in reality with a batch size of one per device, we can come close to approximating it by using more or less parallel hardware per layer.

Original: This is relatively simple to implement given the minimal communication overhead.
Revised: This is relatively simple to implement given the minimal communication overhead.

Original: We additionally take into account optimizer ﬂops, which we we approximate as ten times the number of parameters, but this results negligible.
Revised: We additionally take into account optimizer ﬂops, which we we approximate as ten times the number of parameters, but this results negligible.

Original: A. C ALCULATIONS PER M ODEL
Revised: A. C ALCULATIONS PER M ODEL

Original: MLP: An MLP is parameterized by the hidden size, N , and the number of layers, L .
Revised: MLP: An MLP is parameterized by the hidden size, N , and the number of layers, L .

Original: The ﬁrst layer’s total ﬂops are from matrix vector multiplication, a bias add of size N , and a ReLU (which we assume costs 1 FLOP per entry).
Revised: The ﬁrst layer’s total ﬂops are from matrix vector multiplication, a bias add of size N , and a ReLU (which we assume costs 1 FLOP per entry).

Original: This yields a total size of (2 ∗ I ∗ N − I )
Revised: This yields a total size of (2 ∗ I ∗ N − I )

Original: + N + 2 ∗ N
Revised: + N + 2 ∗ N

Original: FLOPs, where I is the input size (Hunger, 2005).
Revised: FLOPs, where I is the input size (Hunger, 2005).

Original: The auxiliary classiﬁers consist of a matrix vector multiplication to size 10, a bias add, and a softmax cross entropy loss.
Revised: The auxiliary classiﬁers consist of a matrix vector multiplication to size 10, a bias add, and a softmax cross entropy loss.

Original: We assume the softmax costs 5 ﬂops per estimate leading to a ﬂop estimate of (2 ∗ N ∗ 10 − N ) + 10 + 5 ∗ 10 .
Revised: We assume the softmax costs 5 ﬂops per estimate leading to a ﬂop estimate of (2 ∗ N ∗ 10 − N ) + 10 + 5 ∗ 10 .

Original: For this problem, we approximate the backward multiplier to be 1.5.
Revised: For this problem, we approximate the backward multiplier to be 1.5.

Original: For the MLP model used in the main text (with hidden size N = 4096 and L = 8 layers), the average forward cost per layer is 32514176.0 ﬂops, and the auxiliary loss 77884.0 ﬂops.
Revised: For the MLP model used in the main text (with hidden size N = 4096 and L = 8 layers), the average forward cost per layer is 32514176.0 ﬂops, and the auxiliary loss 77884.0 ﬂops.

Original: For the remaining models, we compute our estimates of these components by ﬁrst using JAX to convert our models to TensorFlow functions, and then leveraging TensorFlow’s tf.compat.v1.proﬁler.proﬁler .
Revised: For the remaining models, we compute our estimates of these components by ﬁrst using JAX to convert our models to TensorFlow functions, and then leveraging TensorFlow’s tf.compat.v1.proﬁler.proﬁler .

Original: ResNet50:
Revised: ResNet50:

Original: This model has L = 17 layers and contains 38711720 parameters.
Revised: This model has L = 17 layers and contains 38711720 parameters.

Original: We ﬁnd that the average forward ﬂop count per example, per layer is 5479411.176470588, the auxiliary loss per layer is 3382457.3529411764, and the backward multiplier is 2.0280375672996596.
Revised: We ﬁnd that the average forward ﬂop count per example, per layer is 5479411.176470588, the auxiliary loss per layer is 3382457.3529411764, and the backward multiplier is 2.0280375672996596.

Original: ResNet18:
Revised: ResNet18:

Original: This model has L = 9 layers, and has 13170792 parameters.
Revised: This model has L = 9 layers, and has 13170792 parameters.

Original: We ﬁnd that the average forward ﬂop count per example, per layer is 1640544.352941176, the auxiliary loss ﬂop count per example per layer is 565900.6470588235, and the backward multiplier is 2.08565879129763.
Revised: We ﬁnd that the average forward ﬂop count per example, per layer is 1640544.352941176, the auxiliary loss ﬂop count per example per layer is 565900.6470588235, and the backward multiplier is 2.08565879129763.

Original: Transformer
Revised: Transformer

Original: small:
Revised: small:

Original: This model has L = 4 layers
Revised: This model has L = 4 layers

Original: We ﬁnd that the average forward cost per example, per layer is 13837446.0, the auxiliary loss is
Revised: We ﬁnd that the average forward cost per example, per layer is 13837446.0, the auxiliary loss is

Original: 0, and the backward multiplier is 1.6581083035860107.
Revised: 0, and the backward multiplier is 1.6581083035860107.

Original: Transformer
Revised: Transformer

Original: large:
Revised: large:

Original: This model has L
Revised: This model has L

Original: =
Revised: =

Original: 6 layers.
Revised: 6 layers.

Original: We ﬁnd that the average forward cost per example, per layer
Revised: We ﬁnd that the average forward cost per example, per layer

Original: is
Revised: is

Original: .0, the auxiliary cost
Revised: .0, the auxiliary cost

Original: is
Revised: is

Original: .0, and the backward multiplier is 1.7526391044859857.
Revised: .0, and the backward multiplier is 1.7526391044859857.

Original: A. C ALCULATIONS PER M ETHOD
Revised: A. C ALCULATIONS PER M ETHOD

Original: In all cases, we ﬁrst obtain the total computation cost in terms of ﬂops and then compute time (or sequential ﬂops) by dividing by the max amount of parallelism (assuming that each example and each layer are run concurrently).
Revised: In all cases, we ﬁrst obtain the total computation cost in terms of ﬂops and then compute time (or sequential ﬂops) by dividing by the max amount of parallelism (assuming that each example and each layer are run concurrently).

Original: As stated before, this is not strictly possible to implement in hardware.
Revised: As stated before, this is not strictly possible to implement in hardware.

Original: In reality, however, we expect more than one example to be used per device in combination with data parallelism and thus appropriate load balancing can be done.
Revised: In reality, however, we expect more than one example to be used per device in combination with data parallelism and thus appropriate load balancing can be done.

Original: All of these calculations are a function of the 4 numbers described above (forward cost, auxiliary cost, backward multiplier and the optimizer cost) in addition to batch size and the number of gradients steps until the target loss is reached.
Revised: All of these calculations are a function of the 4 numbers described above (forward cost, auxiliary cost, backward multiplier and the optimizer cost) in addition to batch size and the number of gradients steps until the target loss is reached.

Original: Backprop: Per step, backprop involves running one forward pass and one backward pass of the entire network plus plus one auxiliary head for the last layer loss computation.
Revised: Backprop: Per step, backprop involves running one forward pass and one backward pass of the entire network plus plus one auxiliary head for the last layer loss computation.

Original: The cost per example is computed as follows:
Revised: The cost per example is computed as follows:

Original: [Equation]
Revised: [Equation]

Original: Greedy: Per step, the greedy method requires running one forward and backward pass for L layers and L auxilary loss computations.
Revised: Greedy: Per step, the greedy method requires running one forward and backward pass for L layers and L auxilary loss computations.

Original: cost
Revised: cost

Original: [Equation]
Revised: [Equation]

Original: backward
Revised: backward

Original: [Equation]
Revised: [Equation]

Original: aux cost ) ∗ layers ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ optimizer cost time =
Revised: aux cost ) ∗ layers ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ optimizer cost time =

Original: cost / ( batch size ∗ layers )
Revised: cost / ( batch size ∗ layers )

Original: Overlapping: Because we are using overlapping chunks of layers, additional compute must be performed.
Revised: Overlapping: Because we are using overlapping chunks of layers, additional compute must be performed.

Original: This method uses one full forward pass though the entire network plus two backward passes for each non terminal layer.
Revised: This method uses one full forward pass though the entire network plus two backward passes for each non terminal layer.

Original: The terminal layer only requires one less layer of computation.
Revised: The terminal layer only requires one less layer of computation.

Original: We additionally need one forward and backward pass of each auxiliary loss.
Revised: We additionally need one forward and backward pass of each auxiliary loss.

Original: An additional average of gradients is required which incurs extra compute per layer.
Revised: An additional average of gradients is required which incurs extra compute per layer.

Original: cost per example =( forward cost + aux cost ) ∗ layers + ( layers − 1) ∗ backward multiplier ∗ (2 ∗ forward cost + aux cost )+ backward multiplier ∗ ( forward cost + aux cost ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ ( optimizer cost + 2 ∗ parameters ) time = cost / ( batch size ∗ layers )
Revised: cost per example =( forward cost + aux cost ) ∗ layers + ( layers − 1) ∗ backward multiplier ∗ (2 ∗ forward cost + aux cost )+ backward multiplier ∗ ( forward cost + aux cost ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ ( optimizer cost + 2 ∗ parameters ) time = cost / ( batch size ∗ layers )

Original: Two/Three chunk:
Revised: Two/Three chunk:

Original: In this, we perform a full forward + backward pass for each layer plus two or three auxiliary losses.
Revised: In this, we perform a full forward + backward pass for each layer plus two or three auxiliary losses.

Original: Lets call the number of chunks K for the equations bellow.
Revised: Lets call the number of chunks K for the equations bellow.

Original: cost per example =
Revised: cost per example =

Original: (1 + backward multiplier )( forward cost + K ∗ aux cost ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ optimizer cost time =
Revised: (1 + backward multiplier )( forward cost + K ∗ aux cost ) cost = cost per step example ∗ steps ∗ batch size + steps ∗ optimizer cost time =

Original: cost / ( batch size ∗ K )
Revised: cost / ( batch size ∗ K )

Original: Last One/Two Layers: These methods require a full forward pass, a single auxilary loss computation and then a backward pass on the last K layers.
Revised: Last One/Two Layers: These methods require a full forward pass, a single auxilary loss computation and then a backward pass on the last K layers.

Original: To calculate time, we assume this last K layers is the smallest atomic chunk that can be run and we divide up the remaining layers accordingly.
Revised: To calculate time, we assume this last K layers is the smallest atomic chunk that can be run and we divide up the remaining layers accordingly.

Original: cost per example =
Revised: cost per example =

Original: ( layers ∗ forward costaux cost )
Revised: ( layers ∗ forward costaux cost )

Original: + backward multiplier ∗ ( K ∗ forward cost + aux cost ) cost =
Revised: + backward multiplier ∗ ( K ∗ forward cost + aux cost ) cost =

Original: [Table]
Revised: [Table]

Original: cost num parallel =
Revised: cost num parallel =

Original: ( layers + K ∗ backward mult ) /
Revised: ( layers + K ∗ backward mult ) /

Original: ( K ∗ (1 + backward mult )) time = cost / ( batch size ∗ num parallel )
Revised: ( K ∗ (1 + backward mult )) time = cost / ( batch size ∗ num parallel )

Original: B H YPERPARAMETER AND C ONFIGURATION D ETAILS FOR E XPERIMENTAL R ESULTS B. MLP ON CIFAR- 
Revised: B H YPERPARAMETER AND C ONFIGURATION D ETAILS FOR E XPERIMENTAL R ESULTS B. MLP ON CIFAR-

Original: We sweep the batch size from 64-524,288 in powers of 2.
Revised: We sweep the batch size from 64-524,288 in powers of 2.

Original: At each batch size, we train models using learning rate tuned Adam (with six values log spaced between 1e-4 and 3e-2) as well as the ﬁrst 50 optimizers taken from opt list to provide a stronger baseline (Metz et al., 2020).
Revised: At each batch size, we train models using learning rate tuned Adam (with six values log spaced between 1e-4 and 3e-2) as well as the ﬁrst 50 optimizers taken from opt list to provide a stronger baseline (Metz et al., 2020).

Original: All models are trained for three million examples on an eight core TPU-V2 using gradient accumulation to control memory usage.
Revised: All models are trained for three million examples on an eight core TPU-V2 using gradient accumulation to control memory usage.

Original: We select a sequence of cut off values (the loss for which we attempt to reach in the shortest time) and plot the Pareto frontier of the different training methodology in Figure 3.
Revised: We select a sequence of cut off values (the loss for which we attempt to reach in the shortest time) and plot the Pareto frontier of the different training methodology in Figure 3.

Original: B. T RANSFORMERS ON LM1B
Revised: B. T RANSFORMERS ON LM1B

Original: Our Transformer has 4 layers, 8 heads per attention layer, 128-dimensional query, key, and value vectors, 256-dimensional hidden layers, and 128-dimensional embeddings.
Revised: Our Transformer has 4 layers, 8 heads per attention layer, 128-dimensional query, key, and value vectors, 256-dimensional hidden layers, and 128-dimensional embeddings.

Original: We train on lengthsequences formed from subword tokenization with a vocabulary size of 8k.
Revised: We train on lengthsequences formed from subword tokenization with a vocabulary size of 8k.

Original: Each Transformer layer is treated as a separate parallelizable component.
Revised: Each Transformer layer is treated as a separate parallelizable component.

Original: Our auxiliary classiﬁers consist of layer norm and a linear projection back to the vocabulary, with a softmax cross entropy loss.
Revised: Our auxiliary classiﬁers consist of layer norm and a linear projection back to the vocabulary, with a softmax cross entropy loss.

Original: We sweep batch sizes in powers of two from 32 to 524,288.
Revised: We sweep batch sizes in powers of two from 32 to 524,288.

Original: At each batch-size we either train Adam with six different learning rates taken evenly spaced on a log scale between 1e-4 and 3e-2 and the ﬁrst 50 optimizers from opt list (Metz et al., 2020).
Revised: At each batch-size we either train Adam with six different learning rates taken evenly spaced on a log scale between 1e-4 and 3e-2 and the ﬁrst 50 optimizers from opt list (Metz et al., 2020).

Original: All models are run until they have processed 50 million sequences, an an 8-core TPU-V2 with gradient accumulation to control memory.
Revised: All models are run until they have processed 50 million sequences, an an 8-core TPU-V2 with gradient accumulation to control memory.

Original: We chose four cutoff values computed on validation loss to show early in training (a value of 5.0 and 4.0), the value chosen by
Revised: We chose four cutoff values computed on validation loss to show early in training (a value of 5.0 and 4.0), the value chosen by

Original: Shallue et al.
Revised: Shallue et al.

Original: (2018) (3.9), and a loss value slightly lower (3.8).
Revised: (2018) (3.9), and a loss value slightly lower (3.8).

Original: Results can be found in Figure 4.
Revised: Results can be found in Figure 4.

Original: C A DDITIONAL P ARETO C URVES E XPERIMENTS
Revised: C A DDITIONAL P ARETO C URVES E XPERIMENTS

Original: We provide additional Pareto curves for different architecture models.
Revised: We provide additional Pareto curves for different architecture models.

Original: C. R ES N ETS ON I MAGE N ET
Revised: C. R ES N ETS ON I MAGE N ET

Original: We build our code off of the Haiku implementation (Hennigan et al., 2020).
Revised: We build our code off of the Haiku implementation (Hennigan et al., 2020).

Original: We break the network up by putting the ﬁrst convolution, and each residual block into a separate parallelizable component For auxiliary losses we apply batch normalization (Ioffe & Szegedy, 2015), then ReLU, then compute mean across the spatial dimensions, and ﬁnally perform a linear projection to the output classes.
Revised: We break the network up by putting the ﬁrst convolution, and each residual block into a separate parallelizable component For auxiliary losses we apply batch normalization (Ioffe & Szegedy, 2015), then ReLU, then compute mean across the spatial dimensions, and ﬁnally perform a linear projection to the output classes.

Original: We sweep batch sizes from 8 to 524,288 in powers of 2.
Revised: We sweep batch sizes from 8 to 524,288 in powers of 2.

Original: For each batch size we randomly sample optimizer hyperparameters for both the SGDM optimizer with a staircase schedule described in
Revised: For each batch size we randomly sample optimizer hyperparameters for both the SGDM optimizer with a staircase schedule described in

Original: Goyal et al.
Revised: Goyal et al.

Original: (2017) and from the ﬁrst 50 conﬁgurations in opt list.
Revised: (2017) and from the ﬁrst 50 conﬁgurations in opt list.

Original: The resulting cost wall time
Revised: The resulting cost wall time

Original: Pareto curves for both validation accuracy and training accuracy are shown in Figure 5.
Revised: Pareto curves for both validation accuracy and training accuracy are shown in Figure 5.

Original: C. MLP S
Revised: C. MLP S

Original: We provide MLP’s trained matching Section 4.1 but using a different number of hidden units.
Revised: We provide MLP’s trained matching Section 4.1 but using a different number of hidden units.

Original: In addition to 4096 units, we show 1024, 256, and 64 units in Figure 8.
Revised: In addition to 4096 units, we show 1024, 256, and 64 units in Figure 8.

Original: We ﬁnd the last 2 layers performs well for larger networks, as there is enough capacity, but is considerably less useful as model size shrinks.
Revised: We ﬁnd the last 2 layers performs well for larger networks, as there is enough capacity, but is considerably less useful as model size shrinks.

Original: C. T RANSFORMER L ARGE
Revised: C. T RANSFORMER L ARGE

Original: In this section we explore a larger transformer than that in Section 4.2.
Revised: In this section we explore a larger transformer than that in Section 4.2.

Original: This transformer matches the default settings of of (Flax Developers, 2020).
Revised: This transformer matches the default settings of of (Flax Developers, 2020).

Original: It has has 6 layers, 8 heads per attention layer, 512-dimensional query, key, and value vectors, 512-dimensional hidden layers, and 512-dimensional embeddings.
Revised: It has has 6 layers, 8 heads per attention layer, 512-dimensional query, key, and value vectors, 512-dimensional hidden layers, and 512-dimensional embeddings.

Original: We train on length 128 sequences formed from subword tokenization with a vocab size of 32k.
Revised: We train on length 128 sequences formed from subword tokenization with a vocab size of 32k.

Original: We show results in Figure 9.
Revised: We show results in Figure 9.

Original: Unlike in the small transformer and due to increased compute costs, we random sample conﬁgurations instead of running all of them.
Revised: Unlike in the small transformer and due to increased compute costs, we random sample conﬁgurations instead of running all of them.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: Backprop
Revised: Backprop

Original: Greedy
Revised: Greedy

Original: Overlapping 2 Chunk Greedy 3 Chunk Greedy
Revised: Overlapping 2 Chunk Greedy 3 Chunk Greedy

Original: Last Layer
Revised: Last Layer

Original: Last 2 Layers 3.
Revised: Last 2 Layers 3.

Original: [Figure]
Revised: [Figure]

Original: Backprop
Revised: Backprop

Original: Greedy
Revised: Greedy

Original: Overlapping 2 Chunk Greedy 3 Chunk Greedy
Revised: Overlapping 2 Chunk Greedy 3 Chunk Greedy

Original: Last Layer
Revised: Last Layer

Original: Last 2 Layers
Revised: Last 2 Layers

Original: Valid 0.
Revised: Valid 0.

Original: Train 0.
Revised: Train 0.

Original: C. R ES N ET
Revised: C. R ES N ET

Original: In addition to the ResNet50 we explored in the main text (Section 4.3) we also explore a ResNet trained with the same protocols.
Revised: In addition to the ResNet50 we explored in the main text (Section 4.3) we also explore a ResNet trained with the same protocols.

Original: We ﬁnd similar results in Figure 10.
Revised: We ﬁnd similar results in Figure 10.

Original: D H ARDWARE U TILIZATION D. E XECUTION T RACES
Revised: D H ARDWARE U TILIZATION D. E XECUTION T RACES

Original: The execution traces shown in Figure 11 illustrate hardware utilization.
Revised: The execution traces shown in Figure 11 illustrate hardware utilization.

Original: A single cycle of pipelined backpropagation over 8 microbatches is shown in Figure 11a.
Revised: A single cycle of pipelined backpropagation over 8 microbatches is shown in Figure 11a.

Original: This comprises a ramp-up phase (blue dashed box), a steady state where all processors are used at each step (e.g. green dashed box), and a ramp-down phase (orange dashed box).
Revised: This comprises a ramp-up phase (blue dashed box), a steady state where all processors are used at each step (e.g. green dashed box), and a ramp-down phase (orange dashed box).

Original: Gradients for all steps are accumulated and applied at
Revised: Gradients for all steps are accumulated and applied at

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: the end of the cycle.
Revised: the end of the cycle.

Original: It is clear that there is signiﬁcant device underutilization in the ramp-up and ramp-down phases.
Revised: It is clear that there is signiﬁcant device underutilization in the ramp-up and ramp-down phases.

Original: Moreover, in the steps which follow the ramp-up the gradient signal has not yet reached the earlier layers of the network, preventing backwards passes from being executed on these processors and causing poor load balancing.
Revised: Moreover, in the steps which follow the ramp-up the gradient signal has not yet reached the earlier layers of the network, preventing backwards passes from being executed on these processors and causing poor load balancing.

Original: A similar effect is present in the steps which precede ramp down, where no further forward passes are run in shallower layers of the network.
Revised: A similar effect is present in the steps which precede ramp down, where no further forward passes are run in shallower layers of the network.

Original: These effects contribute to the higher throughput of local parallelism relative to backpropagation.
Revised: These effects contribute to the higher throughput of local parallelism relative to backpropagation.

Original: During the steps where all processors calculate a forward and backward pass (dashed green box in Figure 11a, enlarged in Figure 11b), the utilization of backpropagation is at its highest.
Revised: During the steps where all processors calculate a forward and backward pass (dashed green box in Figure 11a, enlarged in Figure 11b), the utilization of backpropagation is at its highest.

Original: More of these steps may be executed in each cycle, which reduces the fractional overhead of ramp-up and ramp-down.
Revised: More of these steps may be executed in each cycle, which reduces the fractional overhead of ramp-up and ramp-down.

Original: However, this results in gradients being accumulated over more microbatches, and therefore in the use of a larger overall minibatch size, thus reducing the potential for data-parallel replication of the pipeline.
Revised: However, this results in gradients being accumulated over more microbatches, and therefore in the use of a larger overall minibatch size, thus reducing the potential for data-parallel replication of the pipeline.

Original: During the peak-utilization steps, the operations executed are largely the same as that for chunked local parallelism (Figure 11c), with the exception of the auxiliary classiﬁers necessary for local parallelism.
Revised: During the peak-utilization steps, the operations executed are largely the same as that for chunked local parallelism (Figure 11c), with the exception of the auxiliary classiﬁers necessary for local parallelism.

Original: Note that in the pipelined backpropagation case (Figure 11b)
Revised: Note that in the pipelined backpropagation case (Figure 11b)

Original: the backward pass is executed before the forward pass (with different microbatches), whereas in the chunked local update case the forward pass is executed before the backward pass (for the same minibatch).
Revised: the backward pass is executed before the forward pass (with different microbatches), whereas in the chunked local update case the forward pass is executed before the backward pass (for the same minibatch).

Original: This difference in ordering can be seen in the tall, thin “spike-like” operations which occur at the start of the backwards pass.
Revised: This difference in ordering can be seen in the tall, thin “spike-like” operations which occur at the start of the backwards pass.

Original: In the pipelined backpropagation step they are run at the start of the step, while they are run mid way through the local update step.
Revised: In the pipelined backpropagation step they are run at the start of the step, while they are run mid way through the local update step.

Original: D. I NTER - PROCESSOR C OMMUNICATION
Revised: D. I NTER - PROCESSOR C OMMUNICATION

Original: Our proﬁling of the code found that the total data communicated between processors for local parallel training was half that of pipelined backprop, for all network/batch size conﬁgurations presented in Table 1.
Revised: Our proﬁling of the code found that the total data communicated between processors for local parallel training was half that of pipelined backprop, for all network/batch size conﬁgurations presented in Table 1.

Original: As an example, Table 2 reports the data communicated between IPUs for a single training step of ResNet34.
Revised: As an example, Table 2 reports the data communicated between IPUs for a single training step of ResNet34.

Original: Note this corresponds to a single microbatch in the case of pipelined backpropagation.
Revised: Note this corresponds to a single microbatch in the case of pipelined backpropagation.

Original: There is a clear reduction in the total data communicated with local parallelism relative to backpropagation.
Revised: There is a clear reduction in the total data communicated with local parallelism relative to backpropagation.

Original: The reduction is as expected: IPU 1, processing the ﬁrst layers, receives no data from other IPUs as no gradient signal from later layers are communicated backward.
Revised: The reduction is as expected: IPU 1, processing the ﬁrst layers, receives no data from other IPUs as no gradient signal from later layers are communicated backward.

Original: IPU 4 does not transmit any data backward for the same reason.
Revised: IPU 4 does not transmit any data backward for the same reason.

Original: Overall the total data communicated is 43.2MB for chunked local parallelism and 86.4MB for backpropagation.
Revised: Overall the total data communicated is 43.2MB for chunked local parallelism and 86.4MB for backpropagation.

Original: These results are consistent with the expectation that pipelined local parallelism should result in half as much inter-processor communication as pipelined backpropagation.
Revised: These results are consistent with the expectation that pipelined local parallelism should result in half as much inter-processor communication as pipelined backpropagation.

Original: D. M EMORY C ONSUMPTION
Revised: D. M EMORY C ONSUMPTION

Original: Table 3 contains the memory consumption statistics for different network and training conﬁgurations.
Revised: Table 3 contains the memory consumption statistics for different network and training conﬁgurations.

Original: We can draw a number of conclusions.
Revised: We can draw a number of conclusions.

Original: First, activation recomputation drastically reduces the memory consumption for pipelined backpropagation.
Revised: First, activation recomputation drastically reduces the memory consumption for pipelined backpropagation.

Original: Note that we do not observe a reduction in throughput with recomputation, as the operations to read and write stored activations also take
Revised: Note that we do not observe a reduction in throughput with recomputation, as the operations to read and write stored activations also take

Original: [Table]
Revised: [Table]

Original: signiﬁcant numbers of cycles, when not recomputing 1 .
Revised: signiﬁcant numbers of cycles, when not recomputing 1 .

Original: Thus, we are conﬁdent that the comparison presented in Table 1 remains valid.
Revised: Thus, we are conﬁdent that the comparison presented in Table 1 remains valid.

Original: Even with recomputation, we see that local parallelism generally reduces the average memory consumption
Revised: Even with recomputation, we see that local parallelism generally reduces the average memory consumption

Original: [Table]
Revised: [Table]

Original: and corresponds in all cases to lower or similar max memory.
Revised: and corresponds in all cases to lower or similar max memory.

Original: The difference in memory consumption depends on the local batch size and number of processors.
Revised: The difference in memory consumption depends on the local batch size and number of processors.

Original: While recomputation reduces the number of activations that must be stored for pipelined backprop, we must still store the input to each processor for each “live” microbatch.
Revised: While recomputation reduces the number of activations that must be stored for pipelined backprop, we must still store the input to each processor for each “live” microbatch.

Original: Thus for larger microbatches this overhead increases.
Revised: Thus for larger microbatches this overhead increases.

Original: Further, a larger number of processors mean that there are more live microbatches at any one time, also increasing the memory overhead.
Revised: Further, a larger number of processors mean that there are more live microbatches at any one time, also increasing the memory overhead.

Original: Conversely, local parallelism introduces extra parameters in the auxiliary classiﬁers, which are not needed for backpropagation.
Revised: Conversely, local parallelism introduces extra parameters in the auxiliary classiﬁers, which are not needed for backpropagation.

Original: These observations explain why the memory consumption for pipelined backpropagation increases over that of local parallelism as the local batch size and/or the number of processors grow.
Revised: These observations explain why the memory consumption for pipelined backpropagation increases over that of local parallelism as the local batch size and/or the number of processors grow.

Original: Thus local parallelism can reduce the memory consumption in a high throughput regime.
Revised: Thus local parallelism can reduce the memory consumption in a high throughput regime.

Original: E H ARDWARE B ACKGROUND
Revised: E H ARDWARE B ACKGROUND

Original: The Intelligence Processor Unit (IPU) (Jia et al., 2019) is an accelerator designed for machine intelligence workloads.
Revised: The Intelligence Processor Unit (IPU) (Jia et al., 2019) is an accelerator designed for machine intelligence workloads.

Original: An IPU contains several parallel processing elements, called tiles, each of which has its own local high-speed memory (SRAM) and is able to run a number of parallel threads.
Revised: An IPU contains several parallel processing elements, called tiles, each of which has its own local high-speed memory (SRAM) and is able to run a number of parallel threads.

Original: For example, in the second generation IPU (MK2), each IPU chip has 1472 compute tiles; each tile is equipped with six parallel threads and 600KB of SRAM, equivalent to a total of 8832 parallel threads and 900MB of on chip memory with an aggregate 47.5 TB/s memory bandwidth per chip.
Revised: For example, in the second generation IPU (MK2), each IPU chip has 1472 compute tiles; each tile is equipped with six parallel threads and 600KB of SRAM, equivalent to a total of 8832 parallel threads and 900MB of on chip memory with an aggregate 47.5 TB/s memory bandwidth per chip.

Original: This design beneﬁts from efﬁcient execution of ﬁne-grained operations across the very large number of parallel threads, and allows these threads to access data efﬁciently.
Revised: This design beneﬁts from efﬁcient execution of ﬁne-grained operations across the very large number of parallel threads, and allows these threads to access data efﬁciently.

Original: This is particularly advantageous when the data access patterns are irregular, sparse or incoherent.
Revised: This is particularly advantageous when the data access patterns are irregular, sparse or incoherent.