Title: Accelerated On-Device Forward Neural Network Training with Module-Wise Descending Asynchronism
Venue: NeurIPS 2023
Accepted: yes
Keywords: asynchronous algorithm

one-device learning

forward gradient descent

directional derivative

forward algorithms
Abstract: On-device learning faces memory constraints when optimizing or fine-tuning on edge devices with limited resources. Current techniques for training deep models on edge devices rely heavily on backpropagation. However, its high memory usage calls for a reassessment of its dominance.
In this paper, we propose forward gradient descent (FGD) as a potential solution to overcome the memory capacity limitation in on-device learning. However, FGD's dependencies across layers hinder parallel computation and can lead to inefficient resource utilization.
To mitigate this limitation, we propose AsyncFGD, an asynchronous framework that decouples dependencies, utilizes module-wise stale parameters, and maximizes parallel computation. We demonstrate its convergence to critical points through rigorous theoretical analysis.
Empirical evaluations conducted on NVIDIA's AGX Orin, a popular embedded device,  show that AsyncFGD reduces memory consumption and enhances hardware efficiency, offering a novel approach to on-device learning.

Accelerated On-Device Forward Neural Network Training with Module-Wise Descending Asynchronism

Abstract

On-device learning faces memory constraints when optimizing or fine-tuning on edge devices with limited resources. Current techniques for training deep models on edge devices rely heavily on backpropagation. However, its high memory usage calls for a reassessment of its dominance. In this paper, we propose forward gradient descent (FGD) as a potential solution to overcome the memory capacity limitation in on-device learning. However, FGD's dependencies across layers hinder parallel computation and can lead to inefficient resource utilization. To mitigate this limitation, we propose AsyncFGD, an asynchronous framework that decouples dependencies, utilizes module-wise stale parameters, and maximizes parallel computation. We demonstrate its convergence to critical points through rigorous theoretical analysis. Empirical evaluations conducted on NVIDIA's AGX Orin, a popular embedded device, show that AsyncFGD reduces memory consumption and enhances hardware efficiency, offering a novel approach to on-device learning.

1 Introduction

Deep learning models have increasingly gained attraction in a multitude of applications, showcasing exceptional predictive capabilities. Nevertheless, their rapidly expanding size [19] poses a formidable challenge for resource-limited edge devices, such as mobile phones and embedded systems. These devices are pervasive in our society and continuously generate new data. To attain model customization, user privacy, and low latency, these devices necessitate on-device learning, involving training and fine-tuning models on freshly gathered data [46]. However, the restricted memory capacity of these devices emerges as a significant hindrance. For example, the Raspberry Pi Model A, introduced in 2013, only featured 256 MB of memory [17], while the more recent Raspberry Pi 400, released in 2020, modestly increased this capacity to a mere 4 GB.

Various techniques have been proposed to address this issue, encompassing quantized training, efficient transfer learning, and rematerialization. Quantized training curtails memory consumption by utilizing low-precision network representation [22, 15, 42, 43, 8]. Efficient transfer learning diminishes training costs by updating merely a subset of the model [20, 5, 2, 44]. Lastly, rematerialization

1

2

3

4

1

2

A

A

B

C

B

C

𝑨

𝒕

𝑨

𝒕

𝑨

𝒕ା𝟐

𝑨

𝒕ା𝟐

𝑩

𝒕

𝑩

𝒕

𝑪

𝒕

𝑪

𝒕

𝑩

𝒕ା𝟏

𝑩

𝒕ା𝟏

𝑪

𝒕

𝑪

𝒕

Broadcast

𝑪

𝒕

Host

:

Client:

Find

local

tangent

௧

Worker

JVP

Forward

𝑨

𝒕

𝑨

𝒕

Activation,

JVP value

𝒉𝑨

𝒕

, 𝒐 𝑨

𝒕

2

Time

Stamp

Message

Passing

𝒕

𝒕ା𝟏

𝒕ା𝟐

Client:

𝑨

𝒕

𝑨

𝒕

𝒕

Update with

஺

௧

௧

conserves memory at the expense of time by discarding and recomputing intermediate variables [3, 7, 16, 35].

Although most of these techniques can be applied irrespective of the optimization strategy, they are frequently employed in tandem with backpropagation in deep learning. Due to the reciprocal nature of backpropagation (i.e., activations from the forward pass are preserved for the subsequent backward pass), the training of deep models utilizing backpropagation typically commands a memory footprint that is 3-4 × larger than the number of parameters involved. Consequently, reexamining backpropagation within the context of surpassing the memory capacity barrier is crucial due to its elevated memory consumption.

The recent revival of interest in the forward-forward algorithm proposed by [10], along with other algorithms predicated on forward computation [39, 28, 1, 33], has prompted a reevaluation of backpropagation. The question persists: is it essential to store intermediate variables and pause to propagate gradients? An alternative, forward gradient descent (FGD), leverages the Jacobianvector product (JVP) under automatic differentiation [37, 1, 33] to assess the effect of a stochastic perturbation in conjunction with the forward pass, yielding an unbiased approximation of the true gradient. FGD offers benefits in terms of memory consumption and biological plausibility[31], as it solely employs forward passes and substantially reduces the storage of intermediate variables [31]. Moreover, FGD can be combined with other existing techniques, such as quantized training and efficient transfer learning, potentially further diminishing memory costs.

Nonetheless, the intrinsic locking mechanism within FGD, more specifically, the layer-by-layer sequential computation during the forward pass, poses an impediment to parallel computation across disparate layers as shown in Fig. 1. In response to these challenges, this paper aims at training deep models on edge devices with memory constraints from an optimization perspective, while improving resource utilization by breaking the lock inside the forward pass with provable convergence. Thus, we propose AsyncFGD, an asynchronous version of Forward Gradient Descent with module-wise asynchronism in the forward pass to disentangle its dependencies, allowing simultaneous computation on different workers. The induced module-wise decreasing and bounded staleness in the parameters not only accelerates prediction and training but also provides theoretical guarantees in convergence. We empirically validate our method across multiple architectures and devices, including CPUs, GPUs, and embedded devices. Our results indicate that AsyncFGD reduces memory consumption and improves hardware efficiency, achieving efficient training on edge devices.

2 Related Works

2.1 Forward Gradient with Reinforcement Learning

This research builds upon the concept of forward-mode automatic differentiation (AD) initially introduced by [37]. It has since been applied to learning recurrent neural networks [39] and calculating

Hessian vector products [28]. However, exact gradient computation using forward-mode AD requires the complete and computationally expensive Jacobian matrix. Recently, Baydin et al. [1] and Silver et al. [31] proposed an innovative technique for weight updating based on directional gradients along randomly perturbed directions. These algorithms have connections to both reinforcement learning (RL) and evolution strategies (ES), since the network receives global rewards in each instance. RL and ES have been successfully utilized in specific continuous control and decision-making tasks in neural networks [38, 34, 32, 4]. Clark et al. [4] observed that global credit assignment performs well in vector neural networks with weights between vectorized neuron groups. However, in forward-model AD methods, much like in evolution strategies, the random sampling approach can lead to high variance in the estimated gradient, particularly when optimizing over a large-dimensional parameter space [31].

2.2 Parallel Strategies and Asynchronous Approaches

Various parallel frameworks have been developed to accelerate computations by utilizing multiple workers. These frameworks include data parallelism [21], pipeline parallelism [12, 30, 23, 14], and tensor parallelism [29, 24]. Each approach leverages different dimensions of data - batch dimension for data parallelism, layer dimension for pipeline parallelism, and feature dimension for tensor parallelism - to improve computational efficiency.

However, in edge computing environments, data is often collected and processed frame-by-frame rather than in batches. This is particularly true for scenarios requiring real-time adaptation and streaming data learning. In such contexts, edge models must quickly adapt to newly acquired data in real-time, leaving limited scope for collecting, fetching, and batching historical data. Therefore, algorithms that do not depend on batch size are preferable.The ability to align with this pipeline and train the model on the fly with incoming data is crucial. Therefore, our research specifically focuses on pipeline parallelism.

Pipeline parallelism can be categorized into synchronous and asynchronous pipelines, which divide computations into sequential stages and align well with the data stream on edge devices. However, in edge devices with limited computational power, potential idle workers due to synchronization in synchronous pipelines like [12] are not optimal.

