BP-Modified Local Loss for Efficient Training of Deep Neural Networks
Authors: REN Lianhai
Qianxiao Li
Venue: ICLR 2025 Poster (2025)
Keywords: deep learning optimization
local loss training
bias-variance balance
Primary area: optimization
TL;DR: We proposed a novel method that periodically compute the BP gradient and use it to modify the local loss gradient. This method improve the performance of the original local loss methods with negligible additional memory usage.

Abstract: The training of large models is memory-constrained, one direction to relieve
    this is training using local loss, like GIM, LoCo, and Forward-Forward
    algorithms. However, the local loss methods often face the issue of slow or
    non-convergence. In this paper, we propose a novel BP-modified local loss
    method that uses the true Backward Propagation (BP) gradient to modify the
    local loss gradient to improve the performance of local loss training. We
    use the stochastic modified equation to analyze our method and show that
    modified offset decreases the bias between the BP gradient and local loss
    gradient, but introduces additional variance, which results in a
    bias-variance balance. Numerical experiments on full-tuning and LoKr tuning
    on the ResNet-50 model and LoRA tuning on the ViT-b16 model on CIFAR-100
    datasets show 20.5\% test top-1 accuracy improvement for the Forward-Forward
    algorithm, 18.6\% improvement for LoCo algorithm and achieve only on average
    7.7\% of test accuracy loss compared to the BP algorithm, with up to 75\%
    memory savings.

Published as a conference paper at ICLR 2025
-P-MODIFIED LOCAL Loss FOR EFFICIENT TRAINING OF DEEP NEURAL NETWORKS
Lianhai Ren
Qianxiao Li
Department of Mathematics
Department of Mathematics and Institute for Functional Intelligent Materials, National University of Singapore 10 Lower Kent Ridge Road, Singapore 119076 qi gianxiaoênus edu.
National University of Singapore
10 Lower Kent Ridge Road, Singapore 119076 1 Lianhai ren@u. nus edu
ABSTRACT
The training of large models is memory-constrained, one direction to relieve this is training using local loss, like GIM, LoCo, and However, the local loss methods often face the issue of slow non-convergence. In this paper, we propose a novel moddfiie local loss method that uses the true Backward Propagation (BP) gradient to modify the local loss gradient to improve the performance of local loss training. We use the stochastic modified equation analyze our method and show that modified offset decreases the bias between the BP gradient and local loss gradient, but introduces additional variance, which results a bias- variance balance. Numerical experiments on full-tuning and LoKr tuning on the ResNet-50 model and LoRA tuning on the VīT-b16 model on CIFAR- 100 datasets show 20.5% test top-1 accuracy improvement for the Forward- Forward algorithm, 18.6% improvement for LoCo algorithm and achieve only an average 7.7% of test accuracy loss compared the BP algorithm, with up to 75% % memory savings.
algorithms.
INTRODUCTION
Neural networks have seen rapid advancements in recent years, becoming widely applied across various fields. key algorithm for training deep neural networks is the Backward Propagation (BP) algorithm (Rumelhart et al. 1986). which needs to store the intermediate state in the forward pass to compute the gradient in the backward pass and thus leads to a large memory usage. This issue becomes more severe the models become deeper and require larger batch sizes, like SimCLR (Chen) et al. 2020). Consequently, reducing the memory footprint of training and fine-tuning tonitunngg remains an active of research (Krizhevsky et al. 2017 Rhu et al. 2016 Malladi et al. 2023).
To mitigate this memory issue, one common approach is gradient accumulation or distributed training strategies, such Data Parallelism (Zhang al., 1989 Das al. 2016) and Pipeline Parallelism (Huang et al. 2019)), which trades memory usage for additional computation time or enlarge the total memory by using multiple devices. An alternative using local loss (also known as local learning, or layerwise learning) to train the neural networks instead of the full BP (e.g. Greedy InfoMax algorithm (Löwe al. 2019). Forward-Forwart algorithm (Hinton) 2022)), where the backward propagation is applied only to portions of the model, reducing memory consumption. However, this method generally leads to poorer accuracy compared BP. Recently, the Model Predictive Control (MPC) framework (Ren & Li] 2024) has been proposed to unify the BP algorithm and the local loss method, balancing memory usage and performance through different horizons. uses the truncated loss to train the neural networks. However, the MPC framework still faces the trade-off between efficiency and performance: smaller horizons require less memory but result in lower accuracy.
Similar to leveraging global information to improve local loss (Lorberbom et al.) 2023). we propose a BP-modified local loss method. Our method integrates the true BP gradient modify the local loss gradient. Using stochastic modified equations (Li et al.) 2017] 2018). we show that this modification reduces bias between the gradients but introduces additional variance, thus resulting in mias-nariance

---