Asynchronous parallelism, which allows non-simultaneous task execution to enhance resource utilization, is particularly relevant to our method. [23, 14] parallelize tasks by executing them from different iterations. However, the additional memory overhead of storing replicate copies to handle staleness in [23] and the use of multiple activations for backpropagation make the algorithm memory inefficient. This memory overhead is mitigated in strategies such as 'rematerialization' [40, 13] and weight estimation [41]. However, these strategies were originally designed for backpropagation, while our focus is on FGD, which has distinct computational characteristics, rendering the existing '1B1F' strategy for work scheduling in [40, 13, 41, 23] unsuitable.

3 Preliminaries

Gradient-based Method. We commence with a succinct introduction to gradient-based methods deployed in neural networks. Envision the training of a L -layer feedforward neural network, where each layer l ∈ 1 , 2 , . . . , L accepts h l -1 as an input, generating an activation h l = F l ( h l -1 ; w l ) with weight w l ∈ R d l . We represent all parameters by w = [ w 1 ⊺ , w 2 ⊺ , . . . , w L ⊺ ] ⊺ ∈ R d and the output of the network by h L = F ( h 0 , w ) , where h 0 symbolizes the input data x . Given a loss function f and targets y , the objective problem is as follows:

Here, we use f ( x ; w ) in subsequent sections for the sake of brevity.

Gradient-based methods are widely used to optimize deep learning problems. At iteration t , we feed data sample x i ( t ) into the network, where i ( t ) signifies the sample's index. As per the principles of stochastic gradient descent (SGD), we define the network parameters as follows:

Here, γ t ∈ R is the stepsize and ∇ f l,x i ( t ) ( w t ) ∈ R d l is the gradient of the loss function with respect to the weights at layer l and data sample x i ( t ) .

The crux of the matter is obtaining ∇ f l,x i ( t ) ( w t ) . Both scientific and industrial sectors lean towards backpropagation underpinned by automatic differentiation (AD). However, FGD could be another alternative, because it can approximate the true gradient without bias by using forward-mode AD, and importantly has low memory consumption as it only preserve the intermediate variable passed from previous layer, while backpropagation need to store intermediate variable in each layer.

Forward Gradient Descent. Let J F l represent the Jacobian matrix of layer l while u w l represent the random perturbation on w l (more precisely, tangent vector around w l ), we can calculate the JVP value o l , in each layer sequentially by

Mathematically, o l is the directional derivative of F l along u l ; intuitively, o l can be interpreted as the influence of the perturbation u w l on function value h l (we set o 0 to be 0 since we don't need to calculate the JVP value with respect to the input data). This process can be arranged within the forward pass such that (3), (4) are computed in the same pass. Also, since we don't need to propagate backward, h l -1 and o l -1 are thrown away right after the computation of h l and o l . Then we can approximate the true gradient ∇ f ( x ; w ) unbiasedly (Lemma 3.1) by ( ∂f ∂h L o L ) u , where u = [ u ⊺ , . . . , u ⊺ ] ⊺ .

Lemma 3.1. Let u w l ∈ R d l , l ∈ { 1 , 2 , . . . , L } be a normally distributed random Gaussian vectors, then the forward gradient computed through Eq. (3) , (4) is an unbiased estimate of the true gradient

More specifically, each layer l receives o L from the last layer and updates its own parameters locally by using E u w l ( u w l · o L ) . We can then rewrite (2) as :

Forward Locking. It is evident from Eq.(4) that the computation in layer l depends on the activation and JVP value h l -1 , o l -1 from layer l -1 . This dependency creates a "lock" that prevents all layers from updating before receiving the output from dependent layers, thus serializing the computation in the forward pass (refer to Figure 1 for illustration).

4 AsyncFGD

In this section, we propose an innovative approach that utilizes module-wise staleness to untether the dependencies inherent in FGD. This method, which we've named AsyncFGD, facilitates the simultaneous execution of tasks originating from disparate iterations. Suppose a L -layer feedforward neural network is divided into K modules, with each module comprising a set of consecutive layers and their respective parameters. Consequently, we have a configuration such that w = [ w G (0) ⊺ , w G (1) ⊺ , · · · , w G ( K -1) ⊺ ] ⊺ , with G ( k ) denoting the layer indices in group k .

Now, let's delve into the details of how AsyncFGD untethers iteration dependencies and expedites the training process.

4.1 Detaching Iteration Dependency

Forward. At the timestamp t , the data sample x i ( t ) is pumped to the network. In contrast to sequential FGD [1], AsyncFGD permits modules to concurrently compute tasks, each originally belonging to a distinct iteration. All modules, with the exception of the last, operate using delayed parameters. We designate the activation and its JVP value originally ascribed to iteration t in Eq.(3), (4) as ˆ o t l , ˆ h t l respectively. Though the superscript t is no longer time-dependent, it maintains its role in indicating sequential order. Consider L k ∈ G ( k ) as the final layer in module k and f k as the activation of this last layer. The computation in module k ∈ 0 , 1 , · · · , K -1 at timestamp t is defined recursively as follows:

Concurrently, each module receives and stores output from its dependent module for future computations.

Update. The update phase in AsyncFGD parallels that in FGD [1]. The final module broadcasts its JVP value, triggering each module to perform local updates to their parameters. To summarize, at timestamp t , we execute the following update:

Staleness. We measure the time delay in Eq.(6), (7) with g ( k ) = K -1 -k , and g ( K -1) = 0 indicates that last module employs up-to-date parameters.

This approach effectively disrupts the "lock-in" characteristic of FGD, facilitating a parallel forward pass. A comparative illustration of the execution details in sequential FGD and AsyncFGD is provided in Figure 1.

4.2 Stochastic AsyncFGD Alogrithm

To better illustrate the working state of AsyncFGD, we make the following definitions:

Unlike FGD, AsyncFGD forwards the JVP value and activations by parameters in different time delays, which can be concluded as f ( x i ( t -2 K +2) ; w t -2 K +2 G (0) ; w t -2 K +3 G (1) ; · · · ; w t -K +1 G ( K -1) ) . A detailed illustration of the AsyncFGD with K = 3 is shown in Appendix G.2. We summarize the proposed algorithm in Algorithm 1 by example of sampling one tangent vector per iteration.

Algorithm 1 AsyncFGD-SGD

̸

1: for t = 0 , 1 , · · · , T -1 do 2: for k = 0 , 1 , · · · , K -1 in parallel do 3: Read ˆ h t -k L k -1 , ˆ o t -k L k -1 from storage if k = 0 4: Compute ˆ h t -k L k , ˆ o t -k L k 5: Send ˆ h t -k L k , ˆ o t -k L k to next worker's storage if k = K -1 6: end for 7: Broadcast ˆ o t -K +1 L K -1 8: for k = 0 , 1 , · · · , K -1 in parallel do 9: Compute ∆ w t -K +1 G ( k ) = ˆ o t -K +1 L K -1 u t -K +1 w G ( k ) 10: Update w t -K +2 G ( k ) = w t -K +1 G ( k ) -γ t -K +1 ∆ w t -K +1 G ( k ) 11: end for 12: end for

̸

Additionally, the procedure in line 5 could overlap with the operations in line 7,9 and 10. We can also apply this approximation algorithm to gradient-based methods like Adam [18] with little modification to the original. Details can be found in Appendix G.1.

Tangent checkpoint. However, some workers, especially those which are closer to the input, store duplicated tangent vectors. To tackle this issue, we use tangent checkpoint,i.e., storing the seed of tangent vectors and reproducing them in the update stage.

Integration with Efficient Transfer Learning Although AsyncFGD offers several advantages over backpropagation, it shares a common challenge with random Zeroth-Order Optimization and Evolution Strategy methods: the variance of the approximation increases with the dimension of random perturbations[25]. Reducing the learning rate can help but may result in slower convergence. However, we observe that deploying models on edge devices typically involves fine-tuning rather

than training from scratch. Our method can flexibly incorporate the principles of efficient transfer learning by introducing a scaling factor α ∈ [0 , 1] to the randomly sampled tangent:

The modified u ′ w l still supports an approximation of the true gradient, with the expectation of the modified estimator being α 2 w l ∇ f ( x ; w l ) . When α w l is set to 0 , the corresponding parameter is "frozen," resulting in no perturbation or updates and a transparent transmission of JVP values. By adjusting α to various values, we can either fix or reduce the influence and learning of specific layers, aligning our approach with the objectives of efficient transfer learning.

4.3 Acceleration of AsyncFGD

When K = 1 the AsyncFGD is reduced to vanilla FGD without any time delay in parameters. When K ≥ 2 , we can distribute the network across multiple workers. Figure 2 shows the computational time of different algorithms in ideal conditions (i.e. the network is evenly distributed and the communication is overlapped by computation and update). T F denotes the time for forward pass and T U denotes the time for updates. It is easy to see that AsyncFGD can fully utilize the computation resources, thus achieving considerable speedup.

5 Convergence Analysis

In this section, we provide the convergence guarantee of Algorithm 1 to critical points in a non-convex setup. We first make the following basic assumptions for nonconvex optimization:

Assumption 5.1. The gradient of f ( w ) is Lipschitz continuous with Lipschitz constant L > 0 , i.e.,

Assumption 5.2. The second moment of stochastic gradient is bounded, i.e., there exist a constant M ≥ 0 , for any sample x i and ∀ w ∈ R d :

Lemma 5.3 ( Mean and Variance ) . Let t ′ = t -K + 1 and diagonal matrices I 0 , · · · , I k , · · · , I K -1 ∈ R d × d such that all the principle diagonal elements of I k in G ( k ) are 1 , and all the principle diagonal elements of I k in other than G ( k ) are 0 . Then we can obtain the mean value and the variance of the forward gradient as follows,

Remark 5.4 . Note that, from modules 0 to K -1 , the corresponding time delays are from K -1 to 0 . Specifically, in timestamp t ′ , when k = K -1 , E u t ′ w G ( K -1) ( o t ′ L · u t ′ w G ( K -1) ) =

∇ f G ( K -1) ,x i ( t ′ ) ( w t ′ ) indicates that the forward gradient in module K -1 is an unbiased estimation of the up-to-date gradient.

Under Assumption 5.1 and 5.2, we obtain the following descent lemma about the objective function value:

Lemma 5.5. Assume that Assumption 5.1 and 5.2 hold. In addition, let t ′ = t -K + 1 , σ := max t ′ γ max { 0 ,t ′ -K +1 } γ t ′ , M K = KM + σK 4 M and choose γ t ′ ≤ 1 L . The iterations in Algorithm 1 satisfy the following descent property in expectation, ∀ t ′ ∈ N :

Method

Training Progress (

=3)

Computation

Time

FGD

Sync

FGD

Async

FGD

F

U

F

F

F

U

F

F

U

U

F

F

F

U

U

U

F

F

F

U

U

U

…

U

F

Forward

Update

Idle

Theorem 5.6. Assume Assumption 5.1 and 5.2 hold and the fixed stepsize sequence { γ t ′ } satisfies γ t ′ = γ ≤ 1 L , ∀ t ′ ∈ { 0 , 1 , . . . , T -1 } . In addition, we assume w ∗ to be the optimal solution to f ( w ) and let t ′ = t -K +1 , σ = 1 such that M K = KM + K 4 M . Then, the output of Algorithm 1 satisfies that:

Theorem 5.7. Assume Assumption 5.1 and 5.2 hold and the diminishing stepsize sequence { γ t ′ } satisfies γ t ′ = γ 0 t ′ +1 ≤ 1 L . In addition, we assume w ∗ to be the optimal solution to f ( w ) and let t ′ = t -K +1 , σ = K such that M K = KM + K 5 M . Let Γ T = ∑ T -1 t ′ =0 γ t ′ , the output of Algorithm 1 satisfies that:

Remark 5.8 . Since the stepsize sequence γ t = γ 0 t +1 satisfies that lim T →∞ ∑ T -1 t =0 γ t = ∞ , lim T →∞ ∑ T -1 t =0 γ 2 t < ∞ , when T → ∞ , the RHS of Eq.(14) will converges to 0 . Let w s be randomly chosen from { w t ′ } T -1 t ′ =0 with probabilities proportional to { γ t ′ } T -1 t ′ =0 , then we have lim s →∞ E ∥∇ f ( w s ) ∥ 2 = 0 .

6 Experiments

This section embarks on a meticulous examination of our proposed method, AsyncFGD. We assess its performance through three distinct facets: memory consumption, acceleration rate, and accuracy. We initiate our analysis by outlining our experimental setup. To validate the efficacy of applying directional derivatives and utilizing module-wise stale parameters, we contrast AsyncFGD with an array of alternative methods encompassing backpropagation, conventional FGD, and other backpropagation-free algorithms. Subsequently, our focus shifts towards scrutinizing the potential of AsyncFGD within the sphere of efficient transfer learning, conducting experiments on prevalent efficient networks. Memory consumption represents another cardinal aspect that we explore, benchmarking AsyncFGD against popular architectures and unit layers like fully-connected layers, RNN cells, and convolutional layers. Concluding our empirical investigation, we assess the speedup of our method relative to other parallel strategies under a diverse set of conditions across multiple platforms.

6.1 Experimental Setup

Methods. We contrast our proposal's memory footprint with Backpropagation, Sublinear [3], Backpropagation through time, and Memory Efficient BPTT [7]. Accuracy-wise, we compare with backpropagation-free methods: Feedback Alignment (FA) [27], Direct Feedback Alignment (DFA) [26], Direct Random Tangent Propagation, and Error-sign-based Direct Feedback Alignment (sDFA) [6]. We also apply parallelization to FGD through FGD-DP (data parallelism) and FGD-MP (model parallelism).

Platform. Experiments utilize Python 3.8 and Pytorch, primarily on nVidia's AGX Orin. Additional results on alternate platforms are in the appendix.

Training. Batch size is 64 unless noted. The optimal learning rate (chosen from { 1 e -5 , 1 e -4 , 1 e -3 , 1 e -2 } with Adam optimizer [18]) is based on validation performance. The parameter α is initially set to 1 for the classifier, with others at 0 for the first 10 epochs. Subsequently, α is gradually increased to 0 . 15 for specific layers. More details are in the appendix.

6.2 Effectiveness of Directional Derivative and Asynchronism

This section documents our experimentation on the effectiveness of using random directional derivative to approximate the true gradient by contrasting it with other BP-free algorithms. Furthermore, we demonstrate that the consistency remains intact when using module-wise stale parameters to uncouple the dependencies, comparing AsyncFGD with vanilla FGD. Results in Table 1 indicate that AsyncFGD can produce results closely aligned with vanilla FGD. Notably, FGD and AsyncFGD yield

the most optimal outcomes when the network is more complex, or when we employ ReLU as the activation function devoid of batchnorm layers (results from ConvS ReLU and FCS ReLU ), situations where FA, DFA and sDFA often fail to propagate an effective error message. Backpropagation results are also furnished solely as a reference to true gradient results. However, when the network grows larger, all BP-free algorithms grapple with variance. The use of larger networks results in only minor improvements compared to BP. We try to address this challenge through efficient transfer learning in the subsequent section.

50

100

150

200

250

Number of Layers

10 1

10 2

10 3

Extra Memory Consumption (MB)

Async

Async

BP

JVP

Sublinear

(a) FC Layers

500

1000

1500

2000

2500

Number of Time Steps

10 3

2×10 2

3×10 2

4×10 2

6×10 2

Async

BPTT

JVP

Sublinear

(b) RNN Cell

1000

2000

3000

4000

5000

Memory Consumption

70

75

80

85

90

95

EfficientNet

MobileNet

ResNet

MnasNet

ShuffleNet

FGD

Async *

BP

(c) Popular Efficient Architectures

6.3 Efficacy of Efficient Transfer Learning

In this segment, we delve into the efficacy of amalgamating AsyncFGD with efficient transfer learning, focusing on popular architectures like ResNet-18(Res-18) [9]; MoblieNet (Mobile)[11]; MnasNet(Mnas)[36] and ShuffleNet(Shuffle)[45] with their lightweight counterpart. The models are fine-tuned with weights pre-trained on Imagenet. AsyncFGD † denotes AsyncFGD utilizing the efficient training strategy. As can be observed from Table 2, the application of an efficient transfer learning strategy brings about substantial performance enhancement, yielding superior results

compared to training with perturbation on the full model. Ablation studies on α provided in Appendix I.1 also shows that, compared to optimizing subset of the model, FGD suffers more from variance.

6.4 Memory Footprint

As illustrated in Fig 3(c), when we plot accuracy against memory consumption, it is evident that AsyncFGD employs approximately one-third of the memory while maintaining close accuracy. Further experimentation on memory consumption with respect to the computing unit reveals, as shown in Fig 3(a) and Fig 3(b), that the additional memory consumption in AsyncFGD mainly serves as placeholders for random tangents, while the JVP computation consumes a negligible amount of additional memory. Memory consumption of other basic units like CNN and batch norm layers, are provided in Appendix.

6.5 Acceleration on Input Stream

In this final section, we assess the acceleration of ResNet-18 with varying K . In this setting, the batch size is set to 4 to better reflect the mechanism of streamlined input on edge device. As demonstrated in 4, AsyncFGD, by detaching dependencies in the forward pass, can outperform other parallel strategies in terms of acceleration rate. While pipeline parallelism is fast, the locking within the forward pass still induces overhead for synchronization, ultimately leading to lower hardware utilization and speed. Results pertaining to different network architectures and other platforms like CPU and GPU as well as more generalized case for larger batch size can be found in the Appendix I.2.

7 Limitations and Discussion

While AsyncFGD offers computational benefits, it is not without limitations. A key drawback is its inferior performance compared to traditional Backpropagation (BP). This performance disparity is largely attributed to the use of randomly sampled directional derivatives in the forward gradient computation, aligning AsyncFGD with Zeroth-Order (ZO) optimization methods and evolutionary strategies. This sampling introduces a significant source of gradient variance, a challenge that is part of a larger problem in the field of stochastic optimization. However, we take encouragement from

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Acceleration

1.54

2.15

2.84

1.03

1.02

1.03

1.11

1.34

1.88

1.23

1.65

2.11

Async

FGD-MP

FGD-DP

Sync

(a) Cluster with 4 GPUs

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Acceleration

1.41

1.95

2.1

1.01

1.02

1.03

1.02

1.02

1.04

1.2

1.55

1.75

Async

FGD-MP

FGD-DP

Sync

(b) Single Embedded Device

recent advancements aimed at reducing this variance, some of which have even facilitated the training of large-scale models [31].

Another constraint pertains to the availability of idle workers on edge devices-a condition that is not universally applicable given the wide variety of edge computing environments. These can span from IoT chips with limited computational resources, where even deploying a standard deep learning model can be problematic, to high-capacity micro-computers used in autonomous vehicles.

Nevertheless, our experimental findings suggest that AsyncFGD is particularly beneficial for specific edge computing scenarios. In such settings, it may serve as a viable alternative for reducing memory usage while fully leveraging available computational resources.

8 Conclusion

In the present paper, we introduce AsyncFGD, an innovative approach designed to shatter the shackles of locking within the forward pass in FGD. This is achieved by incorporating module-wise stale parameters, simultaneously retaining the advantage of minimized memory consumption. In the theoretical segment, we offer a lucid analysis of this partially ordered staleness, demonstrating that our proposed method is capable of converging to critical points even in the face of non-convex problems. We further extend our algorithm to efficient transfer learning by introducing a scale parameter. Our experimental reveals that a sublinear acceleration can be accomplished, without compromising accuracy as well as huge performance gain when utilizing efficient transfer learning strategy. While the exploration of large models employing extensive datasets will undoubtedly continue to rely on backpropagation [10], we assert that the potential of asynchronous algorithms predicated on forward computation should not be overlooked. It offers a promising avenue for fully harnessing limited resources in on-device learning scenarios.

References

Atılım Güne¸ s Baydin, Barak A Pearlmutter, Don Syme, Frank Wood, and Philip Torr. Gradients without backpropagation. arXiv preprint arXiv:2202.08587 , 2022.

Han Cai, Chuang Gan, Ligeng Zhu, and Song Han. Tinytl: Reduce activations, not trainable parameters for efficient on-device learning. arXiv preprint arXiv:2007.11622 , 2020.

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016.

David Clark, LF Abbott, and SueYeon Chung. Credit assignment through broadcasting a global error vector. Advances in Neural Information Processing Systems , 34:10053-10066, 2021.

Yin Cui, Yang Song, Chen Sun, Andrew Howard, and Serge Belongie. Large scale fine-grained categorization and domain-specific transfer learning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 4109-4118, 2018.

Charlotte Frenkel, Martin Lefebvre, and David Bol. Learning without feedback: Fixed random learning signals allow for feedforward training of deep neural networks. Frontiers in neuroscience , 15:629892, 2021.

Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Marc Lanctot, and Alex Graves. Memoryefficient backpropagation through time. Advances in Neural Information Processing Systems , 29, 2016.

Suyog Gupta, Ankur Agrawal, Kailash Gopalakrishnan, and Pritish Narayanan. Deep learning with limited numerical precision. In International conference on machine learning , pages 1737-1746. PMLR, 2015.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.

Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations. arXiv preprint arXiv:2212.13345 , 2022.

Andrew Howard, Mark Sandler, Grace Chu, Liang-Chieh Chen, Bo Chen, Mingxing Tan, Weijun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, Quoc V. Le, and Hartwig Adam. Searching for mobilenetv3, 2019.

Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems , 32, 2019.

Zhouyuan Huo, Bin Gu, and Heng Huang. Training neural networks using features replay. Advances in Neural Information Processing Systems , 31, 2018.

Zhouyuan Huo, Bin Gu, Heng Huang, et al. Decoupled parallel backpropagation with convergence guarantee. In International Conference on Machine Learning , pages 2098-2106. PMLR, 2018.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-arithmetic-only inference. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2704-2713, 2018.

Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, and Joseph E. Gonzalez. Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2020.

Steven J Johnston and Simon J Cox. The raspberry pi: A technology disrupter, and the enabler of dreams, 2017.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Skanda Koppula, Lois Orosa, A Giray Ya˘ glıkçı, Roknoddin Azizi, Taha Shahroodi, Konstantinos Kanellopoulos, and Onur Mutlu. Eden: Enabling energy-efficient, high-performance deep neural network inference using approximate dram. In Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture , pages 166-181, 2019.

Simon Kornblith, Jonathon Shlens, and Quoc V Le. Do better imagenet models transfer better? In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 2661-2671, 2019.

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems , 25, 2012.

Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Goldstein. Training quantized nets: A deeper understanding. Advances in Neural Information Processing Systems , 30, 2017.

Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. Pipedream: generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles , pages 1-15, 2019.

Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-lm. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-15, 2021.

Yurii Nesterov and Vladimir Spokoiny. Random gradient-free minimization of convex functions. Foundations of Computational Mathematics , 17(2):527-566, 2017.

Arild Nøkland. Direct feedback alignment provides learning in deep neural networks. Advances in neural information processing systems , 29, 2016.

Paul Orsmond and Stephen Merry. Feedback alignment: effective and ineffective links between tutors' and students' understanding of coursework feedback. Assessment & Evaluation in Higher Education , 36(2):125-136, 2011.

Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation , 6(1):147160, 1994.

Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis , pages 1-16. IEEE, 2020.

Benjamin Recht, Christopher Re, Stephen Wright, and Feng Niu. Hogwild!: A lock-free approach to parallelizing stochastic gradient descent. Advances in neural information processing systems , 24, 2011.

Mengye Ren, Simon Kornblith, Renjie Liao, and Geoffrey Hinton. Scaling forward gradient with local losses. arXiv preprint arXiv:2210.03310 , 2022.

Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as a scalable alternative to reinforcement learning. arXiv preprint arXiv:1703.03864 , 2017.

David Silver, Anirudh Goyal, Ivo Danihelka, Matteo Hessel, and Hado van Hasselt. Learning by directional gradient descent. In International Conference on Learning Representations , 2022.

Kenneth O Stanley and Risto Miikkulainen. Evolving neural networks through augmenting topologies. Evolutionary computation , 10(2):99-127, 2002.

Benoit Steiner, Mostafa Elhoushi, Jacob Kahn, and James Hegarty. Olla: Optimizing the lifetime and location of arrays to reduce the memory usage of neural networks, 2022.

Mingxing Tan, Bo Chen, Ruoming Pang, Vijay Vasudevan, Mark Sandler, Andrew Howard, and Quoc V. Le. Mnasnet: Platform-aware neural architecture search for mobile, 2019.

Robert Edwin Wengert. A simple automatic derivative evaluation program. Communications of the ACM , 7(8):463-464, 1964.

Darrell Whitley, Stephen Dominic, Rajarshi Das, and Charles W Anderson. Genetic reinforcement learning for neurocontrol problems. Machine Learning , 13:259-284, 1993.

Ronald J Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation , 1(2):270-280, 1989.

AnXu, Zhouyuan Huo, and Heng Huang. On the acceleration of deep learning model parallelism with staleness. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2088-2097, 2020.

Bowen Yang, Jian Zhang, Jonathan Li, Christopher Ré, Christopher Aberger, and Christopher De Sa. Pipemare: Asynchronous pipeline parallel dnn training. Proceedings of Machine Learning and Systems , 3:269-296, 2021.

Guandao Yang, Tianyi Zhang, Polina Kirichenko, Junwen Bai, Andrew Gordon Wilson, and Chris De Sa. Swalp: Stochastic weight averaging in low precision training. In International Conference on Machine Learning , pages 7015-7024. PMLR, 2019.

Yukuan Yang, Lei Deng, Shuang Wu, Tianyi Yan, Yuan Xie, and Guoqi Li. Training highperformance and large-scale deep neural networks with full 8-bit integers. Neural Networks , 125:70-82, 2020.

Elad Ben Zaken, Shauli Ravfogel, and Yoav Goldberg. Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models. arXiv preprint arXiv:2106.10199 , 2021.

Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient convolutional neural network for mobile devices, 2017.

Qihua Zhou, Zhihao Qu, Song Guo, Boyuan Luo, Jingcai Guo, Zhenda Xu, and Rajendra Akerkar. On-device learning systems for edge intelligence: A software and hardware synergy perspective. IEEE Internet of Things Journal , 8(15):11916-11934, 2021.

A Appendix

The organization of the Appendix is as follows: Appendix sections B to F provide detailed proofs of the lemmas and theorems presented in the main text. This is followed by additional insights, including operational specifics of working with Adam and details on how AsyncFGD can be extended for parallel processing on recursive networks, covered in Appendix G. Lastly, Appendix H and I provide a comprehensive view of the training process and additional experimental results, respectively.

B Proof of Lemma 3.1

According to the update rule of Eq.(3), (4), we have

Then for any l ∈ { 1 , 2 , . . . , L } , take expectation with respect to u w l , we have

̸

Note that

where

̸

with m ∈ { 1 , 2 , . . . , d h L } , n ∈ { 1 , 2 , . . . , d l } . Since each u w l ∼ N (0 , I ) , we have

̸

Similarly, we can prove that E u w l [ ∑ k = l ( J F L ( w k ) u w k ) · u w l ] = 0 ∈ R d h L × d l . So we have,

C Proof of Lemma 5.3

Proof of Mean .

· · · · · ·

· · · · · ·

· · · · · ·

Take expectation with respect to u t -K +1 w l , where l ∈ G ( k ) , we have

So we have,

Lemma C.1 ([25], Theorem 3) . Let g u ( x ) = ⟨∇ f ( x ) , u ⟩ u , where u ∈ R d is a normally distributed Gaussian vector, then we have

Lemma C.2. Let g u 1 ,u 2 ( x ) = ⟨∇ f ( x ) , u 1 ⟩ u 2 , where u 1 ∈ R d 1 , u 2 ∈ R d 2 are two i.i.d. normally distributed Gaussian vectors, then we have

Proof.

̸

̸

̸

̸

where the first inequality is due to Lemma 1 in [25].

̸

where the inequality is due to Lemma C.1 and C.2.

D proof of Lemma 5.6

Proof. For the convenience of analysis, we denote t ′ = t -K +1 , then the update rule of algorithm 1 can be rewritten as

Take expectation with respect to u t ′ w G ( k ) , we have

Define diagonal matrices I 0 , · · · , I k , · · · , I K -1 ∈ R d × d such that all the principle diagonal elements of I k in G ( k ) are 1 , and all the principle diagonal elements of I k in other than G ( k ) are 0 . Then we have

Since f ( · ) is L -smooth, the following inequality holds that:

From the update rule of Algorithm 1, we take expectation with respect to all random variables on both sides and obtain:

Using the fact that ∥ x + y ∥ 2 ≤ 2 ∥ x ∥ 2 +2 ∥ y ∥ 2 and xy ≤ 1 2 ∥ x ∥ 2 + 1 2 ∥ y ∥ 2 , we have

+ Lγ 2 t ′ ∥ ∥ ∥ ∥ ∥ K -1 ∑ k =0 I k ∇ f G ( k ) ( w t ′ -K + k +1 ) -∇ f ( w t ′ ) ∥ ∥ ∥ ∥ ∥ ︸ ︷︷ ︸ Q 4 Q 2 = -( γ t ′ -Lγ 2 t ′ ) K -1 ∑ k =0 ∇ f ( w t ′ ) ⊺ I k ( ∇ f G ( k ) ( w t ′ -K + k +1 ) -∇ f G ( k ) ( w t ′ )) ≤ γ t ′ -Lγ 2 t ′ 2 ∥ ∥ ∥ ∇ f ( w t ′ ) ∥ ∥ ∥ 2 + γ t ′ -Lγ 2 t ′ 2 ∥ ∥ ∥ ∥ ∥ K -1 ∑ k =0 I k ∇ G ( k ) f ( w t ′ -K + k +1 ) -∇ f ( w t ′ ) ∥ ∥ ∥ ∥ ∥ 2 Using E ∥ ξ -E [ ξ ] ∥ 2 ≤ E ∥ ξ ∥ 2 , we have Q 3 = E ∥ ∥ ∥ ∥ ∥ K -1 ∑ k =0 I k · ˆ o t ′ L K -1 · u t ′ w G ( k ) -K -1 ∑ k =0 I k ∇ f G ( k ) ( w t ′ -K + k +1 ) ∥ ∥ ∥ ∥ ∥ 2 ≤ E ∥ ∥ ∥ ∥ ∥ K -1 ∑ k =0 I k · ˆ o t ′ L K -1 · u t ′ w G ( k ) ∥ ∥ ∥ ∥ ∥ 2 ≤ ( d +4) ∥ ∥ ∥ ∥ ∥ K -1 ∑ k =0 I k ∇ f G ( k ) ,x i ( t ′ ) ( w t ′ -K + k +1 ) ∥ ∥ ∥ ∥ ∥ 2 =( d +4) K -1 ∑ k =0 ∥ ∥ ∥ ∇ f G ( k ) ,x i ( t ′ ) ( w t ′ -K + k +1 ) ∥ ∥ ∥ 2 ≤ ( d +4) KM, where the second inequality is due to Lemma 5.3. Then we bound Q 4 ,

where the second inequality is from Assumption 5.1, the third inequality is due to Lemma 5.3, the fourth inequality follows from Lγ t ′ < 1 , the last inequality follows from the inequality of arithmetic and geometric means, Assumption 5.2 and σ := max t ′ γ max { 0 ,t ′ -K +1 } γ t ′ . Integrating the upper bound together, we have

where we let M K = KM + σK 4 M .

E Proof of Theorem 5.6

Proof. Let γ t = γ be a constant, taking total expectation in Lemma 5.5, we have

where σ = 1 and M k = KM + K 4 M . summing the above inequality from t ′ = 0 to T -1 we have

Then we have

F Proof of Theorem 5.7

Proof. Let { γ t ′ } be a diminishing sequence and γ t ′ = γ 0 t ′ +1 , such that σ < K and M K = KM + K 5 M . Taking expectation in Lemma 5.5 and summing it from t ′ = 0 to T -1 , we have

Letting Γ T = ∑ T -1 t ′ =0 γ t ′ , then we have

G Details of AsyncFGD

G.1 Working with Adam

We provide example for AsyncFGD working with Adam in Algorithm 2. Minimal changes are made on Adam by substituting the gradient the estimator using Forward Gradient.

G.2 Execution Details

Details are presented in Figure 5. By pipelining over time dimension, we can preserve buffer for input in only one timestamp and still achieve parallel computation.

Algorithm 2 AsyncFGD-Adam

0

, v

G

(

k

)

= 0

, β

1

= 0

.

9

, β

2

= 0

.

999

, η

= 1

1: for t = 0 , 1 , · · · , T -1 do

2: for k = 0 , 1 , · · · , K -1 in parallel do

̸

3: Read ˆ h t -k L k -1 , ˆ o t -k L k -1 from storage if k = 0

4: Compute ˆ h t -k L k , ˆ o t -k L k

̸

5: Send ˆ h t -k L k , ˆ o t -k L k to next worker's storage if k = K -1

7: Broadcast ˆ o t -K +1 L K -1

6: end for

8: for k = 0 , 1 , · · · , K -1 in parallel do

15: end for

16: end for

𝑠𝑠

𝐿𝐿1

0

,

�

𝑜𝑜

𝐿𝐿1

0

Initialize

Module

2

Module

3

𝑤𝑤

𝐿𝐿1

1

𝑡𝑡

= 1

𝑠𝑠

𝐿𝐿1

1

,

�

𝑜𝑜

𝐿𝐿1

1

𝑤𝑤

𝐿𝐿1

1

𝑡𝑡

= 2

𝑠𝑠

𝐿𝐿1

2

,

�

𝑜𝑜

𝐿𝐿1

2

𝑤𝑤

𝐿𝐿1

1

𝑡𝑡

= 3

𝑠𝑠

𝐿𝐿1

3

,

�

𝑜𝑜

𝐿𝐿1

3

𝑤𝑤

𝐿𝐿1

2

𝑠𝑠

𝐿𝐿1

5

,

�

𝑜𝑜

𝐿𝐿1

5

𝑤𝑤

𝐿𝐿1

5

0, 0

𝑤𝑤

𝐿𝐿2

1

0, 0

𝑤𝑤

𝐿𝐿3

1

0, 0

𝑤𝑤

𝐿𝐿3

1

𝑠𝑠

𝐿𝐿2

0

,

�

𝑜𝑜

𝐿𝐿2

0

𝑤𝑤

𝐿𝐿2

1

𝑠𝑠

𝐿𝐿3

0

,

�

𝑜𝑜

𝐿𝐿3

0

𝑤𝑤

𝐿𝐿3

1

𝑠𝑠

𝐿𝐿2

1

,

�

𝑜𝑜

𝐿𝐿2

1

𝑤𝑤

𝐿𝐿2

1

𝑠𝑠

𝐿𝐿3

1

,

�

𝑜𝑜

𝐿𝐿3

1

𝑤𝑤

𝐿𝐿3

2

𝑠𝑠

𝐿𝐿2

2

,

�

𝑜𝑜

𝐿𝐿2

2

𝑤𝑤

𝐿𝐿2

2

𝑠𝑠

𝐿𝐿3

3

,

�

𝑜𝑜

𝐿𝐿3

3

𝑠𝑠

𝐿𝐿2

4

,

�

𝑜𝑜

𝐿𝐿2

4

𝑤𝑤

𝐿𝐿3

5

𝑤𝑤

𝐿𝐿2

5

Module

1

warm-up

steady

Forward

Skip

𝑡𝑡

= 6

𝑠𝑠

𝐿𝐿1

4

,

�

𝑜𝑜

𝐿𝐿1

4

𝑡𝑡

= 5

𝑤𝑤

𝐿𝐿1

4

𝑠𝑠

𝐿𝐿3

2

,

�

𝑜𝑜

𝐿𝐿3

2

𝑠𝑠

𝐿𝐿2

3

,

�

𝑜𝑜

𝐿𝐿2

3

𝑤𝑤

𝐿𝐿3

4

𝑤𝑤

𝐿𝐿2

4

𝑤𝑤

𝐿𝐿1

3

𝑤𝑤

𝐿𝐿3

3

𝑤𝑤

𝐿𝐿2

3

𝑡𝑡

= 4

𝑠𝑠

𝐿𝐿1

6

,

�

𝑜𝑜

𝐿𝐿1

6

𝑤𝑤

𝐿𝐿1

6

𝑠𝑠

𝐿𝐿3

4

,

�

𝑜𝑜

𝐿𝐿3

4

𝑠𝑠

𝐿𝐿2

5

,

�

𝑜𝑜

𝐿𝐿2

5

𝑤𝑤

𝐿𝐿2

6

𝑤𝑤

𝐿𝐿2

6

𝑡𝑡

= 7

𝑠𝑠

𝐿𝐿1

7

,

�

𝑜𝑜

𝐿𝐿1

7

𝑤𝑤

𝐿𝐿1

7

𝑠𝑠

𝐿𝐿3

5

,

�

𝑜𝑜

𝐿𝐿3

5

𝑠𝑠

𝐿𝐿2

6

,

�

𝑜𝑜

𝐿𝐿2

6

𝑤𝑤

𝐿𝐿2

7

𝑤𝑤

𝐿𝐿2

7

𝑡𝑡

= 8

𝑖𝑖

= 0

𝑖𝑖

= 1

Update

̂

̂

̂

̂

̂

̂

G.3 Extension: AsyncFGD-Recursive

In this section, we extend the potential of AsyncFGD by exploring the parallelization of sequential inputs in RNNs with reduced memory footprint, necessitating the preservation of input for only a single timestamp.

We adopt a one-to-many RNN network for ease of illustration and denote the equal length of each sequence as n . We begin by refactoring the original loss for RNNs in terms of cumulative loss and new activation. Here, s t l signifies the hidden state at timestamp t on layer l . At timestamp t , each layer ingests ( s t l -1 , s t -1 l ) as input, generating s t l = F l ( s t l -1 , s t -1 l , w l ) . We represent the stacked latent states passed from t -1 as s t -1 = [ s t -1 1 , s t -1 2 , . . . , s t -1 L ] and the output as y t = F ( s t 0 , s t -1 ; w ) , where s t 0 symbolizes the input data x t . The cumulative loss from timestamp 1 ∼ T is given by:

We next refactor equation 2 for the i th sequential input in iteration i, i ≥ 0 as:

e

-

8

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

̂

where L i := ∑ t = in +1 ( i +1) n f ( F ( x t , s t -1 ; w ) , y t ) represents the loss for the i th sequence. We break the dependency between timestamps and iterations by employing dynamic staleness in AsynFGD. Specifically, the computation in module k ∈ 1 , 2 , · · · , K at timestamp t is defined as follows:

Given that tasks belonging to the same iteration use identical parameters, we use δ ( k, t, i ) = t -ni -k -1 , t ∈ [ in +1 , ( i +1) n ] to quantify this difference for the i th sequential. If δ ( k, t, i ) ≤ 0 , then module k uses stale parameters from iteration i -1 at timestamp t . AsyncFGD-RNN only updates the parameter upon the completion of the last computation in the sequence. Specifically, we use:

Refer to figure 5 for detailed execution instructions. By combining training and prediction, we can process data from different timestamps of sequential input, maintain a buffer for just a single timestamp, and still achieve parallelization among various workers.

H Training Details

In this section, we explain some details in the trainning process.

H.1 Random Seed

The seeds for all experiments are fixed to 0.

H.2 Description of Network Architecture

H.2.1 Models Deployed in Section 6.2.

The network structures of ConvS,ConvL,FCS and FCL are enumerated in Tables 8, 11, 9, and 10 correspondingly. In these designations, the subscripts ReLU and Tanh signify the particular activation function used in the model. Moreover, the larger models, denoted as counterparts, incorporate batch normalization layers for enhanced performance.

H.2.2 Models used in Section 6.3.

We delve into the specific models employed in Section 6.3. For MobileNet, we utilize the structure of MobileNet_V3_Small. The ResNet-18 structure is used when implementing ResNet. The model EfficientNet_B0 signifies the EfficientNet architecture. The MNASNet0_5 is used for MNasNet. Lastly, we adopt ShuffleNet_V2_X0_5 for ShuffleNet.

H.3 Model Splitting

In this section, we provide details for how to split the model into consective moduels and distribute them in different worker, we will first provide how to split models in Section 6.2, then we provide how to split model in 6.3.

H.3.1 Model Splitting in Section 6.2

In Section 6.2, all models are split with K = 3 . Details for how to split the ConvS, ConvL, FCS, FCL are repented in Table 5 Table 6, Table3 and Table 4, respectively.

H.3.2 Model Splitting in Section 6.3

In section 6.3, all models are divided into four parts ( K = 4 ). Detailed descriptions of how each model is split are provided below. Note that 'head' and 'tail' refer to the layers before and after the main blocks of each architecture, respectively, which are assigned to the first and the last worker:

ResNet-18: The core of ResNet-18 consists of 4 Residual Blocks, each distributed to one of the four workers.

EfficientNet: The core of EfficientNet consists of 7 Mobile Inverted Bottlenecks (MBConv). The first worker handles MBConv 1, the second handles MBConv 2 to 3, the third manages MBConv 4 to 6, and the last one manages MBConv 7.

MoblieNet: The core of MoblieNetV3-small includes 13 layers of bottlenecks. The first worker handles layers 1 to 3, the second manages layers 4 to 7, the third manages layers 8 to 11, and the last worker handles layers 12 to 13.

MnasNet: The core of MnasNet consists of 6 blocks of inverted residuals. Blocks 1 to 2, 3 to 5, and 6 are assigned to workers 2, 3, and 4 respectively, while the first worker only handles the head.

ShuffleNet: The core of ShuffleNet consists of 3 stages, each assigned to workers 2, 3, and 4, respectively.

I Additional Experimental Results

I.1 Ablation study in α

We have incremented the value of α bias gradually, with a step size of 0.0075, over 20 epochs. This process can be generalized using the following equations:

Here, α ∗ bias control the rate of increase and the maximum attainable value of α bias , respectively. The ablation study with respect to α ∗ bias is presented in Table 7.

We observe that reducing α ∗ bias to 0, which corresponds to only updating the classifier, still results in performance gains compared to updating the full model. This improvement can be attributed to reduced variance. As α ∗ bias increases, we generally see better results, since the norm of the gradient approximation increases. However, when α ∗ bias exceeds 0.25, we sometimes observe a performance drop due to the corresponding increase in variance.

I.2 Acceleration across Various Platforms and Architectures

In Section 6.5, we examined the acceleration of AsyncFGD in comparison to vanilla FGD on ResNet18, using two hardware platforms: 1) NVIDIA AGX Orin, an embedded device, and 2) a cluster of four NVIDIA 1080 Ti GPUs. These platforms were chosen to reflect real-world edge device scenarios and to simulate situations of ample computational power, such as in the case of stacked chips.