Published as a conference paper at ICLR 2025
trade-off. Through analysis based on the Orntein-UUllenbeck process and one-step loss, further illustrate this balance and derive the explicit update equations for the scaling factor A and a, by minimizing the balance. Numerical experiments on full-tuning and LoKr tuning on the ResNet-50 model and LoRA tuning on the ViT-b16 model on the CIFAR-100 dataset show up to 28% improvement in test top-1 accuracy for the Forward algorithm and 29% improvement for the LoCo algorithm, with only 7.7% average test accuracy loss compared to the BP algorithm while saving 75% memory.
The main contributions of this paper are as follows:
We introduce novel local loss training algorithm that significantly improves performance with minimal additional memory overhead.
We provide a theoretical analysis using the stochastic modified equation, illustrating the birs-variance trade-off and deriving optimal scaling factors.
Numerical experiments both CNN and vision transformer models using different tuning methods verify that our method can improve the performance of the local loss methods with minimal additional memory usage.
The rest of the paper organized as follows. In Section 2 introduces the proposed BP-modified local loss method. The theoretical analysis is in Section In Section we present the numerical experiments. Section discusses the limitations and future directions.
we review the related work. Section
2 LITERATURE REVIEW
To address the memory issue of the BP algorithm. many local loss methods have been pro- posed (Belilovsky al. 20T9). For example, Löwe et al. (2019) proposed the Greedy InfoMax (GIM) learning approach which uses local constructive loss to train the model in self- learning. Belilovssy et al. (2019) greedily train each convolution layer at one time and construct the model layerwise. Hinton| (2022) proposed the Forward-Forwant algorithm using simple local loss to train the MIP model. However, these methods suffer performance issues compared to the BP algorithm since the local loss cannot give the true BP gradient. Recently, Ren & Li] (2024) proposed the MPC framework that unifies the BP algorithm and the local loss method. The MPC framework uses the "loss split" method to create local loss from terminal loss and truncated loss. It uses different lengths ("horizon" the language of MPC) of truncated loss to balance memory usage and performance. It finds the convergence of the gradient to the true gradient in the large horizon. However, the MPC framework still has the efffciency-performaan trade-off: small horizons use less memory but result in larger losses.
To improve the performance of the local loss methods, one natural idea is to offer additional infor- mation. The LoCo algorithm (Xiong et al. 2020) uses the adjacent two blocks in the model for back-propagation to indirectly pass the global information. Lorberbom et al. (2023) uses the infor- mation from other layers to update the threshold in the Forward-FFrwart algorithm. The proposed BP- BP-modified local loss method uses the difference between the true gradient and local loss gradient in another small batch modify the original local gradient and reduce its bias thus improving the performance of the original local loss training method.
The implementation of our method uses the delayed update strategy, i.e. the modified term only updates periodically in larger batch, which is similar to the control variate methods such as SVRG (Johnson & Zhang, 2013). SARAH (Nguyen et al., (017). and Control variate forward gradient (Arisaka & Li 2024). Both our methods and theirs include some additional parameters that are expensive in computation (the control variate in their method and the modified term in our method), and an inner loop where these parameters stay unchanged. However, the motivation of the control variate methods is to reduce the variance while our method tends to reduce the bias in the expense of adding variance.
2

---

Published as a conference paper at ICLR 2025
3 METHOD: BP-MODIHIED LOCAL LOsS
In this section, we first introduce the deep learning training problem and the MPC framework. Then, we present our BP-modified local loss approach. We consider the feed-forward" neural network with I' layers or blocks, where each block is sequentially connected as follows:
I++1 = ft ut),
(I)
where E (O...,T- 1} }} denotes the block index, E TR denotes the input of the 1-th block, (with To being the model input and 2T the model output). and u, E Th are the trainable parameters the t-th block, the function ft R x _T - [R*++ represents the forward mapping of the f-th block. Moreover, we denote the combination of head and final loss as (xT), omitting possible parameters, targets/labels, and regularization terms for ease of notation. We assume L is compatible with x, for all 1. T. By compatibility of the loss L. we mean it can accept the outputs of all blocks, i.e. [xx)) is valid for Vt 1.... T. For further details of L. refer to Appendix
Recently, Ren & Li (024) introduced the MPC framework that unifies BP and the local loss method. algorithm (Hinton
Since this framework includes many other local losses like Forward- 2022) and LoCo algorithm (Xiong et al.] 2020), we use this framework to present our proposed method, which applies generally local loss training algorithms.
The MPC framework uses a "loss split" technique to decompose the terminal loss into local losses:
(12, ue) L(fe(Ee, un)) ((xx) L(z+4)) L(xe).
(2)
and then define the truncated loss:
h-1
(1,, u) Σ (ael( L(xn+t) - L(xe)
- E I represents all the weights in the model, m = IT- m, and E (1, ... } is the Ia = Ty when t > T and remove the constraint T and the minimum argument with T for simplicity. The gradient of u, under horizon h is defined as the corresponding gradient of La,e:
where u = (u
that controls the length of the truncated loss. We assume
Ih.t Vu, Ln. (z,, u) V., L(xn+t).
From Eq. we know that when h = T, .T, correspond to the gradient the terminal loss, reducing to traditional BP, i.c. gT' 9BP; when h L; 1 (Iz, 2)) Le(1,, uz) L(x++1) L(zz), which reduces to the Forward-Forward algorithm (Hinton 2022). Further, the LoCo algorithm (Xiong et al., 2020) can be
the case when h 2. For detailed example and comparison, please refer to (Ren & Li 2024).
The MPC framework adjusts the accuracy-memory trade-off through horizon h. Since the gradient of the f-th block 9h.t is V.,, +x)) a smaller h uses less global information thus less accurate and a larger horizon hi receives more global information and gets better performance. However, the gradient Ih only depends on the loss Tt+ so the gradient only needs to back-prop from L(x++h) to we through h blocks, reducing memory usage to (h), i.e., a smaller horizon will save more memory.
3.1 BO-MIIIFIED LOCAL LOSS
To retain the memory efficiency while improving the performance of small horizons, we propose a BP-modified local loss. This especially useful for the h = 1 case, i.e. Forward-Forward case, which has minimal memory usage. The method introduces an offset to modify the local loss gradient t, as illustrated in Figure T
gh.1 = 0h, Eug [9n + Bug [ōr,, cr. ax [ow,el + AE,
[Arr.l.
(5)
T 1, where (AB represents the empirical distribution of the mini-batch with size B, Le. Aa = B _ Ez,, and μæ is the empirical distribution of another independent mini-batch with size B', and letting g denotes the gradient obtained from (B. ğ denotes the gradients obtained from pr and ğ is the modified gradient. a₁ E R, An,tt E [0, 1] are two scaling factors, where a scales the original gradient and A scales the offset.
for all
3

---

Published as a conference paper at ICLR 2025
Offe
Mod fled Gradlent
Optiizer
Local Grediere
Figure 1: Diagram of the BP-modified local loss method
Usually, there is large deviation between true gradient gr and local (loss) gradient gh. The BP-modified local loss method introduces an additional offset Agh using local gradient and true gradient in batch 3² modify local gradient gi. modified gradient 9h. in other batches B.
The additional information of the BP- 3o-modifeed local loss method lies in the offset term Agr. which is the difference between the true gradient 9r, and the local loss gradient gh. of another If the difference of gradient in different samples are similar, i.e. Bgg [9h.t ST.e) R Bygr [9,, 9T, then the modified term Ağh, will reduce the bias of local gradient. However, the offset term itself will introduce additional variance. In the following section, we will theoretically analyze the impact of the offset.
pgr.
Furthermore, to acquire this offset term, we need to sample another mini-batch μa and compute the local gradient Ih., and true gradient 9I. on it. This will increase computation cost and memory used in the algorithm to mitigate the impact both time and memory
usage. Some techniques usage. For the detailed implementation, please refer to Section and Alg. Storing the offset Agh. will also need additional memory in the size of trainable parameters, which is negligible especially for CNN models and low-rank tuning cases such as LoRA (Hu et al. 2022).
4 THEORETICAL ANALYSIS
In this section, we theoretically analyze the performance of the proposed BP-modified local loss method using the stochastic modified equation of stochastic gradient descent (SGD) (Li et al. 2017 2018 Mandt et al., (017) which is applicable for various models and tasks. For clarity, we consider the full set of weights u, where the local gradient Sh. and true gradient gr the concatenation of the block gradient (i.e., = (0h.o: ... - block t individually. Further, we will omit the dependency of gradients g. and their variance C, on weights in the following for simplicity.
1 ). However, the following analysis can apply to each
Similar to traditional BP, the stochastic modified equation of SGD for horizon h within the original MPC framework can be written
dU(T) Epllg dr V No dW,
(6)
where is the first order continuous approximation of weights u as U(nk) Re u(k), denotes the empirical distribution of the whole training data, time index for the stochastic modified equation, is the learning rate, B is the batch size, and C₁ Coo(g) 11 is the autocovariance matrix of gh. and W, represents a standard Brownian motion. By Li et al. (2017). we can derive that U (njk) is first-order approximation of u(k) in a weak sense. Please refer to B.I for details.
To simplify the analysis, we assume that the mini-batch B' for computing offset is independently sampled at each step. The stochastic modified equation for the proposed BP-modified local loss method then becomes:
no Cn + B
T O dw., B'
(U = -E, - Nn)angh + Angri dr + V
denotes the covariance matrix of random vectors X and Y, and Cov(X) Cov( X) he the autocovariance matrix of random vector X.
In the paper, denote Cor(X.Y) F[(X (X [[
4

---

Published as a conference paper at ICLR 2025
where C₁ = Cov (gh) dj Ch Cr ar Cn, r is the autocovariance matrix of Agh. -r 94 Ch.T = Cor Ør + Coolgr enolggrghh... 9h is the covariance matrix between Ih and gr. The detailed derivation of Eq. and Eq. 7 please refer to Appendix B.1
By comparing Eq. 7 with Eq. (6, we can see that the proposed BP-modified local loss method reduces the bias between the applied gradient with the true gradient from [| Ex ll to
2A2 Cr. This leads to bias- [I(1 An. [gh] E [gr]ll2 while introducing additional variance F"' variance trade-off in the proposed method. We will further illustrate this trade-off using a toy model and then analyze it in the one-step setting in general.
Remark 4.1 Note that when ) = O. Eq. Z reduces to Eq. (6) since there is no addition offset. When d O, An 1, Eq. 7 describes the stochastic modified equation for traditional BP with batch size B' since there is only gT sampled from B' in Eq.
4.1 TOY EXAMPLE:
PROCESS
To illustrate the bias-variance balance in Eq. we consider an Ornstein-Uhlenbedd (OU) pro- cess (Uhlenbeck & Ornstein 1930). Although it is generally not the true training process of neural networks, this simple model allows for exact calculations which gives important insights into the effect of ' and bias-variance trade-off that is aligned with the discoveries in the following analysis.
Assuming u E I2" and the global loss and local loss with respect to u are:
1 1 T Au _ - 11 Ju) 2 "i
1
J(u) = 2
(8)
where A RRxn is symmetric positive-definite matrix and the optimal solution up:: " respectively, and un = O. Assuming the gradient variance V E "R"X" remains constant throughout the dynamic, the stochastic modified equation of SGD for both global loss J and local loss Jn follows Ornstein-Uhlenbeck (OU) process
dUr(T) - (J(UT(T))d + vVaw, = -UUr(T)dd + VVaW..
(9)
(T) -VJL(U + VVdW, --A(U,(T) ug,ddr VVaW, Assuming their initial point are the same,i.e. U7-(0) = (0) ugg IR...
(10)
For the BP-midiiied version, mimicking Eq. 7 and assuming 4 = 1 for simplicity, we add the modified term Au, (T) - (T) - -ff) and the additional variance V (which also remains constant) into Eq. 10).. we have:
dún(r) ((1 A,) wi) + A, Aū,(())))d + /v + V
(11)
Where the drift term modified towards true gradient (i.e. Au), with additional variance and rescaling. The optimal Xj and optimal expected loss can be explicitly solved:
Nx (e-AT AT un) T Au (T) + [[
(e~AT AT (C-ATun) AT Bias
(12)
ir (r(ETA) 1/ (7)
Var + Bias
(E) Au (T))
E[J(U((T)) min E[J E(JU,, 2 (tr(EtA) + (T) An
(13)
where - AT Juy ΣT= Sh e A dr. The derivations of Eq. 12 and Eq.
(13 please refer Appendix B.2 The Bias, Var terms in Eq. 12) just to indicate || (T) is bias-related and A) variance-related.
For large T. since em -- + 0 as 7 oc, we have (T) E (0, 1), demonstrating both An. 0 (original MPC) and X = (unbiased modification) are suboptimal choices. When bias is dominate, L.c. ||ug (T) < Ar( >; tends to 1 reduce the bias, while if the additional variance A) is large compared to the bias IT) Aj, tends to to reduce the variance. This the basic idea of the bias-variance balance which will also appear in the following one-step analysis. Further from Eq. (13), we know that for optimal Xi, the modified version will have better-expected
loss than the original.
5

---

Published as a conference paper at ICLR 2025
4.2 ONE-STEP BIAS- VARIANCE BALANCE
In the previous section, we demonstrate the bias-variance balance inherent in the BP-modified local loss method. However, modeling this through OU process may oversimplify the dynamics. In this section, we delve deeper into the local bias-variance balance of the proposed method using one-step loss analysis. We postpone the derivation of this section Appendix B.3 for clarity.
Let J()) E2 [[TT)] = En LT, (zo, u)] be the empirical loss with respect to weights u, where VJ(u) assumed be 3 3-Lipschitz continuous. Assuming (()(k);;) to be general gradient estimator whose randomness is controlled by the random variable &. Using the update rule u(k u(k++)) 1) u(k) ())
Taking the expectation conditioned on u(k) (denoted as Fg - [((-)(()))) we have:
7/3
n71
n/3
Ex|J(u(k + 1))1 S J(((k)) -
V J((()))
2
2x
(14)
n²3
+
R2 [I - {(u(k);;
where Ex [x[lV)((((H) - ((u(():: indicates the bias-variance balance since:
E (ellVV (1 ((((())))) ( (R(;;; Exll (o(())))) +]
(15)
The first term the norm of bias between true gradient Vu(((()) and the expected gradient Ex (g((());; E)], while the second term represents the variance of the gradient. Substitute g by In we have:
& [(((V)((())))) g((((((()))))))))))()) :) 1| -
27 Ch
X
(16)
(a Cn Cr - a.r-a))r.))
+ B'
which is a quadratic equation for both an, Ano Optimizing Eq. 16 over an and An. we obtain:
d T b + d Biii + b+ d h.aar
(17)
E& (angh 9r]l
Bias
X = I B (a, 20, + Or an Cr,T) + | [angn 9T
(18)
Var + Bias
where bioss T are the optimums of the bias and variance respectively, and b, d are defined as follows:
[on "TT] argmin [E
An) -9r)]12
0 bias
E 19
ABB
tr tr(Cu,r)
a B
X
argmin
C₁
(a Cr -r-an (n,Tr)))
(19)
X B B' (r((C))
B
X
b = (1-
(Bx [oN 0,d
+
tr(C 20
B'
we observe directly that the oj is the linear combination of the optimum of bias and variance, which strikes a balance between them. The weights d are related the bias and the variance of the local gradient gh- In the early training stage where the gradient is large and so is b, oj will tend to ai hbbas to reduce the bias. When the variance dominates, larger and leads az to har reduce the variance.
From Eq. 17p.
As for Mo. since it has similar Bias Variance structure as Eq. 13 observations: both A = 0 and X = are sub-optimal and A can adjust the bias-variance balance for fixed an.
Bias
we can derive the same
Remark 4.2 Simultaneously solving a,"" and X results in quintic function, making the solution highly complex. However, since Eq. (15) is convex in both ag and An we can simply optimize a, and A alternatively till their convergence.

---

Published as a conference paper at ICLR 2025
In conclusion, the balancing of bias and variance is the crucial point in the BP- BP-modified local loss since it reduces the bias of the local gradient at the expense of introducing additional variance. This balance exists throughout the whole analysis. If the variance is large, the estimation of the offset is inaccurate, leading the method to fail. Conversely, the balancing bias and variance gives a useful way sign the scaling factors.
Algorithm 1 BP-modifieed Local Loss Algorithm
Require: Starting point ug. Local loss gradient estimator Gh.r, Global loss gradient estimator gr. Total training step T, Gradient-based Optimizer Opt, Batch size B, B3, Sampling period K, Function M to get a, À
u(0)= Uo
2: k =
3: for Training period i E [1,.... IT do
4: Sample amall-batches MB with size B'
Lazy update of Agh, ahit An,t
5: for Layer t {0.-.. ,
do 6: Compute L 91 Ch.t: Cr,,, CA,T. on juB
7: Compute C,,,, An. = M (gh,e. (gh,t. 91,1, Ch. Cr.,, C,,r,))
D e.g. Eq.
Eq. 18
8: Compute gg [gh,)) gg (ĀT, ah
Eq.
9:
end for
10: while k <is K and & T do with size B
11:
Sample small-batches for Layer t {..... T' 1 do
13:
Compute Ih., μB Compute Gh.t a Ep [9h.el + An, ( Update weight (k) Opt(ue (k: - 1h.)) end for
Dr get local gradient, Eq. +> modify local gradient, Eq.
15:
16:
17:
lc k ++1
18:
end while
19:
end for
20:
return Final weight T(
5 IMPLEMENTATION OF THE BP-MODIFIED LOCAL LOSS ALGORITHM
In this section, we present the detailed implementation of the BP-modified local loss algorithm. Instead of simply applying Eq. (5) in every step, we use the following two techniques further mitigate the impact on both time and memory: (1) lazy update of the offset Ağn, and scaling factors ans As: (2) split mini-batch B' into smaller mini-batches to compute the offset. The final algorithm presented in Algorithm[]
Lazy update redueing computation overhead Assuming the gradient does not change rapidly, we periodically update the offset and scaling factors and use the same values within each period to reduce the time required for computing the offset. In other words, we sample pap and compute the offset Agh. and scaling factors 0n, Å₁ at the beginning of every K steps, and apply the same offset in Eq. 55 to modify the local gradient in the next K steps. The additional relative computation overhead
will be about (1+r) where is the ratio of the computation time of computing the true gradient to KBB
the computation time of acquiring the local gradient.
Batch split reducing memory overhead Further, the total memory overhead of the BP-modified local loss method will be on the scale of O(Bh, B'h, B' We need B' < + B to sustain comparable memory usage as the original algorithm, i... O(B)). However, small B' will increase the additional variance introduced by the offset (Eq. (7). Combined with the impact of lazy update strategy (see following discussion and Appendix for details), the variance of the offset will be
Firstly, the memory usage will increase linearly with batch size B and it is reported that the memory usage will grow linearly with respect to horizon h (Ren & Li 2024). and these two factors are independent. Secondly, acquiring different gradients is asynchronous only need to consider the maximum memory for different gradients.

---

Published as a conference paper at ICLR 2025
unacceptable. Therefore, further split the mini-batch pg into smaller mini-batches and use the gradient accumulation technique to compute the offset, and the memory overhead will reduce to O(Bh, B'T, hhh) where 6' is the size of the smaller mini-batch.
However, some problems will appear after applying these two strategies. For example, the noise in the offset will be accumulated in the period thus increasing the equivalent variance. As same offset in one period, the same noise will accumulate, leading to a larger variance in the long run. For further discussion of the impact of delayed gradient adjustment, please refer Appendix C Secondly, further splitting aay into smaller mini-batches requires more computation time. However, the lazy update relieves this issue. As long as B' is smaller than the scale of K B. the impact of computation time will be small, thus B' no longer needs to be smaller than B.
the
From the above discussion, we know that larger batch size B' and quicker update (i.e. smaller period K) will reduce the variance of the offset and the delayed bias thus improving performance, at the expense of more computation. We have shown the numerical results of different choices of K and B' in Appendix D.2 and its result is aligned with the analysis.
There is an inner loop (i.e. the update period) in the implemented BP-modified local loss method, and B' will become larger by introducing the lazy update. This procedure is similar to control variate methods like SVRG (Johnson & Zhang 2013), where an additional term introduced (the variance reduction term in SVRG and the offset Agn our method), which is updated periodically. However, the basic motivation of the control variate method to reduce the variance, while method aims to reduce the bias. Furthermore, the variance reduction methods based on acquiring the true gradient, i.e., BP, while our method is based on the local loss gradient.
-onouuuunnnn and
NUMERICAL RESULTS
To evaluate the effectiveness of the proposed BP-modified local loss method, we conducted several
experiments on CIFARI00 (Krizhevsky) 2009) and ImageNet- Tiny (mnmoustafa & Ali (017) dataset:
(1) Fine-tuning pre-trained ResNet-50 model (He et al.) 2016); ; (2) LoKr fine-tuning (Hyeon- -Hooon-Wool
ct al. 2022) of a pre-trained ResNet-50; (3) LoRA fine-tuning (Hu et al. 2022) pre-trained ViT-b16
model. We test the proposed method on these tasks with Forward-Forwar algorithm (FF) (Hinton
2022). LoCo algorithm (Xiong et al., 2020), and MPC algorithm (Ren & Li 2024) with horizon
h 5 and the traditional BP algorithm. All the experiments are done by ourselves since no related
results are provided in the original papers.
All the experiments use batch size B = 64, learning rate 0.001 and epoch 30, and SGD optimizer with momentum 0.9. For the BP-modified local loss method, the batch size to compute the offset is B' = 320 and split into mini-batches with size 8, and the period of the offset update is one epoch, i.e. K = _ where N is the size of the training dataset. The input image is resized to (224.224) and then normalized without any other enhancement before training. For the LoRA and LoKr, the rank and alpha are set to be a, = 4. All the experiments were conducted using single NVIDIA GeForce RTX 3090 GPU using Py Torch (Paszke et 2019).
In summary, the numerical experiments show that
1. The BP-modified local loss method outperforms the original local loss in terms of perfor- mance.
2. The performance improvement more significant in small horizons, i.e. algorithm and LoCo algorithm.
3. Ablation studies on the impact of the offset show the importance of the offset in improving the local loss.
We further provide more experiment results in Appendix I iincluding:: (1) memory/time usage of the BP- moddiid local loss method (Appendix DDI:: (2) sensitivity analysis of the update period K and the offset batch size B' (Appendix(D22p (3) number of epochs the modified local loss method needs to achieve the original test accuracy (Appendix D.3
8

---

Published as a conference paper at ICLR 2025
Table 1:
Test accuracy of BP-modified local loss on different methods and tasks on CIFARI00 Dataset. The numbers in parentheses indicate the improvement over the original local loss method.
Task
Method
ResNet--0 Full Tune
ResNet-30
ViT-b16 LoRA
BP
82.90
76.47
91.49
original
40.10
54.43
63.65
FF
BP-modified
68.30 (+28.20)
65.34 (+10.91)
86.30 (+22.65)
original
45,29
55,84
74.31
LoCo
BP-modified
74.61 (+29.32)
66.15 (+10.31)
90.54 (+16.23)
original
66.02
65.65
84.03
MPC (h=5) -modified
77.24 (+11.22)
67.77 (+2.12)
90.29 (+6.26)
6.1 PERFORMANCE IMPROVEMENT OF THE BP-MODIFIED LOCAL LOSS METHOD
In this section, we highlight the performance of the BP-modified local loss method in improving local loss. We compare the results of the BP-modified method against the original methods the same model and tasks with identical training settings.
Table shows the test accuracy of the BP-modified local loss method on different methods and tasks. The results indicate that the BP-modified local loss method significantly outperforms the original local loss method across all tasks and methods. It achieves an average improvement of 20.53% on the FF method, 15.29% on the LoCo method, and 6.53% on the MPC method with horizon h = 5. The test accuracy of the BP-modified local loss method is comparable to the full BP method.
Additionally, we observed that the performance improvement is more significant in the FF method than in the LoCo and MPC methods. Since the FF method is the MPC method with horizon h this result suggests that the BP- moddiiee local loss method is more effective when the original performance is worse. As the FF method uses the least global information, the bias of the local gradient is more significant, leading to larger offset as shown in Eq. (18) thus more significant performance improvement.
Further, we observed that the test accuracy of the BP-modified local loss method is higher in larger horizons, i.c. LoCo and MPC h = 5 cases. These original local loss methods already perform better, indicating smaller bias and maybe more alignment with the true gradient. This results in smaller variance of the offset (Eq. 7) However, the performance improvement of the BP-modified local loss method in the LoKr task with horizon h 5 is much smaller than the other tasks. This is largely due to the estimated variance dominating the biri-nar trade-off, leading to a smaller offset.
We also conducted experiments the ImageNet-Tiny dataset with the same models and settings used in the CIFAR 100 dataset, except we set the total epochs to 10 and period K 500. The results are shown in Table 2 Consistent results are found that the BP-modified local loss method can improve the original method when the original performance is worse. In the LoKr tuning case where the original methods have achieved good performance, then the proposed method cannot improve their performance.
We also provide the memory and time usage of the above experiments in Appendix D The results show that the memory and time usage of the mo-miiif local loss method are comparable to the original local loss method. Additionally, the memory overhead is significantly lower than the full BP method.
6.2 ABLATION STUDY ON THE IMPACT OF THE OFFSET
In the ablation study, we investigate the impact of the offiet on the performance of the BP-modified local loss method. We set > = 0 to remove the offset from the BP-modified local loss method, leaving the scaling factor aj, = aj hias- We conducted the study on the FF method and show the results in Table Although the scaling factor aj improve the performance of the original local loss method without the offset, adding the offset can further improve the performance. This result indicates the importance of both the offset and the sealing factor in improving the local loss method.
9

---

Published as a conference paper at ICLR 2025
Table Test accuracy of BP-modified local loss on different methods and tasks on ImageNet-Tiny Dataset. The numbers in parentheses indicate the improvement over the original local loss method.
Task
Method
ResNet--5 Full Tune
NesNet-50 LoKr
ViT-b16 LoRA
BP
78.56
74.22
87.31
original
24.58
61,23
53.03
FF
BP-modified
68. (+43.56)
53.58 (-7.65)
75.92 (+20.89)
original
31.02
61.97
58,55
LoCo
BP-modified
68.49 (+37.47)
66.50 (+4.93)
74.13 (+15.58)
original
55.11
66.27
72.18
MPC (h=5)
BP-modified
66.03 (+10.92)
65.30 (-0.97)
76.48 (+4.30)
Table 3:
Test accuracy of the BP-modified local loss method with and without the offset the FF method. The numbers in parentheses indicate the improvement over the original local loss method.
Task
Method BP
ResNet-50 Full Tune
ResNet-50 LoKr
ViT-b16LoRA
82.90
76.47
91.49
original
40.10
54.43
63.65
BP-modified (A = 0)
46.21 (+6.11)
59.71 (+5.28)
74.36 (+10.71)
FF
BP- -o-modified
68.30 (+28.20)
63.71 (+9.28)
86.30 (+22.65)
Further sensitivity experiments on the update period K and the offset batch size B' are provided in Appendix D.2
DISCUSSION
In this paper, we introduced novel BP-modified local loss method. This method incorporates an additional offset to adjust the local gradient. Through the use of the stochastic modified equation, we demonstrated that the offset helps reduce the bias between the local gradient and the true gradient. However, it introduces additional variance. By employing the Ornstein- Seeeinn-Uheebeckk (OU) process and one-step loss analysis, we analyzed the bias-variance trade-off and derived optimal scaling factors, An and Og The impact of delayed versions of the offset was also explored. Numerical experiments across various models, tasks, and optimizers illustrate the potential of the proposed method. However, certain limitations should be acknowledged.
The OU process model used Section is simplified approximation of the stochastic gradient descent (SGD) dynamics in neural network training. Its primary role is to provide insight into the global bias-variance trade-off in the training dynamies. While several simplifications were made obtain an explicit solution, the OU model is sufficient to capture the key phenomena of bias- variance balance and the sub-optimality of the unbiased modification. These insights were later validated in our one-step analysis and numerical experiments.
Some important aspects, however, remain under-explored. For instance, the bias introduced by the delayed offset, discussed in Section and certain hyperparameters, such as the offset batch size B' and the sampling period K, are not fully analyzed. Although these factors are critical, their complexity exceeds the scope of this work and is left for future investigation.
The computation of the offset requires access to the true gradient. This process involves further discussed in Section This adds complexity
splitting the batch size B' into smaller sub-batches to the training process. Integrating this approach with methods like forward gradient techniques (c.g..
MeZO (Malladi al. 2023)) and control variate gradient estimation (e.g., (Arisaka & Li 2024
could be a promising direction. We reserve this for future exploration.
Due to limited computational resources, did not experiment with larger models datasets, and all experiments were conducted on a single GPU without distributed implementation. Nevertheless, experiments across various models and tasks consistently demonstrated the potential of the
10

---

Published as a conference paper at ICLR 2025
BP-modified local loss method. Further optimization of the method and implementation, including distributed settings and tests on larger datasets, will be explored in future work.
ACKNOWLEDGMENTS
This research is supported by the National Research Foundation, Singapore, under the NRF fellowship (project No. NRF-NRFF13-2021-000
REFERENCES
Sohei Arisaka and Qianxiao Li. Accelerating legacy numerical solvers by non-intrusive gradient- meta-solving, 2024.
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to imagenet. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning. ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 583-593. PMLR, 2019.
Ting Chen, Simon Kormblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. Proceedings of the 37th International Conference Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 15.115977 1597-1607. PMLR, 2020.
Dipankar Das, Sasikanth Avancha, Dheevatsa Mudigere, Karthikeyan Vaidynathan, Srinivas Srid- haran, Dhiraj Kalamkar, Bharat Kaul, and Pradeep Dubey. Distributed deep learning using synchronous stochastic gradient descent, 2016.
Kaiming Hc. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pp. 770-778. IEEE Computer Society, 2016. doi: 10.1 00.1000/VP....01..99.
Geoffrey Hinton. The forward-forward algorithm: Some preliminary investigations, 2022.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In The Tenth Interna- fional Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Bus Bucé-Buu,, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 103-112, 2019.
Nam Hyeon- Hyeon-Woo, Moon Ye-Bin, and Tae- Hyun Oh. Fedpara: Low-rank hadamard product for communication-efiiienn federated learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Rie Johnson and Tong Zhang. Accelerating stochastic gradient descent using predictive variance reduction. In Proceedings of the 26th International Conference Neural Information Processing Systems Volume 1, NIPS' 13, pp. 315-323, Red Hook, NY, USA, 2013. Curran Associates Inc.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. URL https:
//api, s...emanttt //api.semante i wtw.cacac.......oooooooo.ooo....
DD: 8268744
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep con- volutional neural networks. Commun, ACM, 60(6),84-90 may 2017. ISSN 0001-0782. doi: 10.1 145/3065386. hwwp:ntteee///ii..rr//0....1111.

---

Published as a conference paper at ICLR 2025
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and adaptive stochastic gradient algorithms. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th Interna- tional Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp- 2101-2110. PMLR, 2017.
Qianxiao Li, Cheng Tai, and Weinan E. Stochastic modified equations and dynamics of stochastic gradient algorithms i: Mathematical foundations, 2018.
Guy Lorberbom, Itai Gat, Yossi Adi, Alex Schwing, and Tamir Hazan. Layer collaboration in the forward-forwate algorithm, 2023.
Sindy Löwe, Peter O D'Connor, and Bastiaan S. Veeling. Putting an end to end-to-end: Gradient- isolated learning of representations. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Alcé--Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 3033-3045, 2019.
Sadhika Malladi, Tianyu Gao, Eshaan Nichani, Alex Damian, Jason D. Lee, Danqi Chen, and Sanjeev Arora. Fine-tuning language models with just forward passes, 2023.
Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic gradient descent as approximate bayesian inference, 2017.
mnmoustafa and Mohammed Ali. Tiny imagenet. Einy-imagenee 2017. Kaggle.
Lam M. Nguyen, Jie Liu, Katya Scheinberg, and Martin Takác. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 2613- 2621. PMLR, 2017. URL/tttp://proceeiii
mir wtw.:/ppreees////nnnueeeee 7tmmml
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, perforreform deep learning library. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Acché- Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8024-8033 2019.
Lianhai Ren and Qianxiao Li. Unifying back preck--rppggationn and haward-forwant algorithms through model predictive control. arXiv preprint arXiv: 19561, 2024.
Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W. Keckler. vdnn:
Virtualized deep neural networks for scalable, memory-efficient neural network design. In 2016 49th Annual IEEE/ACM International Symposium on Microarchitecture (MICRO), pp. 1-13, 2016. doi: 10.1 10.1009/MCRR022016.77733
David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by back-propagating errors. 323(6088):5333536 1986. ISSN 1476-4687. doi: 10.1038/323533a))
Publisher: Nature Publishing Group.
G. E. Uhlenbeck and L. S. Omstein. On the theory of the brownian motion. Phys. Rev. 36:823 841, Sep 1930. doi: 10.1103/PhysRev.3..222. URL PhysRev. 36. 823
Yuwen Xiong, Mengye Ren, and Raquel Urtasun. Loco: Local contrastive representation learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Xiru Zhang, Michael McKenna, Jill Mesirov, and David Waltz. An efficient implementation of the back- back-propase propagation algorithm on the connection machine CM-2. In Advances in Neural Information Processing Systems, volume 2. Morgan-Kaufmann 1989.
12

---

Published as a conference paper at ICLR 2025
A EXPLANATION FOR LOSS L (REN& & LI 2024)
In Ren & Li 2024). the author states that the loss function includes the head block and other possible weights, and target y is omitted as stated in the main paper. Moreover, it vary through the block, i.c. the loss function L in ((xe) and L(x++1) ) in Eq. (2) can be different. However, in this paper, we still assume the loss needs have the same structure.
For example, assuming an image classification task, is allowed that the two loss functions Lg in Lx(x)) and Lu+1 (I++1) ( with the same global pooling-linear www.s-ssooooa-osss-ooeee to have independent weights in its linear layer and the weight even size of the linear layer can be different (e.g., in CNN networks, the number of channels will increase through the block). We still think of them as the 'same' loss and thus do not distinguish them in the main paper. However, it is not allowed that Lt use cross-entropy loss while L++1 use contrastive loss.
loss structure
The study of using different losses in different blocks is left for future work.
B DETAILED DERIVATIONS
In this section, we will give detailed derivations of Section
B.1 DERIVATIONS OF THE STOCHASTIC MODIFIED EQUATIONS
AND 7
In this section, we will give the derivation of the stochastic modified equation Eq. (6) and Eq. using the analysis and result of et al. (2017).
For the SGD update of general gradient estimator {(())));(())) where the random variable {(k)
controls its randomness:
ukk + 1) = u(k) ng((((); = - (1 +
(20)
where v) is the learning rate and V(u(k);E(A+) V (() - vig(u(();E(())) (()))))))()))))))))))))))()))))) By
simple computation, we can know:
Ex(x)) (v())
= 0,
(21)
Thus Eq. 20 can be seen as the Euler-Maruyana disseettaaaaasceetttiaaiiai
discretization of the SDE:
dU(+) -b(,, + VnE(U,T)dW,
(22)
for a time step o, = and u(k) s U (mn)), where b(U,T) = Bcc)loU((())))))) Cou(g(U(r):((r))
=
Denoting G is the set of functions that have most polynomial growth, i.e.
E G + 3xr K2, L 0, lg < L(xi + ||>=1/".),,
(23)
and Gt is the set n-times continuously differentiable functions with its partial derivatives belonging in G, and T is the weak differentiable analog (detailed definition please refer to (Li et al. (017)).
From Theorem 9 and Corollary 10 in (Li et al. 2017), if the gradient estimator ((((k))E(()) the following assumptions:
satisfies
1. Uniform linear growth for mean and variance in U, T. i.c.
[|b(U, IEE(,, < (1 + || (2), VU E RR", E [0.1]
(24)
for some 81 > 0,
2. Uniform Lipschitz condition for mean and variance in U, r, i.e.
( 7)) (MUU, -)( T)) < 92 12 2 YU RD TE E 1]
(25)
3. b(U,)) is continuously differentiable and b(U, E G 4. Vg(U((). is Lipschitz continuous almost surely in U with Lipschitz constant (3(T)) and Rg(r) ) xo for all [0, 1) and 2
13

---

Published as a conference paper at ICLR 2025
then (r) is a 1-order weak approximation of u(k), i.e. for each g E G?, exists C' > 0 independent of 17, such that:
max
[B())(((()) (ll))lll= - Cn
(26)
k=
Then we just substitute g(u: <) by local gradient gh to get the Eq. 6 As for Eq. the mean and
covariance modified gradient On. are:
E [ = En )nnan + Angrl
(27)
a Cor(gs) + Cou(Agn).
X
Cov(@n) =
(28)
B
B'
Then we can get the I 1-order stochastic modified equations of the proposed method Eq. 7
B.2 DERIVATION OF OU PROCESS
In this section, we will derive the Eq. 12 and Eq.
13 in Section 4.1
By Eq. (10)
(T)(T) = V.) + Vvaw, = -A(U -(fr)))-;jddd ",, VVdw,
(29)
and the property of the OU process, the distribution of UT(T) satisfies the Gaussian distribution: Un(T) N(
AT ( + (T),
(30)
where uT( = e and En = So e- -- A dr. The expectation loss at time T is:
[[J(UN = Ar + ((LU
(31)
which consist of a bias-related term ((B(Un(T)) and variance-related term 2 A).
Similarly by Eq.
A,,Aū VV + ARV (T) = -((1 (n) - un)
(32)
the distribution of ún (T) is:
(T) N (e-AT
AT + (1- An) (T), 13.2%)
(33)
where = e- -AVe- A dr. The expected loss is:
X
(T) Au)( (34)
To minimize the expectation loss, we need to balance the bias and variance by Abs which results in: (TT x²()) (e- AT un) T Au (T) + [,;
(35)
tr A) (T)
(B(Un, AuR (T))²
min (() = =((( 2 (Ir(E (T) (1)
(36)
An
which is Eq.
12 and Eq.
B.3 DERIVATIONS IN ONE-STEP Loss ANALYSIS
In this section, we will derive Eq. and 16 in Section 4.2
By the 8 Lippstcihh continuity of the gradient VJ, we have:
J(u)) - J(wz) = V.[(tz + - (1=)) T (w1 wz) dt
f (VJ(tu) + +1- V T(u
(2i = =V](2))T T u2) + j
(m) - u2) dt
9
<VJ(uz)] ((((2)(([ (w1 - (2) マタ
[1 U2 11
(37)
14

---

Published as a conference paper at ICLR 2025
Substitute u1, d2 by u(k + 1), u(k), we have:
Ju(k + 1)) - J(((k))
I ()(((())) + +) u())) + 3 |lu(k + +1) u
2
-n g()(()):))+ g(u())))) + 3²²3
2 Il(((()):E]l1ã
(38)
-(1 m3) (((())())(((((()))))) T/
(I) -
n
n - (v)(())) +
3
2 + ; |[ g(((k); 1
S 11
2x
where the second inequality is the Cauchy-Schie Schwartz inequality for arbitrary > 0. Taking expectation condition on ukk) on both sides we can get Eq. (14)
n8
n
n9
[x[]( + 1)M1 S J(((k)) - 1)
[V(((
2
(39)
1/3 + #22 Ex [x]lV)(u(()) - g(u(()); 2
+
For the derivation of 0i Xi. consider the last term in the previous equation (i.e. E& [(((()(((()))
( and substitute g g(u(k); E) by Gr. we have:
E& ((u()) - ((u(k); 3
[J(()))
+ (x(ll((u(kk; Eelg(u(k);;
X
=|1Ba(1 Xangnn - 9r)1112 Ir
Ch
+
Ch
(40)
B
B
ah
X1
[[1 (angh 9r) +
Ch
+
(ah CT
-a, a..r.r)).
B
B'
which is Eq.
16
C IMPACT OF DELAYED GRADIENT ADJUSTMENT
Since the computation of gradient adjustment in the BP-modified local loss method every step can be computationally expensive, in practice, we update the sample offset Ag, and a, Ad, only every K steps, which reduces the computational cost. This delayed gradient adjustment introduces two effects on the dynamics: noise accumulation and biased adjustment.
Applying the same offset noise over multiple steps in a period introduces higher variance. To illustrate this, consider the OU process in Section 4.1 using the Fuler- Maruyama method with timestep dt and A = I, & (0, Vdt), E. ~ N(0, X Vtt), then the update are:
Accumulation of Noise
(T (T + dt) = ux(() dt((u(() uuf) + Er, ūn + dt) = - -i,)) E. -dtuj + (1- (t- (T) +
(41)
(42)
After K steps, the solutions are:
K-1
up un(Ka (Kdt) = (1 -1- dt) un (1 - (1 di) ddt)*)ni +
Σ
(1 dt) K K-r-1g ET
(43)
7=0
K
K-1
ün (Kdt) (1 d uo Σ (1 - d1 K- +
Σ (1 dt) --
(44)
T=0
If we apply the first modify term - -1, dl + 5o over the entire K steps, the solution for un (Kdl) becomes:
K
1 (1 dt K
(Kdt) = (1 (1d un +
Σ
& (1 d +
So.
(45)
dl
7=0
15

---

Published as a conference paper at ICLR 2025
The variance of the noise terms can be computed as:
K-
Σ (1-dt) NN N(,, 1- dt)22 - 7=0
-- (1-11 dt)x to ~ N
(0.(1-(17
dt) 2
2dl + di2 X v
dl
dt
(46)
1-(1-dt)2 K and (اااادد ~ K² indicating that the variance
When dt < 1, we have 2dt increases by factor of approximately K. This noise accumulation will also happen in the stochastic modified equations (Eq. (7).. and thus in practice, we multiply the variance introduced by the offset
by factor of K (i.e. nd
B' Ch..
Biased Adjustment The delayed gradient adjustment introduces additional bias E[A,, on (g)(T) (T(T") An (T) gr()") where 7' is the last offset-sampling step. As mentioned in Remark 4.1 when A 1. the modified gradient is an unbiased estimator of the true gradient. However, in the delayed adjustment case, this is no longer true. This additional bias introduced by this delay will also affect the bias-variang trade-off and thus the optimal values of ag and Xn. which we leave for future study.
D FURTHER EXPERIMENT RESULTS
This section provides additional experimental results to further demonstrate the effectiveness of the BP-modified local loss method. We report the memory usage and runtime of the BP-modified local loss method compared to the original local loss and full BP methods of the experiments in Section 6.1 We also provide additional results on the impact of period and batch size B' on the performance and efficiency of the BP-modified local loss method.
D.1 COMPARISON OF MEMORY USAGE AND RUNTIME
We compared the memory usage and runtime of the BP-modified local loss method against the original local loss and full BP methods. Memory usage was evaluated by profiling several training steps, including the original training step and the computation of offsets, and recording peak memory usage. Detailed procedures for memory profiling are provided in Appendix E
indicates that the memory usage of the BP-modified local loss method comparable to the original local loss method, which is significantly lower than the full BP method. The additional memory required by the BP-modified method is minimal, as it primarily involves storing offsets (equal to the size of trainable parameters) and scalar scaling factors. As described in Section 5 most memory usage during training stems from storing intermediate states for backpropagation, making the memory overhead introduced by the BP- modifiee method negligible. This is particularly evident in low-rank methods like LoRA and LoKr, where the trainable parameter size is significantly smaller than the model size, rendering the memory overhead negligible.
On the other hand, Table 5) compares the runtime per epoch of the original and BP-modified methods. The BP-modified local loss method incurs small runtime increase due to the computation of offsets and scaling factors. This overhead is mitigated by the lazy update strategy. In our experimental setup with B' = 320 and period K = epoch, the additional runtime under 6% of the original local loss method. This increase higher than the expected batch size ratio ( 0.64%) due to the mani-batch split technique.
Combining memory and runtime efficiency with performance results from the previous section, the BP-modified local loss method demonstrates strong balance of accuracy and resource usage.
D.2 SENSITIVITY TO PERIOD K AND MODIFY BATCH SIZE B'
We further examine the sensitivity of the moddiiie local loss method to the update period K and modify batch size B' in the mP-modified local loss method. In the FF algorithm experiments, we vary = {0.1 epoch, 0.3 epoch, 1 epoch} and B' = {32, 320, 3200), while keeping other hyperparameters consistent with the settings in 6
16

---

Published as a conference paper at ICLR 2025
Table 4: Memory usage (MB) of BP-modified local loss on different methods and tasks. The numbers in parentheses indicate the additional memory usage and its percentage of the BP-modified local loss method compared the original local loss method
Task
Method
ResNet-50 Full Tune
ResNet-50 LoKr
ViT-b16 LORA 5586
BP
5617
5529
original
1555
1366
1519
FF
BP-modified
1642 (+88 (5.63%)
1370 (+4.1 I (0.30%)
1520 (+0.15 (0.01 %))
original
2239
2251
1926
LoCo
BP-modified
2329 (+90 (4.02%)
2252 (+1.1 1 (0.05%))
1926 (+0.15 (0.01 ())
original
3965
3965
3146
MPC (h=5)
BP-modified
4055 (+91 (2.29%))
3968 (+3.1 (0.08%))
3146 (+0.15 (0.005%)
Table 5: Runtime per epoch (s) of BP-modified local loss on different methods and tasks. The numbers in parentheses indicate the additional runtime and its percentage of the BP-modified local loss method compared to the original local loss method
Task
Method
ResNet-50 Full Tune
50 LoKr
ViT-bló LoRA
BP
127.07
193.51
322.84
original
123.03
171.11
291.99
FF
BP-modified
126.79 (+3.76 (3.06%))
177.1 (+5.99 (3.50%) 298.00
296.37 (+4.38 (1.50%)
original
201.53
527.05
LoCo
BP-modified
205.67 (+4.14 (2.05%))
314.58 (+16.58 (5.56%))
515.08 (-11.97 (-2.27%)
original
363.71
566.94
967.10
MPC (h=5) BP-modified
367.48 (+3.77 (1.04%))
578.08 (+11.14 (1.96%))
979.28 (+12.18 (1.26%))
Theoretically, larger period K coresppondig to a less frequent update of Agh.t and scaling factors a,,t, Th..-inttroduce greater variance and amplifies the delayed bias, potentially impacting performance, discussed in Appendix C Conversely, smaller B' increases variance, suggested by Eq. 7 In contrast, using a smaller K or a larger B' improves performance but comes at the cost of higher computational requirements.
Table 6 and 7 summarize the results for different B' and K values. These results align with our theoretical analysis: larger B' and smaller K lead to better performance but increases computation. A detailed comparison shows that if the gradient will be large and the performance might be worse, while it seems that has more tolerance for larger K. If the period K is too small B' is too large, the computation overhead will be large with diminishing performance improvemen( Our experience suggests that should be large enough to reduce the variance of the modified gradient while K can be relatively large such that the computation overhead is small. B' being 1% of BK are usually good choices.
size B' is too small, the variance of the modified
Table Sensitivity of modify batch size B'. Test accuracy (in percentage) and training time (second epoch) are reported.
ResNet-50 Full Tune
ResNet--0 LoKr
ViT-bl6
LoRA
Method
Test Acc
Time(s)
Test Acc
Time(s)
Test Ace
Time(s)
original
40.10
123.03
54.43
171.11
63.65
291.99
BP-modified (B'=32)
38.97
125.32
57.84
173.19
74.88
317.63
FF
BP-modified (B'=320)
68.30
126.79
65.34
177.10
86.30
296.37
BP-modified (B'=3200)
74.62
152.88
71.91
219.16
90.47
333.50
Notted that there is no change in memory overhead because of batch split technique. Larger B3' means more computation but memory usage.
17

---

Published as a conference paper at ICLR 2025
Table Sensitivity period K. Test accuracy (in percentage) and training time (second per epoch) are reported.
ResNet-50
Full Tune
ResNet-50 LoKr
ViT-bl6 LoRA
Method
Test Acc
Time(s)
Test Acc
Time(s)
Test Acc
Time(s)
original
40.10
123.03
54.43
171.11
63.65
291.99
BP-modified (K=0. Epoch)
73.98
152.30
72.99
218.52
89.97
337.07
FF
BP-modified (K=0.3Epo) 3
68.18
133.28
65.34
186.26
87.83
324.80
mediioie (K=1Epoch)
14
126.79
65.34
177.10
86.30
296.37
Table 8: Average first epoch reach the same test accuracy as the original local loss method in the CIFAR100 dataset. The numbers in parentheses indicate the final test accuracy of the original methods.
Task
Method
ResNet-30 Full Tune
ResNet-50 LoKr
ViT-b16 LoRA
FF
1.0 (40.10)
3.0 (54.43)
4.0 (63.65)
LoCo
1.0 (45.29)
3.3 (55.84)
2.0 (74.31)
MPC (h=5)
2.0 (66.02)
13.0 (65.65)
4.0 (84.03)
D.3 EPOCHS TO GET THE ORIGINAL PERFORMANCE.
Here we provide the number of epochs needed for the BP-modified local loss method to achieve the same test accuracy as the original local loss method in CIFAR 100 dataset. We report the results in 8 We find that the BP-modified local loss method requires much less epochs to reach the same performance as the original local loss method, demonstrates the efficiency of the proposed method. also consistent with the significant improvement in test accuracy.
E RECORDING THE MEMORY USAGE
We did not directly check the GPU memory usage by PyTorch, instead, we used the forch.profiler module to capture the GPU memory allocation by torch following this blog") This method can show detailed memory usage through training steps. We construct the same model training and call the training procedure for a few steps including some offset computation steps. Then we use the peak memory allocation the memory usage of that method to train the model. Each experiment was done three times and reported their average.
18


Reviews:
[{"review_id": "l8hNhZpBZl", "reviewer": "Reviewer_Zuvr", "summary": "The paper proposes a method to compute local loss based on BP which significantly improves performance of local loss based training methods with negligible memory overhead.", "questions": "Is there any overhead in latency or throughput for BP-modified loss compared to other local loss-based training methods?\n\nIs the bias-variance trade-off analysis generalizable across various model architectures and tasks, or are there conditions under which it may not apply?\n\nCan you further explain how using true gradient information doesn’t increase memory usage?\n\nThe lazy update and mini-batch splitting techniques are introduced to manage memory and computational costs. How sensitive is the method to the choice of period K and mini-batch size B'?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Significantly improved performance compared to SoTA local loss based learning methods\n\nResults in negligible additional memory usage\n\nStrong theocratical foundations for proposed method", "weaknesses": "Compare throughput with SoTA local loss-based learning methods \n\nMisc: \nParaphrase line 31, 32, and 33 for better readibility", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OVyrrYqEJd", "reviewer": "Reviewer_z11R", "summary": "This paper proposes a new BP modified local loss method, which aims to adjust the local gradient by introducing an additional offset to reduce the deviation between the local gradient and the true gradient and introduce additional variance. This approach improves the performance of local loss training while maintaining memory efficiency. The theoretical analysis and experimental results of the paper show that this method can effectively improve performance on different models and tasks, while significantly reducing memory usage. This is an interesting and promising research direction with important implications for the deep learning community.", "questions": "The paper mentions hyperparameters such as offset batch size B' and sampling period K, but does not discuss the selection of these parameters in detail. It is recommended that the authors provide more analysis on how these hyperparameters affect model performance and memory usage.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Innovation: The proposed BP modification of local loss method improves training efficiency while reducing memory usage, which is a valuable contribution.\nTheoretical analysis: The paper provides a theoretical analysis based on the random modification equation, deeply explores the bias-variance trade-off, and derives the optimal scaling factor.", "weaknesses": "Impact of delayed gradient adjustment: The paper discusses the noise accumulation and bias adjustment problems that may be caused by delayed gradient adjustment, but does not provide a detailed analysis. It is recommended that the authors further study the impact of these factors on training dynamics and explore possible solutions.\nMore extensive experimental validation: Although the paper conducts experiments on multiple models and tasks, it is still recommended to conduct further experimental validation on larger models and datasets to demonstrate the universality and scalability of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5WUpmCvpXz", "reviewer": "Reviewer_EzNV", "summary": "In this paper, authors proposed a novel training algorithm called BP-modified local loss, which aims to address non-convergence issues and poorer performance in the previous local loss methods, e.g., GIM, LoCo, and Forward-Forward algorithm. It combines backward propagation (BP) with local loss methods by modifying the local gradient using the BP gradient to balance bias and variance.", "questions": "1. The method introduced several new hyper-parameters, such as batch size $B^{'}$ for the BP gradient, scaling factors $\\lambda_{h}$ and $\\alpha_{h}$, as well as the lazy update period $K$. I wonder how sensitive is the performance to these hyper-parameters, can authors report sensitivity analyses for these hyper-parameters? And what guidance or insight can authors provide on how to tune them for different models and datasets? Specifically, the optimal scaling factors, $\\lambda_{h}$ and $\\alpha_{h}$, are derived theoretically, but their practical implementation seems challenging. \n\n2. Although the authors adopted some strategies like lazy update and split mini-batch, computing the offset $Δg_{h,t}$ involves additional operations. I wonder what’s the overhead in terms of runtime and memory usage as the depth of the network increases? Can authors quantify the impact of these additional operations on the overall training efficiency? Such as total training time, per-epoch time as well as GPU memory usage as the depth of the network increases?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Strong empirical results: The BP-modified local loss shows significant improvement, e.g., up to 36% improvement in the test accuracy for Forward-Forward algorithm, 20% for LoCo and 11% for MPC.\n\n2. The use of scaling factors $\\alpha_{h}$ and $\\lambda_{h}$ to manage the bias-variance balance is well motivated. The use of stochastic differential equations to model the modified training dynamic helps in understanding the benefits, as well as the limitations of the method.", "weaknesses": "1. Generalisability is one of the major concerns, as the scope of the experiments is a bit limited. The results are primarily conducted on CIFAR-100, a relatively small dataset. While the improvements are evident, a more robust evaluation that includes larger datasets like ImageNet, as well as tasks other than image classification, would make the proposed method more convincingly.\n\n2. The authors leveraged the Ornstein-Uhlenbck process to model the bias-variance dynamic, however, the OU process is a simplified linear model, and deep neural network training is inherently non-linear and non-convex. The validity of the approximation that discussed in the paper would require some empirical evidence to support, as well as some discussions in more depth. For example, deep neural networks exhibit different training dynamics across layers, with shallow and deep layers experiencing different gradient behaviours. Can authors provide how gradients evolve for different layers and test if the OU process holds similarly for both shallow and deep layers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QzvzFz6zKK", "reviewer": "Reviewer_UZvZ", "summary": "This paper proposes a novel BP-modified local loss to improve the performance of local loss training, which decreases the bias between the BP gradient and local loss gradient but introduces additional variance. Then authors provide a theoretical analysis using the stochastic modified equation, illustrating the bias-variance trade-off and deriving optimal scaling factors. Experiment results show that the proposed method can effectively improve the performance of local loss algorithm while increase little GPU memory.", "questions": "Overall, I think this is a good paper and I would raise rating if authors could address my concerns above. In addition, I suggest adding some figure descriptions if possible so that the readers could more clearly understand the authors’ intention.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The theoretical derivation of the proposed BP-Modified Local Loss algorithm is convincing.\n- Experimental results indicate that the proposed method significantly improves the performance of local loss training algorithm.\n- The paper is highly readable and describes the implementation details of BP-Modified Local Loss algorithm in detail.", "weaknesses": "- The main motivation of local training algorithms is to train large datasets or large models in resource-constrained environments. So I think authors could scale to some datasets larger than CIFAR-100, such as the ImageNet, to better illustrate the applicability of the proposed method and meet the design motivation.\n- The optimization of the Loss function tends to significantly affect the training convergence performance. So I think it is necessary to provide model training time comparison results between the BP-Modified Local Loss method and other Local Loss baseline methods, such as total epochs needed for convergence.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)