In this section, we expand our scope of investigation by incorporating two additional devices: 1) NVIDIA A100, and 2) Intel(R) Xeon(R) CPU E5-2678 v3 @2.50GHZ. These additions allow us to further examine acceleration under various conditions. We also provide supplementary results on acceleration with respect to different batch sizes to reflect variable input streams. Moreover, to emulate streamlined input, the mini-batch size of the synchronized pipeline is set to 1.

The performance of ResNet-18 with different batch sizes on the four NVIDIA 1080Ti GPUs, A100, and AGX Orin platforms is illustrated in Figures 6, 7, and 8, respectively. Results for MobileNetV3small on AGX Orin are presented in Figure 10. A notable observation is that AsyncFGD performance appears largely insensitive to batch size. In contrast, other algorithms typically exhibit poorer performance with smaller batch sizes. Particularly, when the batch size is reduced to 1, these algorithms offer negligible performance improvements over vanilla FGD. Furthermore, the overall

acceleration on a single device is constrained by computational power. For instance, while AsyncFGD achieves a speedup of 2.84 × on a four GPU cluster, it only delivers a 2.11 × speedup on a single AGX Orin. Communication also imposes a limit on the overall acceleration, as demonstrated by the superior performance on the A100 in comparison to the four-GPU cluster. This is attributable to the elimination of communication overhead on a single device, except for the sending and receiving operations of CUDA kernels.

Results for MobileNetV3-small with different batch sizes on CPU are depicted in Figure 9. Due to the inherently sequential execution pattern of CPUs, the acceleration is constrained, resulting in only modest speedup and advantage over other algorithms.

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.53

2.12

2.81

0.95

0.95

0.96

0.98

0.98

1.01

0.93

0.92

0.94

Async

FGD-DP

FGD-MP

Sync

(a) Batch size = 1

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.52

2.17

2.8

1.1

1.14

1.21

1.02

1.01

1.01

1.15

1.27

1.2

Async

FGD-DP

FGD-MP

Sync

(b) Batch size = 2

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.56

2.11

2.83

1.15

1.33

1.33

1.01

0.99

0.99

1.2

1.53

1.53

Async

FGD-DP

FGD-MP

Sync

(c) Batch size = 3

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.54

2.15

2.84

1.11

1.34

1.88

1.01

1.02

1.03

1.23

1.65

2.11

Async

FGD-DP

FGD-MP

Sync

(d) Batch size = 4

I.3 Memory Profiling on Other Basic Units of Convolutional Neural Networks

This section outlines memory profiling for basic units within a Convolutional Neural Network (CNN). Commonly, a CNN layer is coupled with a batch normalization layer and an activation layer using ReLU, so we've combined these elements for our memory testing. We examine the memory consumption against the number of layers and present the results in Figure 11(a). For further examination, we also assess the memory consumption against the number of output channels and batch size, with results shown in Figures 11(b) and 11(c), respectively.

Our findings reveal that implementing forward gradients can significantly reduce memory consumption. Generally, the majority of memory usage in CNNs stems from intermediate results, since CNNs often operate in a 'broadcast then product' pattern (to be specific, they are referred as 'img2col'). Consequently, the additional memory required by the random tangent in AsyncFGD is minimal. As such, the memory consumption appears to be invariant to the number of layers, mainly because in the forward pass we discard almost all the intermediate variables.

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.6

2.35

3.13

1.02

0.98

1.02

0.98

0.98

0.96

0.93

0.92

0.94

Async

FGD-DP

FGD-MP

Sync

(a) Batch size = 1

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.61

2.34

3.13

1.01

1.02

1.01

1.02

1.01

1.0

1.2

1.12

1.53

Async

FGD-DP

FGD-MP

Sync

(b) Batch size = 2

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.63

2.41

3.15

1.0

0.99

0.98

1.01

0.99

1.01

1.37

1.67

1.88

Async

FGD-DP

FGD-MP

Sync

(c) Batch size = 3

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.65

2.4

3.26

1.02

1.01

1.0

1.0

1.0

1.04

1.45

1.87

2.35

Async

FGD-DP

FGD-MP

Sync

(d) Batch size = 4

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.38

1.87

2.2

1.02

0.98

1.02

0.98

0.98

0.96

0.93

0.92

0.94

Async

FGD-DP

FGD-MP

Sync

(a) Batch size = 1

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.4

1.91

2.27

1.01

1.02

1.01

1.02

1.01

1.0

1.12

1.12

1.23

Async

FGD-DP

FGD-MP

Sync

(b) Batch size = 2

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.4

1.94

2.31

1.0

0.99

0.98

1.01

0.99

1.01

1.18

1.44

1.65

Async

FGD-DP

FGD-MP

Sync

(c) Batch size = 3

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.41

1.95

2.31

1.02

1.01

1.0

1.0

1.0

1.04

1.2

1.55

1.75

Async

FGD-DP

FGD-MP

Sync

(d) Batch size = 4

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.2

1.59

1.9

0.97

0.94

0.97

0.94

0.94

0.92

0.7

0.69

0.7

Async

FGD-DP

FGD-MP

Sync

(a) Batch size = 1

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.21

1.58

1.9

0.96

0.97

0.96

0.97

0.96

0.95

0.9

0.84

1.15

Async

FGD-DP

FGD-MP

Sync

(b) Batch size = 2

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.22

1.63

1.91

0.95

0.94

0.94

0.96

0.94

0.96

1.03

1.25

1.41

Async

FGD-DP

FGD-MP

Sync

(c) Batch size = 3

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.24

1.62

1.98

0.97

0.96

0.95

0.95

0.95

0.98

1.09

1.4

1.76

Async

FGD-DP

FGD-MP

Sync

(d) Batch size = 4

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.49

2.19

2.91

0.95

0.92

0.95

0.91

0.92

0.9

0.87

0.86

0.88

Async

FGD-DP

FGD-MP

Sync

(a) Batch size = 1

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.5

2.18

2.92

0.94

0.96

0.94

0.96

0.95

0.93

1.13

1.05

1.42

Async

FGD-DP

FGD-MP

Sync

(b) Batch size = 2

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.52

2.25

2.93

0.93

0.93

0.92

0.94

0.93

0.95

1.28

1.56

1.75

Async

FGD-DP

FGD-MP

Sync

(c) Batch size = 3

2

3

4

K

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

Speedup

1.54

2.24

3.03

0.96

0.95

0.94

0.94

0.94

0.97

1.36

1.75

2.19

Async

FGD-DP

FGD-MP

Sync

(d) Batch size = 4

10 2

10 3

Number of Layers

0

500

1000

1500

Memory Consumption (MB)

Async

Async

BP

Sublinear

(a) Different number of layers

0

250

500

750

1000

Channels

0

2500

5000

7500

10000

12500

Memory Consumption (MB)

Async

BP

(b) Different channels

0

200

400

600

800

1000

0

100

200

300

400

Memory Consumption (MB)

Async

BP

(c) Different batch size

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper focuses on enabling training on memory-constrained platforms. Instead of using conventional backpropagation-based methods, the proposed approach is based on forward gradient descent (FGD), which approximates the gradients through only forward passes. The paper claims that a brute-force utilisation of FGD leads to resource underutilisation in memory-constrained devices. To counteract this, the work proposes AsyncFGD, a method that enables multiple parallel workers to process different parts of the network using different samples, while allowing varying levels of staleness in the parameters of these parts.\n\nSTRENGTHS:\n1) Reducing the memory requirements of the training stage is a key and still open challenge towards enabling on-device training. In this context, FGD-based algorithms constitute an interesting approach that is still largely unexplored.  \n2) The paper provides both theoretical analysis and empirical evaluation of the proposed approach.\n\nWEAKNESSES:\n1) In the motivation of work, the paper makes a number of arguments that are hard to justify. In Section 2.2, the paper claims that \"in edge devices with limited computational power, potential idle workers due to synchronization in synchronous pipelines like [12] are not optimal.\" However, this is rarely, or not at all, the case. Specifically, edge devices typically utilise their processing resources fully by exploiting the intra-layer parallelism, *e.g.* by performing convolutions or matrix multiplications. For the processor of such platforms to remain underutilised, either the device lies more towards server-grade or the target model is very lightweight, with small layers that contain minimal parallelism. This could make more sense in server settings, where layer pipelining has already been proposed (GPipe, PipeDream) as also mentioned in the paper. Therefore, the argument about Forward Locking and the serialisation of computations in the forward pass of FGD-based methods does not seem to be relevant for memory-constrained edge devices, where layer pipelining is hardly used during training.\nAs such, there are important questions with respect to the suitability of AsyncFGD. The presented implementation uses multiple workers to process different parts of the network, *i.e.* it employs module pipeling. For this to make sense as a strategy, the computational resources of the target device must be underutilised in the first place, *i.e.* when module pipelining is not used. In this context, what about batching? Batch processing is not touched upon in the paper. Why not have different workers processing different samples from the same batch, thus not requiring any stale parameters?  \n2) In Section 4.3, the paper makes a claim with respect to the attainable speedup. Specifically, it states that \"It is easy to see that AsyncFGD can fully utilize the computation resources, thus achieving considerable speedup\". This statement is very strong and cannot be supported solely by Fig. 2. In Fig. 2, we see a simple performance model for AsyncFGD, which indicates the potential acceleration *under certain conditions*. These conditions include (among others) *i)* whether the processor is really underutilised, *i.e.* whether there are really free cycles to have more than one module running, and *ii)* whether the multiple workers become external memory bandwidth-bound, in which case AsyncFGD can lead to slowdown - this is very common in edge devices where multiple components of the target chip share the same bandwidth to the external memory. In the scope of this work that aims at memory-constrained devices, the assumption that the memory bandwidth is abundant contradicts the whole setup. Overall, the logic of the paper with respect to acceleration leaves out many important considerations.  \n3) In Section 6.6, one would expect the speedup to plateau or even lead to slowdown, when *K* increases above a certain point, where the workers would start to context-switch excessively. How is *K* selected? What exactly affects it? Both the model architecture, the target dataset and the target platform would play important roles (see also point 2 above on this). In its current version, the paper is missing a thorough investigation of this fact.  \n4) In Section 6.2, while the objective of the experiments and the selection of baselines are sound, the network architectures are arbitrary and simple. It is hard to draw conclusions on the efficacy of AsyncFGD.  \n5) In Section 6.4, it is not clear what the setup is. The models are pre-trained on ImageNet and then fine-tuned on a given dataset using one of the algorithms in Table 2. Why does AsyncFGD* give higher accuracy than FGD and plain AsyncFGD? Furthermore, even Async* falls short compared to the BP-based baseline. Especially, given the issues of point 1) above, it is not evident that AsyncFGD has useful merits or under which conditions it would be a better choice than competing methods.\n6) In Section 6.5, the memory footprint comparison between FGD and AsyncFGD is missing in Fig. 3a and b - how do they compare? Importantly, how does the selection of *K* affect the memory footprint? Having more parallel workers can have a significant impact on memory consumption, through thread bookkeeping and cache misses.  \n7) As specified in the Appendix, only a single random seed is used in the experiments.\n\nQUESTIONS:\n1) Is something accidentally missing from Section 6.1?  \n2) Please also see points 5) and 6) above for things that were not clear.  \n3) Please proof-read the paper as there are some typos.\n\nLIMITATIONS:\nThe paper does not explicitly discuss its limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-26"}

comments: SUMMARY:
ThIS paper introduces AsyncFGD, a forward-only based training method. The key idea behidn this paper is integrating asynchronous updates with Forward Gradient Descent (FGD). The authors test their method on several small-scale datasets.

STRENGTHS:
- The proposed proof of the convergence guarantee of AsyncFGD-SGD is new.
- According to the experiment results, this framework achieves a good acceleration rate on small-scale datasets.

WEAKNESSES:
- The novelty of this approach is not introduced clearly. There is no fundamental difference between the original forward-only training methods by introducing a pipeline mechanism among different layers and this methods.
- The experimental results are only shown on small-scale datasets.
- The comparison with other works is not sufficient [1].
- The proposed method has on-par memory consumption with BP based training yet with lower accuracy. The memory cost is crucial for some edge devices.

[1] Krithivasan S, Sen S, Venkataramani S, et al. Accelerating DNN Training Through Selective Localized Learning[J]. Frontiers in Neuroscience, 2022, 15: 759807.

QUESTIONS:
- Can the author explain the fundamental difference between the original forward-only training methods by introducing a pipeline mechanism among different layers and this methods?
- Can the author show the large-scale dataset experiments?

LIMITATIONS:
Please refer to the weakness sections.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an asynchronous version of the forward gradient descent (FGD) method, in order to alleviate the forward locking exhibited with FGD and enable more efficient implementation. Specifically, AsyncFGD decomposes a network into K decoupled subnetworks that work as a delayed pipeline and enable asynchronous forward propagation. \nThe authors benchmark their method on an NVIDIA AGX Orin, using a variety of network architectures and datasets. Their method seems to perform better in the fine-tuning setup, where learning rates are small, due to the high variance of forward gradient methods in general.\n\nSTRENGTHS:\n- The paper is well motivated, and tackles an important subject: efficient on-device training\n- The authors properly explain their proposed method, the presentation is clear\n- The theoretical analysis on the convergence of the proposed method is appreciated\n- The authors benchmark their method on actual device and report measured performance\n\nWEAKNESSES:\n- The original FGD paper [1] does not report improvements in memory consumption (see Fig 6). I find it interesting that the proposed method does not suffer from the same. I also found that the overall memory discussion a bit lack-luster. It would greatly improve the quality of this paper, and emphasizes the importance of the proposed method, if some details were provided on how memory consumption is computed (or measured?) for all methods (BP, FGD, and AsyncFGD). How bad is the additional memory overhead due to the proposed asynchronous implementation?\n\n- The method has not been benchmarked on big datasets, such as ImageNet. Furthermore, it appears that, much like other forward methods, AsyncFGD cannot scale well to larger problems. This explains why the authors focused on transfer learning applications. This, in my opinion, limits the usefulness of such methods.\n\n\n\n[1] https://arxiv.org/pdf/2202.08587.pdf\n\nQUESTIONS:\n- typo in line 104: define \n- What is the dataset being evaluated in Fig 3c? \n- Why isn't the efficient learning strategy applied to FGD in Table 2?\n\nLIMITATIONS:\nNo discussion on limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposes a novel forward gradient descent method named AsyncFGD to decouple dependencies between layers and thus maximize parallel computation. The authors demonstrate that their method can reduce memory consumption and enhance hardware efficiency through empirical evaluations on AGX Orin.

STRENGTHS:
The proposed method of using asynchronized forward gradient descent is interesting and effective, and the authors provide both theoretical analysis and empirical verification, where the latter includes both accuracy and efficiency. The authors also provided extensive analysis and detailed setup.

WEAKNESSES:
One potential issue is that it is not very clear if this work is very suitable for a machine learning venue, as a large portion of the paper are talking about architecture- or system-level issues. Also, from Table 1, it seems the proposed method is not able to achieve consistently better results,like the sDFA or DFA can be better for MNIST, among others.

QUESTIONS:
Could this method also improve the energy efficiency? Could there be some comparison, even if basic analysis?

LIMITATIONS:
This paper did not talk about limitations. One potential issue might be related to the environment issue, like carbon emission of the proposed method compared to conventional ones.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08