Scaling Forward Gradient With Local Losses
Authors: Mengye Ren
Simon Kornblith
Renjie Liao
Geoffrey Hinton
Venue: ICLR 2023 poster (2023)

Abstract: Forward gradient learning computes a noisy directional gradient and is a biologically plausible alternative to backprop for learning deep neural networks. The standard forward gradient algorithm suffers from the curse of dimensionality in the number of parameters. In this paper, we propose to scale forward gradient by adding a large number of local greedy loss functions. We consider block-wise, patch-wise, and channel group-wise local losses, and show that activity perturbation reduces variance compared to weight perturbation. Inspired by MLPMixer, we also propose a new architecture, LocalMixer, that is more suitable for local learning. We find local learning can work well with both supervised classification and self-supervised contrastive learning. Empirically, it can match backprop on MNIST and CIFAR-10 and significantly outperform backprop-free algorithms on ImageNet.

Published as a conference paper at ICLR 2023
SCALING FORWARD GRADIENT WITH LOCAL LOSSES
Mengye Ren't ; Simon Kornblith", Renjie Liao", Geoffrey Hinton NYU, ?Google, 'UCC, "Vector Institute
ABSTRACT
Forward gradient learning computes noisy directional gradient and is biologi- cally plausible alternative to backprop for learning deep neural networks. However, the standard forward gradient algorithm, when applied naively, suffers from high variance when the number of parameters to be learned is large. In this paper, we propose a series of architectural and algorithmic modifications that together make forward gradient learning practical for standard deep learning benchmark tasks. We show that it is possible to substantially reduce the variance of the forward gradient estimator by applying perturbations to activations rather than weights. We further improve the scalability of forward gradient by introducing a large number of local greedy loss functions, each of which involves only small number of learnable pa- rameters, and new ML.Mixer-insppree
architecture, Lixal that is more suit- able for local learning. Our approach matches backprop on MNIST and CIFAR- 10 and significantly outperforms previously proposed backprop- free algorithms on Im- ageNet. Code is released at https/////httppettttttttbbooo googl research fraa
INTRODUCTION
Most deep neural networks today are trained using the backpropagation algorithm (a.k.a. back- prop) (Werbos 1974] LeCun 1985 Rumelhart et al. 1986), which efficiently computes the gradients of the weight parameters by propagating the error signal backwards from the loss function each layer. Although artificial neural networks were originally inspired by biological neurons, backprop has always been considered as "biologically implausible" as the brain does not form symmetric backward connections or perform synchronized computations. From an engineering perspective, backprop is incompatible with a massive level model parallelism, and restricts potential hardware designs. These concerns call for drastically different learning algorithm for deep networks.
In the past, there have been attempts to address the above weight transport problem by introducing random backward weights (Lillicrap et al.) 2016 Nøkland 2016), but they have been found to scale poorly on larger datasets such as ImageNet (Bartunov et al.] 2018). Addressing the issue of global synchronization, several papers showed that greedy local loss functions can almost as good as
end-to-end learning (Belilovsky et al.) 2019 Löwe et al.. 2019 Xiong et al. 2020). However, they still rely backprop for learning number of internal layers within each local module.
Approaches based on weight perturbation, on the other hand, directly send the loss signal back to the weight connections and hence do not require any backward weights. In the forward pass, the network adds a slight perturbation to the synaptic connections and the weight update is then multiplied by the negative change in the loss. Weight perturbation was previously proposed
biologically plausible alternative to backprop (Xie & Seung] 1999, [Seung 2003 Fiete & Seung. 2006). Instead of directly perturbing the weights, it is also possible to use forward-mode automatic differentiation (AD) to compute a directional gradient of the final loss along the perturbation direction (Pearlmutter) 1994). Algorithms based on forward- forward-mode AD have recently received renewed interest in the context of deep learning (Baydin et al. 2022; Silver et al. (2022). However, existing approaches suffer from the curse of dimensionality, and the variance of the estimated gradients is too high to effectively train large networks.
In this paper, we revisit activity perturbation (Le Cun et al. 1988 Widrow & Lehr 1990 Fiete
& Seung, 2006) as alternative weight perturbation. As previous works focused specific settings, we explore the general applicability to large networks trained on challenging vision tasks.
WWrk done visiting faculty researcher Google. Correspondence to: mengyelcs. nyu.edu.

---

Published as a conference paper at ICLR 2023
We prove that activity perturbation yields power-arrange gradient estimates than weight perturbation, and provide a continuous- tottuuuuus-tmeeeeeee rate-based interpretation of our algorithm. We directly address the scalability issue of forward gradient learning by designing an architecture with many local greedy loss functions, isolating the network into local modules and hence reducing the number of learnable parameters per loss. Unlike prior work that only adds local losses along the depth dimension, found that having patch-wise and channel group-wise losses is also critical. Lastly, inspired by the design of MLPMixer (Tolstikhin et al. (02T), we designed a network called LocalMixer, featuring a linear token mixing layer and grouped channels for better compatibility with local learning.
We evaluate our local greedy forward gradient algorithm on supervised and self-supervised image classification problems. On MNIST and C10,R-10, our learning algorithm performs comparably with backprop, and on ImageNet, it performs significantly better than other biologically plausible alternatives using asymmetric forward and backward weights. Although we have not fully matched backprop on larger-scale problems, we believe that local loss design could be a critical ingredient for biologically plausible learning algorithms and the next generation of model-parallel computation.
2 RELATED WORK
Ever since the perceptron era, the design of learning algorithms for neural networks, especially algo- rithms that could be realized by biological brains, has been a central interest. Review papers by Whhtt [ington & Bogacz (2019) and Lillicrap et al. (2020) have systematically summarized the progress of biologically plausible deep learning. Here, we discuss related work in the following subtopics.
Forward gradient and reinforcement learning. Our work leverages forward mode automatic differentiation (AD), which was first proposed by Wengerr (164)). Later it was used to learn recurrent neural networks (Williams & Zipser 1989) and to compute Hessian vector products (Pearlmutter 1994). Computing the true gradient using forward-mode AD requires the full Jacobian, which is often large and expensive to compute. Recently, [Baydin et al. (2022) and (Silver et al.. (2022) proposed to update the weights based on the directional gradient along a random or learned perturbation direction. They found that this approach is sufficient for small-scale problems. This general family of algorithms is also related to reinforcement learning (RL) and evolution strategies (ES), since in each case the network receives a global reward. RL and ES have a long history of application in neural networks (Whitley 1993 Stanley & Miikkulainen, 2002 Salimans et al. (017)) and they are effective for certain continuous control and decision decision-making tasks. Clark et al. (202T) found global credit assignment can also work well in vector neural networks where weights are only present between vectorized groups of neurons.
Greedy local learning. There have been numerous attempts to use local greedy learning objectives for training deep neural networks. Greedy layerwise pretraining (Bengio et al. 2006) Hinton et al., 2006; Vincent et al. 2010) trains individual layers modules one time to greedily optimize an objective. Local losses are typically applied to different layers or residual stages, using common supervised and self-supervised loss formulations (Belilovsky et al. 2019 Nøkland & Eidnes 2019
Löwe et al.] 2019; Belilovsky et al. (0200). Xiong et a (2020): Gomez et al. (2020) proposed to use overlapped losses to reduce the impact of greedy learning. Patel al. (2022) proposed split
network into neuron groups Laskin et al. (2020) applied greedy local learning on model parallelism training, and Wang et al. (2021) proposed add local reconstruction loss for preserving information. However, most local learning approaches proposed in the last decade rely on backprop to compute the weight updates within a local module. One exception the work of Nokladd & Eidnes (2019). which avoided backprop by using layerwise objectives coupled with similarity loss or a feedback alignment mechanism. Gated linear networks and their variants (Veness et al.) 2017 2021 [Sezener et al. I) ask every neuron to make prediction, and have shown interesting results on avoiding catastrophic forgetting. From a theoretical perspective, Baldi & Sadowski (2016) provided insights and proofs why local learning can be worse than global learning.
Asymmetric feedback weights. Backprop relies on weight symmetry: the backward weights are the same as the forward weights. rescarch has looked at whether this constraint is necessary. Lillicrap feedback et al. al. (2016) feedback alignment (FA) that uses random and fixed backward weights and found it can support error driven learning in neural networks. Direct FA (Nokland) 2016) uses single backward layer to wire the loss function back to each layer. There have also been methods that aim explicitly update backward weights. Recirculation (Hinton & McClelland 1987) and target
propagation (TP) (Bengio 2014
Lee et al. 2015
Bartunov et al. 2018) use local reconstruction
2

---

Published as a conference paper at ICLR 2023
objective to learn separate forward and backward weights approximate inverses of each other. Ladder networks (Rasmus et al.] 2015) found local reconstruction objectives and asymmetric weights can help achieve strong seml-supervised learning performance. However, [Bartunov al. (2018) reported both FA and TP algorithms do not scale to larger problems such as ImageNet, where their error rates over 90%. [Liao al. (2016); Xiao et al. (2019) proposed sign symmetry (SS) where each backward connection weight share the same sign as the forward counterpart, [Akrout et al. (2019) proposed weight mirroring and the modified Kolen-Pollack algorithm (Kolen & Pollack (994) to
align forward and backward weights. WWoo et al. (021)) proposed to update using activities from several layers below to avoid bidirectional connections. Compared to these works, we circumvent the issue of weight symmetry, and more generally network symmetry, by using only reward (and the change rate thereof), instead of backward weights.
Biologically plausible perturbation learning. Forward gradient is related to perturbation learning in the biology context. Traditionally, neural plasticity learning rules focus on deriving weight updates as function of the input and output activity of neuron (Hebb) 1949 Widrow & Hoff 1960, Oja
1982.
Bienenstock et al. 1982 Abbott& & Nelson 2000). Weight perturbation learning (Jabri & Flower
1992).
the other hand, is much more general as it permits any form of global reward (Schultz
et al.
1997). It was developed in both rated-based and spiking-based formuations (Xie & Seung
1999
Seung. 2003). Activity (or node) perturbation was proposed in shallow networks (Le Cun
et al. Scung
1988 Widrow & Lehr 1990) and later in a spike-based continuous time network (Fiete &
2006). where was interpreted as the perturbation of the conductance of neurons. |Werfel
et al. (003) showed that backprop has a faster convergence rate than perturbation learning, and
activity perturbation wins over weight perturbation by another factor. In our work, we show activity perturbation has lower gradient estimation variance compared to weight perturbation.
3
FORWARD GRADIENT LEARNING
In this section, we review and establish the technical background for our learning algorithm. We first review the technique of
automatic differentiation (AD). Second, we formulate two different types of perturbation in the weight space or activity space.
3.1 FORWARD-MODE AUTOMATIC DIFFERENTIATION (AD)
Let f RT + R". The Jacobian of f, Jy,, is a matrix of size n x m. Forward-mode AD computes the matrix-vector product Jjv, where v E RT It is defined as the directional gradient along v evaluated at x:
JAv lim f(x+6v) ov) 6v) 5-+0
8
For comparison, backprop, also known as reverse-mode AD, computes the vector-Jacobian product vJJ, where v E IR"", which corresponds to the last term in the chain rule. In contrast to reverse-mode forward pass, which is augmented with the derivative
AD, forward-mode AD only requires information. To compute the Jacobian vector product of a node in a computation graph, first the input node will be augmented with v, which is the vector be multiplied. Then for other nodes, we send in a tuple of (x, x') as inputs and compute a tuple (y, y') as outputs, where x' and y' are the dy intermediate derivatives at node and node y, i.e. y' = dx x', and dy is the Jacobian between y and x. the JAX library (Bradbury et al. 2018), forward-mode AD is implemented as jax. jup.
dx
3.2
WEIGHT-PERTUREE
FORWARD GRADIENT
Weight perturbation generate weight updates was originally explored (Barto al.) 1983 Xie & Seung 1999, Seung 2003). Baydin et al. (2022) uses the technique of forward-mode AD to implement weight perturbation, which is better than finite differences in terms of numerical stability. Let Way be the weight connection between unit i and j. and be the loss function. We can estimate the gradient by sampling a random matrix with iid elements Vig drawn from a zero-mean unit-variance Gaussian distribution. The estimator is
(2) Intuitively, this estimator samples a random perturbation direction va and tests how it aligns with the true gradient V by using forward-mon to perform the dot product, and then multiplies the scalar alignment with the perturbation direction again. Baydin et al. (2022) referred this form of gradient estimation using forward-mon -ode AD as "forward gradient". To distinguish with another form
Iw (wis) = (Eey Vueyeyy
Vij-
3

---

Published as a conference paper at ICLR 2023
Unbiased??
Avg.
Variance (shared)
Avg. Variance (independent)
(w(() (.)
Yes
po V + (pg + 1)S 2 (q + 1)5
P042 V V
S
9n(-)
Yes
S
Table 1: Comparing weight (gw) and activity (ga) perturbation. V=dimension-wise avg. gradient
variance, S=dimension-wiss avg. squared gradient norm; p=fan-in; g=fan-out; N=batch size.
of perturbation we detail later, we refer this to as "weight--prturbed forward gradient", or simply as "weight perturbation".
3.3
ACTIVITY-
FORWARD GRADIENT
An alternative to perturbing the weights to instead perturb the activities, which can reduce the num- ber of perturbation dimensions per example. Activity perturbation was originally explored in Le Cun et al. (1988); Widrow Lehr (1990) under restrictive assumptions. Here, we introduce a general way to estimate gradients using activity perturbation. It is potentially biologically plausible, since it could correspond to perturbation of the conductance in each neuron (Fiete & Seung 2006). Here, we focus on diserete-time rate- base- formulation for simplicity. Let s;; denote the activity of the i-th pre- synaptic neuron and 2j denote that of the j-th post-synaptic neuron before the non-linear activation function, and u, be the perturbation of 2j. The sertiity--prtuureed forward gradient estimator is
(w:y)) =
EE Vzyruy)) Uj.
(3)
where the inner product between Vz and is again computed by using forward-mode AD.
3.4 THEORETICAL PROPERTIES
In this section we aim to analyze the expectation and variance properties of forward gradient estimators. We focus our analysis on the gradient of one weight matrix {wis) ). but the conclusion holds for network with many weight matrices too.
Table summarizes the theoretical results] With a batch size of N, independent perturbation achieve 1/N reduction of variance, whereas shared perturbation has constant variance term dominated by the squared gradient norm. However, when performing independent weight perturbation, matrix multiplications cannot be batched because each example's s activation vector is multiplied with different weight matrix. By contrast, independent activity perturbation admits batched matrix multiplications. Moreover, activity perturbation enjoys a factor of fan-in (p) times smaller variance compared to weight perturbation since the number of perturbed elements is the number of output units instead of the size of the whole weight matrix. The only drawback of activity perturbation the memory required for storage of intermediate activations, in exchange for a factor of Np reduction variance. However, for both activity and weight perturbation, the variance still grows with larger networks. In Section 4 we will further reduce the variance by introducing local loss functions.
3.5 CONTINUOUS- TTIM RATE-BASED MODELS
Forward-mode AD can be viewed computing the first-order time derivative in a continuous-time physical system. Suppose the tuples passed between nodes of the computation graph are (x, x). where x is the change in x over time. The computation is then the same as forward-mode AD. For each node,
xx' where dy is the Jacobian between the output and the input. Note that in physical system we don't have to explicitly perform the differentiation operation by running two forward passes. Instead the first-order derivative information is readily available in the analog signal, and we only need to plug the output signal into differentiator circuit.
The activity-perturbes learning rule for continuous time system is thus whi xmtiff, where on the pre-synaptic activity, and ýJ is the rate of change in the post-synaptic activity, which is the perturbation direction for small period of time, and is the rate of change of reward (or the negative loss). The reward controls whether learning is Hebbian or anti-Hebbian. Both Hinton et al. (007)) and Bengio et (2017) propose use product of pre-synaptic activity and the rate of change of postsynaptic activity. However, they did not consider using the rate of change of reward as modulator and instead relied on another set of feedback weights communicate the error signal through inputs. In contrast, we show that by broadcasting the rate of change of reward, we can actually bypass the weight transport problem.
All proofs can be found in Appendix and Numerical simulation results be found in Appendix
4

---

Published as a conference paper at ICLR 2023
Local Losses
Local Losses
Local Losses
Channel Mixing
Token Mixing
Channel
Token Mixing
Channel
Mixing
Mixing
Figure 1: A LocalMixer network consists of several mixer blocks. A=Actvaation function (RcLU).
Mixing
Mixing
PaiDes
Nix HM c
C
N MD
H HW G
N MD
Figure 2: A LocalMixer residual block with local losses. Token mixing consists of a linear layer and channels are grouped in the channel mixing layers. Layer norm is applied before and after every linear layer. LN=Layer Norm; FC=Fully Connected layer, A=Activation function (ReLU); T=Transpose.
3.6 ACTIVATION SPARSITY AND NORMALIZATION FUNCTIONS
In networks with ReLU activations, can leverage ReLU sparsity to achieve further variance reduction, because the inactivated units will have zero gradient and therefore we should not perturb these units, and set the perturbation to be zero.
Normalization layers are often added in deep neural networks after the linear layer. To compute the correct gradient activity perturbation, we also need to account for normalization in the weight update rule. Since there is no backward weight connections, one option is to simply apply backprop on normalization layers. However, we also found that it is also fine to ignore the gradient of normalization layer when using layer normalization.
SCALING WITH LOCAL LOSSES
As we have explained in the previous section, perturbation learning can suffer from curse of dimensionality: the variance grows with the number of perturbation dimensions, and in deep networks there are often millions of parameters changing at the same time. One way to limit the number of learnable dimensions is to divide the network into submodules, each with a separate loss function. In this section, will explore several ways to increase the number of local losses to tame the variance.
1) Blockwise loss. First, we will divide the network into modules in depth. Each module consists of several layers. At the end of each module, compute loss function, and that loss is used to update the parameters in that module. This approach is equivalent of adding a "stop gradient" operator in between modules. Such local greedy losses were previously explored in Belilovsky al. (2019) and Löwe et al. (2019).
2) Patchwise loss. Sensory input signals such images have spatial dimensions. We will apply a separate loss patchwise along these spatial dimensions. In the Vision Transformer architec- ture (Vaswani et al. 2017 Dosovitskiy al.) 2021), each spatial token represents a patch in the image. In modern deep networks, parameters in each spatial location are often shared to improve data efficiency and reduce memory bandwidth utilization. Although naive weight sharing not biolog- ically plausible, we still consider shared weights in this work. It may be possible to mimic the effect of weight sharing by adding knowledge distillation (Hinton et al., 2015) losses between patches.
3) Groupwise loss. Lastly, we turn to the channel dimension. To create multiple losses, we split the channels into a number of groups, and each group attached to a loss function (Patel et al. (022). To prevent groups from communicating between each other, channels are only connected other channels within the same group. grouped linear layer is computed as žn.j = Σ; " for individual group g. Whereas previous work used channel groups to improve computational efficiency (Krizhevsky et al. 2012 [oannou et al. [2017] [ et all] (017), in our work, adding groups contributes to the total number of losses and thus reduces variance.
Feature aggregators. Naively applying losses separately to the spatial and channel dimensions leads to suboptimal performances, since each dimension contains only local information. For losses of
5

---

Published as a conference paper at ICLR 2023
A. Conventional
Avg
Patches
B. Replicated
C/G
NXHWXGXC
HW
Avg
Groups
Patches
Avg
Avg
Avg
Stack StopGradient
AvgPoal StopGradient
Figure 3: Feature aggregator designs. A) In the conventional design, average pooling is performed aggregate features from different spatial locations. B) We propose the replicated design, features are first concatenated across groups and then averaged across spatial locations. We create copies of the same feature with different stop gradient masks so that we obtain more local losses instead of global one. The stop gradient mask makes sure that perturbation in one spatial group corresponds to its loss function. The numerical value of the loss function the same as the conventional design.
standard tasks such as classification, the model needs a global view of the inputs to make a decision. Standard architectures obtain this global view by performing global average pooling layer before the final classification layer. We therefore explore strategies for aggregating information from other groups and spatial patches before the local loss function.
We would prefer to perform aggregation without reducing the total number of dimensions. We thus propose replicated design for feature aggregation, shown in Figure B First, channel groups
copied and communicated to one another, but every group except the active group itself is masked with stop gradient so that other groups do not affect the forward gradient computation:
Xp.g [StopGrad(Ip.. 1...............1 Zp StopGrad
xy,c)],
(4)
where p and index the patches and groups respectively. Similarly, each spatial location is also copied, communicated, and masked, and then averaged locally:
Xp.g P (Xpig* + Ep'pp StopGrad(xx..))
(5)
The output of feature aggregation is the same as that of the conventional global average pooling layer. The difference is that here the loss is replicated and different patch groups are activated in each loss.
Learning objectives. We consider the supervised classification loss and the contrastive InfoNCE loss (van den Oord et al. 2018 Chen et al., 2020). which are the two most commonly used losses in image representation learning. For supervised classification, we attach shared linear layer (shared across p, g) on top of the aggregated features for a cross entropy loss: IH E. 4 log softmax (.Xp..))............>>>>)))).( k. The loss is of the same value across each group and patch location.
For contrastive learning, the linear layer becomes linear feature projector. Suppose x and (2)
Xn are the two different views of the n-th example, the InfoNCE loss for contrastive learning is:
p. (Wx(1) T StoppGrad(Wxx)
Σ log E. (Wx!) p.p.)) T Lg
Orgiral
Different Perturation
Stop
15
Tariigg Lote
Note that we add stop gradient operator on the second view. It is usually unnecessary to add this stop gradient in the InfoNCE loss; however, we found that perturbation- based methods require stop gradient and otherwise the loss will not go down. This likely because we share the perturbations on both views, and having the same perturba- tion will increase the dot product between the two views but is not desired from representation learning perspective. Figure 4 shows a comparison of the loss curves. Non-shared perturbations also work but worse than stop gradient.
20000
60000
80000
Stp
Figure 4: Importance of StopGradi- ent the InfoNCE loss, using M/8 on CIFAR- 10 with 256 channels group.

---

Published as a conference paper at ICLR 2023
Type
Blocks
Patches
Channels
Groups
Params
Dataset
LocalMixer S/1/1
1x1
256
1
272K
MNIST
LocalMixer M/1/16
1x1
512
16
429K
MNIST
LocalMixer M/8/16
8x8
512
16
919K
CIFAR-10
LocalMixer L/8/64
4
8x8
2048
64
13.1M
CIFAR-10
LocalMixer L/32/64
4
32x32
2048
64
17.3M
ImageNet
Table 2: LocalMixer Architecture Details
IMPLEMENTATION
Network architecture. We propose the LocalMixer architecture that is more suitable for local learning. It takes inspiration from MLPMiixr (Tolstikhin et al., 2021), which consists of fully connected networks and residual blocks. We leverage the fully connected networks so that each spatial patch performs computations without interfering with other patches, which is more compatible with our local learning objective. An image is divided into non-overlapping patches (i.e. tokens), and each block consists of token and channel mixing layers. Figure the high level architecture, and the detailed diagram for one residual block. We add linear projector/classffcatti layer to attach a loss function the end of each block. The last layer always uses backprop to update weights. For token mixing layers, we use one linear fully connected layer instead of an MIP, since we would like to make each block as shallow as possible. Before the last channel mixing layer, features reshaped into number of groups, and the last layer is fully connected within each feature group. Table Sisowws architectural details for the different sizes of models we investigate.
Normalization.
There are many ways of performing normalization within neural network across different tensor dimensions (Krizhevsky et 2012] loffe & Szegedy 2015 Ba et al., 2016 Ren et al. [ & He|[2118) We opted for a local variant of layer normalization that normalizes within each local spatial patch of features (Ren et al. 2017). For grouped linear layers, each group is normalized separately (Wu & He, 2018). Empirically, we found such local normalization performs better on contrastive learning experiments and about the same as layer normalization on supervised experiments. Local normalization is also more biologically plausible
does not perform global communication. Conventionally, normalization layers are placed after linear layers. In MLPMixer (Tolstikhin et al. 202T), layer normalization is placed at the beginning of each residual block. We found it the best to place normalization before and after each linear layer, as shown in Figure 2 Empirically this design choice does not make much difference for backprop, but allows forward gradient learning to learn much faster and achieve lower training errors.
Efficient implementation of replicated losses. Due to the design of feature aggregation and repli- cated losses, na've implementation of groups can be very inefficient in terms of both memory consumption and compute. However, each spatial group actually computes the same aggregated feature and loss fune- tion. This means that it is possible to share most of the computation across loss functions when perform- ing both backprop and forward gradient. We imple- mented our custom JAX JVP/VJP functions (Brad? bury et al. 2018) and observed significant memory savings and compute speed-ups for replicated losses, which would otherwise not be feasible to run on mod- ern hardware. The results are reported in Figure 5. code snippet is included in Appendix [2
Natve
Fused
Navva
Fused
40
200
GPU Memory (G)
30
Compute (seclepoch)
150
20
100
10
80
Number Groups
Number Groups
Figure 5: Memory and compute usage of naïve and fused implementation of replicated losses.
EXPERIMENTS
We compare our proposed algorithm to a set of alternatives: Backprop, Feedback Alignment and other global variants of Forward Gradient. Backprop a biologically implausible oracle, since it computes true gradients whereas we compute noisy gradicnts. Feedback alignment computes approximate gradients by using a set of random backward weights. We explain each method below.
1) Backprop (BP). We include the standard backprop algorithm as well as its local variants. Local Backprop (L-BP) adds local losses as proposed, but still permits gradient to flow in an end-to-end fashion. Local Greedy Backprop (LG-BP) in addition adds stop gradient operators in between blocks. This is to provide a comparison to our methods by computing true local gradients. LG-BP is similar in spirit to recent local learning algorithms (Belilovsky et al. 2019: Löwe et al. 2019)
7

---

Published as a conference paper at ICLR 2023
Dataset
MNIST S/1/1
MNIST
CIFAR-10 M/8/16
ImageNet I/32/64
Network Metric
M/1/16
Test / Train Err. (%)
Test / Train Err. (%)
Test / Train Err. (%)
Test Train Err. (%)
BP
2.66 / 0.00
2.41 / 0..00
33.62 / 0.00
36,82 14.69
L-BP
2.38 / 0.00
2.16 / 0.00
30.75 / 0,00
42.38 / 22.80
LG-BP
2.43 / 0.00
2.81 / 0.00
33.84 / 0.05
54.37 / 39.66
BP-free algorithms
FA
2.82/ 0.00
2.90/ / 0.00
39.94 / 28.44
94.55 194
L-FA
3.21 / 0.00
2.90/ 0.00
39.74/ 28.98
87.20 / 85,69
LG-FA
3.11 0.000
2.50 / 0.00
39.73 /2.332
85,45 82.83
DFA
3.31 0.00
3.17 0.00
38.80 33.69
91.17 99.228
FG-W
9.25 9.25/8.93
8.56 8.64
55,95 54.28
97.71 997.58
FG-A
3.24 1.53
3.76 1.75
59.72 41.29
98.83 98.80
LG-FG-W
9.25 8.93
5.66 / 4.59
52.70 / 51.71
97.39 97.29
LG-FG-A
3.24/ 1.53
2.55 / 0.00
30.68 / 19.39
58.37 44.86
Table Supervised learning for image classification
Dataset Network
CIFAR-10
CIFAR-10
ImageNet L/32/64
M/8/16
L/8/64
Metric
Test / Train Err. (%)
Test Train Err. (%)
Test Train Err. (%)
BP
24.11 / 21.08
17.53 13.35
55.66 49.79
L-BP
24.69 21.80
19.13 13.60
59.11 / 52.50
LG-BP
29.63 / 25.60
23.62 16.80
68.36 / 62.53
BP-free algorithms
FA
45.87 / 44.06
67.93 65.32
82.86/ 80.21
L-FA
37.73 / 36.13
31.05 26.97
83.18 / 79.80
LG-FA
36.72/ / 34.06
30.49/ 25.56
82.57 79.53
DFA
46.09 / 42.76
39.26 37.17
93.51 92.51
FG-W
53.37 51.56
50.45 45.64
91.94 89.69
FG-A
54.59 / 52.96
56.63 56.09
97.83 / 99.79
LG-FG-W
52.66 / 50.23
52.27 / 48.66
91.36 / 88.81
LG-FG-A
32.88 / 29.73
26.81 23.90
73.24 66.89
Table 4: Self-supervised contrastive learning with linear readout
2) Feedback Alignment (FA). The standard FA algorithm (Lillicrap et al. 2016) adds a set of random and fixed backward weights. We assume that the gradients to normalization layers and activation functions are known since they do not have weight connections. Local Feedback Alignment (L-FA) adds local losses proposed, but still permits error signals flow back. Local Greedy Feedback Alignment (LG-FA) adds a stop gradient to prevent error signals from flowing back, similar the backprop-free algorithm in Nøkland & Eidnes (2019).
3) Forward Gradient (FG). This family of methods comprises our proposed algorithm and related approaches. Weithetpe forward gradient (FG-W) was proposed by [Baydin et al. (0022). In this paper, we propose the activity perturbation variant (FG-A). We further add local objective functions, producing LG-FG-W and LG-FG-A, which stand for Local Greedy Forward Gradient Weight/Activity- -eetttty-eerureed... For local perturbation to work, we have to add a stop gradient in between blocks so each perturbation has single corresponding loss. We expect LG-FG-A achieve the best performance among other variants because it can leverage the variance reduction benefit from both activity perturbation and local losses.
Datasets. We use standard image classification datasets benchmark the learning algorithms. MNIST (LeCun) (998) contains 70,000 28 handwritten digit images of class 0--9 CIFAR- 10 (Krizhessky al. 2009) contains 60,000 32x32 natural images of 10 semantic classes. Ima- geNet (Deng et al. 2009) contains 1.3 million natural images 1000 classes, which we resized to 224x224. For CIFAR-10 and ImageNet, learning. For MNIST, we applied supervised learning only. We designed different configurations of the LocalMixer architecture for each dataset, listed in Table
applied both supervised learning and contrastive
Data augmentation. For MNIST and CIFAR-10 supervised experiments, we do not apply data augmentation. Data augmentation ImageNet follows the open source implementation by Grill
8

---

Published as a conference paper at ICLR 2023
70
60
Glbbl
Dicck
Block
Pash
Block
Slock
Bldcc
Parhh
Bock
Glbbal
Block
Parth
Blokk Sloc
Crow
Peah
GHAy
Group
Pake
(a) CIFAR-I0 Supervised M/8 (b) CIFAR-I0 Contrastive M/8 (c) ImageNet Supervised L/32
Figure 6: Effect of adding local losses different locations on the performance of forward gradient
Group
Goup
Dos PMR-75
Group
Skyp
Sapa
Slpps
Sepss
(a) Supervised Test
(b) Supervised Train
(c) Contrastive Test
(d) Contrastive Train
Figure 7: Error rate of M/8/* during CIFAR-10 training using different number of groups.
et al. (2020). Because forward gradient suffers from variance, we apply weaker augmentations for contrastive learning experiments, increasing the area lower bound for random crops from 0.08 to 0.3-0.5. We find that this change has relatively little effect on the performance of backprop.
Main results. Our main results shown in Table and Table 4 In supervised experiments, there and is almost no cost of introducing local greedy losses, and our local forward gradient method can match the test error of backprop on MNIST and CIFAR. Note that LG-FG-A fails overfit the training set to 0% етгоr when trained without data augmentation. This suggests that variance could still be an issue. For CIFAR-10 10 contrastive learning, our method obtains an error rate approaching that obtained by backprop (26.8 81 % vs. 17.53%), and most of the gap is due to greedy learning vs. gradient estimation (6.09% vs. 3.19%). On ImageNet, we achieve reasonable performance compared to backprop (58.37% vs. 36.82% for supervised and 73.24% vs. 55.66% for contrastive). However, we find that the error due to greediness grows as the problem gets more complex and requires more layers to cooperate. We significantly outperform the FA family on ImageNet (by 25% for supervised and 10% for contrastive). Interestingly, local greedy FA also performs better than global feedback alignment, which suggests that the benefit of local learning transfers to other types of gradient approximation. TP-based methods were evaluated in Bartunov et al. (0018) and were found to be worse than FA on ImageNet. In sum, although there is still some noticeable gap between our method and backprop, we have made a large stride forward compared to backprop-free algorithms. More results are included in the Appendix 14]
Effect of local losses. In Figure 6 we ablate the benefit of placing local losses at different locations: blockwise, patchwise and groupwise. A combination of all three is the strongest. Global perturbation learning fails to learn as the accuracy is similar to initializing with random weights.
Effect of groups. In Figure 7 we investigate the effect of different number of groups by showing the training curves. Adding more groups bring significant improvement to local perturbation learning in terms of lowering both training and test errors, but the effect vanishes around channels group.
7 CONCLUSION
often believed that perturbation-baser learning cannot scale large and deep networks. We show that this is to some extent true because the gradient estimation variance grows with the number of hidden dimensions for activity perturbation, and is even worse for shared weight perturbation. But more optimistically, we show that a huge number of local greedy losses can help forward gradient learning scale much better. We explored blockwise, patchwise, and groupwise local losses, and a combination of all three, with a total of quarter of million losses in one of the larger networks, performs the best. Local activity perturbed forward gradient performs better than previous backprop- free algorithms larger networks. The idea of local losses opens up opportunities for different loss designs and sheds light on the search for biologically plausible learning algorithms in the brain and alternative computing devices.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENT
We thank Timothy Lillicrap for his helpful feedback on our carlier draft.
REFERENCES
Larry Abbott and Sacha B Nelson. Synaptic plasticity: taming the beast. Nature neuroscience, (11): (19:1178-1183, 2000.
Mohamed Akrout, Collin Wilson, Peter C. Humphreys, Timothy Lillicrap, and Douglas B. T'Teed. Deep learning without weight transport. In Advances in Neural Information Processing Systems 32, NeurlPS, 2019.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. abs/1607.06450, 2016.
Layer normalization.
CoRR,
Pierre Baldi and Peter Sadowski. A theory of local learning, the learning channel, and the optimality ne1
of backpropagation. Neural Networks, 83:51-74, 2016. doi:
Andrew G Barto, Richard S Sutton, and Charles W Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Transactions on Systems, Man, and Cybernetics, (5):834-846 846, 1983.
Sergey Bartunov, Adam Santoro, Blake A. Richards, Luke Marri, Geoffrey E. Hinton, and Timo- thy P. Lillicrap. Assessing the scalability of thologically-motivaie deep learning algorithms and architectures. In Advances in Neural Information Processing Systems 31, NewrIPS, 2018.
Atilim Günes Baydin, Barak Pearlmutter, Don Syme, Frank Wood, and Philip H. S. Torr. Gradients without backpropagation. CoRR, abs/2202.08587, 2022.
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to imagenet. In Proceedings of the 36th International Conference Machine Learning, ICML, 2019.
Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Decoupled greedy learning of enns. In Proceedings of the 37th International Conference on Machine Learning, ICML, 2020.
Yoshua Bengio. How auto-encoders could provide credit assignment in deep networks via target propagation. CoRR, abs/1407.7906, 2014.
Yoshua Bengio, Pascal Lamblin, Dan Popovici, and Hugo Larochelle. Greedy layer-wise training of deep networks. Advances in Neural Information Processing Systems 19, NIPS, 2006.
Yoshua Bengio, Thomas Mesnard, Asja Fischer, Saizheng Zhang, and Yuhuai Wu. Stdp presynaptic activity times rate of change of postsynaptic activity approximates back-propagation. Neural Computation, 10, 2017.
Elie L Bienenstock, Leon N Cooper. and Paul W Munro. Theory for the development of neuron selec- tivity: orientation specificity and binocular interaction in visual cortex. Journal of Neuroscience, (1):32-44 48, 1982.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
com/ www.//ooolee//////oomooo
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. Proceedings of the 37th International Conference on Machine Learning, ICML, 2020.
David G. Clark, L. F. Abbott, and Sue Yeon Chung. Credit assignment through broadcasting a global error vector. In Advances in Neural Information Processing Systems 34, NeurlPS, 2021.
10

---

Published as a conference paper at ICLR 2023
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A scage-scale hierarchical image database. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In 9th International Conference on Learning Representations, ICLR, 2021.
lla R Fiete and H Sebastian Seung. Gradient learning in spiking neural networks by dynamic perturbation conductances. Physicul review letters, 97(4):048104 2006.
Aidan N. Gomez, Oscar Key, Stephen Gou, Nick Frosst, Jeff Dean, and Yarin Gal. Interlocking backpropagation: Improving depthwise model-parrallelism CoRR, abs/2010.04116, 2020.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent A new approach to self-supervised learning. In Advances in Neural Information Processing Systems 33, NeurlPS, 2020.
Donald Olding Hebb. The organization of behavior: a neuropsyychoogicall theory. Wiley; Chapman & Hall, 1949.
Geoffrey Hinton et al. How to do backpropagation in brain. In Invited talk at the NIPS 2007 deep learning workshop, 2007.
Geoffrey E. Hinton and James L. McClelland. Learning representations by recirculation. In Neural Information Processing Systemx, 1987.
Geoffrey E. Hinton, Simon Osindero, and Yee Whye Teh. A fast learning algorithm for deep belief nets. Neural Comput., 18(7):1527-155, 2006. doi: 1116//ecco.200.18...1111
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. CoRR, hls/1503.02531,, 2015.
Yani loannou, Duncan P. Robertson, Roberto Cipolla, and Antonio Criminisi. Deep roots: Improving CNN efficiency with hierarchical filter groups. In IEEE Conference Computer Vision and Pattern Recognition, CVPR, 2017.
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Proceedings of the 32nd International Conference on Machine Learning, ICML, 2015.
Marwan Jabri and Barry Flower. Weight perturbation: An optimal architecture and learning technique for analog vlsi feedforward and recurrent multilayer networks. IEFE Transactions on Neural Networks, 3(1):154-157) 1992.
John F Kolen and Jordan B Pollack. Backpropagation without weight transport. In Proceedings of IEEE International Conference on Neural Networks, ICNN, 1994.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep con- volutional neural networks. In Advances in Neural Information Processing Systems 25, NIPS 2012.
Michael I Laskin, Luke Metz, Seth Nabarrao, Mark Saroufim, Badreddine Noune, Carlo Luschi, Jascha Sohl-Dickstein, and Pieter Abbeel. Parallel training of deep networks with local updates. CORR, abs/2012.03837, 2020.
Yann Le Cun, Conrad Galland, and Geoffrey E Hinton. Gemini: Gradient estimation through matrix inversion after noise injection. In Advances in Neural Information Processing Systems, NIPS, volume 1. Morgan-Kaufmant 1988.
11

---

Published as a conference paper at ICLR 2023
Yann LeCun. A learning scheme for asymmetric threshold networks. Proceedings of COGNITIVA, 85(537):599-666 -044 1985.
Yann LeCun. The mnist database handwritten digits, 1998. URL http::/yann.eecun. com/
exdb/mniss
Dong-Hyun Lee, Saizheng Zhang, Asja Fischer, and Yoshua Bengio. Difference target propagation. In Joint european conference on machine learning and knowledge discovery databases, pp. 515. Springer, 2015.
Qianli Liao, Joel Z. Leibo, and Tomaso A. Poggio. How important is weight symmetry in back- propagation? In Proceedings of the Thirtieth AAAI Conference on Artificial Intelligence, AAAI, 2016.
Timothy P. Lillicrap, Daniel Cownden, Douglas B. Tweed, and Colin J. Akerman. Random synaptic feedback weights support error backpropagation for deep learning. Nature Communications, 7(1): 13276, Nov 2016. ISSN 2041-1723. doi: 10. 00.1038/commms 13276.
Timothy P. Lillicrap. Adam Santoro, Luke Marris, Colin J. Akerman, and Geoffrey Hinton. Backprop- agation and the brain. Nature Reviews Neuroscience, 21(6):335-32) 346, Jun 2020. ISSN 1471-0048. doi: 1038/44/ 101001///113--00---000 1583-020-0277-33
Sindy Löwe, Peter O'Connor, and Bastiaan S. Veeling. Putting end end-to-end: Gradient- isolated learning of representations. In Advances in Neural Information Processing Systems 32, NeurIPS, 2019.
Arild Nøkland. Direct feedback alignment provides learning in deep neural networks, In Advances in Neural Information Processing Systems 29, NewrlPS, 2016.
Arild Nøkland and Lars Hiller Eidnes. Training neural networks with local error signals. In Proceedings of the 36th International Conference on Machine Learning, ICML, 2019.
Erkki Oja. Simplified neuron model principal component analyzer. Journal of mathematical biology, 157))2677- 1982.
Adeetya Patel, Michael Eickenberg, and Eugene Belilovsky, Local learning with neuron groups. In From Cells to Societies: Collective Learning Across Scales ICLR 2022 Workshop. 2022.
Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1):147-160, 1994.
Antti Rasmus, Mathias Berglund, Mikko Honkala, Harri Valpola, and Tapani Raiko, Semi-supervised learning with ladder networks. In Advances in Neural Information Processing Systems 28, NIPS, 2015.
Mengye Ren, Renjie Liao, Raquel Urtasun, Fabian H. Sinz, and Richard S. Zemel. Normalizing the normalizers: Comparing and extending network normalization schemes. In Proceedings of the 5th International Conference on Learning Representations, ICLR, 2017.
D. E. Rumelhart, G. Hinton, and R. J. Williams. Learning internal representations by error propagation. In Parallel Distributed Processing: Explorations in the Microstructure Cognition, Vol. I: Foundations, pp. 318-362, Cambridge, MA, USA, 1986. MIT Press. ISBN 026268053X.
Tim Salimans, Jonathan Ho, Xi Chen, and Ilya Sutskever. Evolution strategies as a scalable alternative to reinforcement learning. CoRR, abs/1 fbs//703.03864 2017.
Wolfram Schultz, Peter Dayan, and P Read Montague. A neural substrate of prediction and reward. Science, 155$530)):1533-1113 1599, 1997.
Sebastian Seung. Learning in spiking neural networks by reinforcement of stochastic synaptic transmission. Neuron, 40(6): 1060-10733. 2003.
Eren Sezener, Agnieszka Grabska-Barwar Barwinssk, Dimitar Kostadinov, Maxime Beau, Sanjukta Krish- nagopal, David Budden, Mareus Hutter, Joel Veness, Matthew Botvinick, Claudia Clopath, et al. A rapid and efficient learning rule for biological neural circuits. BioRxiv, 2021.
12

---

Published as a conference paper at ICLR 2023
David Silver, Anirudh Goyal, Ivo Danihelka, Matteo Hessel, and Hado van Hasselt. Learning by directional gradient descent. In Proceedings of the 10th International Conference on Learning Representations, ICLR, 2022.
Kenneth O Stanley and Risto Mikkulainen. Evolving neural networks through augmenting topologies. Evol Compud, 10(2)999-227, 2002.
Ilya O. Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucie, and Alexey Dosovitskiy. Mip-mixer: An all-mip architecture for vision. In Advances in Neural Information Processing Systems 34, NeurlPS, 2021.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs/ 01711107703353 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz. Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Information Processing Systems 30, NIPS, 2017.
Joel Veness, Tor Lattimore, Avishkar Bhoopehand, Agnieszka Grabska-Barwinn Christopher Mat- tern, and Peter Toth. Online learning with gated linear networks. arXiv preprint arXiv:1712.01897, 2017.
Joel Veness, Tor Lattimore, David Budden, Avishkar Bhoopchand, Christopher Mattern, Agnieszka Saooska-Barwinske Eren Sezener, Jianan Wang, Peter Toth, Simon Schmitt, et al. Gated linear networks. In Proceedings of the AAAI Conference Artificial Intelligence, AAAI, 2021.
Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in deep network with local denoising criterion. J. Mach. Learn. Res., 1 1-337,,3308, 2010. doi: 10.5555/1756066. 1953039,
Yulin Wang. Zanlin Ni, Shiji Song, Le Yang, and Gao Huang. Revisiting locally supervised learning: an alternative to end-to-end training. In 9th International Conference on Learning Representations ICLR, 2021.
Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger B. Grosse. Flipout: independent weight perturbations on mini-batches. In 6th International Conference on Learning Representations, ICLR, 2018.
Efficient pseudo-
R. E. Wengert. A simple automatic derivative evaluation program. Commun. ACM, 7(8):463-46 1964.
Paul Werbos. Beyond regression:" new tools for prediction and analysis in the behavioral sciences. Ph. D. dissertation, Harvard University, 1974.
Justin Werfel, Xiaohui Xie, and Seung. Learning curves for stochastic gradient descent in linear feedforward networks. Advances in Neural Information Processing Systems 16, NIPS, 2003.
L. Darrell Whitley. Genetic reinforcement learning for neurocontrol problems. Mach. Learn., 13: 259-284. 1993.
James C.R. Whittington and Rafal Bogacz. Theories of error back propaaaapann in the brain. Trends in Cognitive Sciences, 23(3):235-25 2500, Mar 2019. ISSN 1364-6613. doi: t0.10166j.tcss.2113 201..20......005. 12.005.
Bernard Widrow and Marcian E Hoff. Adaptive switching circuits. Technical report, Stanford Univ Ca Stanford Electronics Labs, 1960.
Bernard Widrow and Michael A. Lehr. 30 years of adaptive neural networks: perceptron, madaline, and backpropagation. Proc. IEEE, 1899)1415-1442, 1990.
Ronald Williams and David Zipser. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1(2):270-280, 1989.
13

---

Published as a conference paper at ICLR 2023
Sunghyeon Woo, Jeongwoo Park, Jiwoo Hong, and Dongsuk Jeon. Activation sharing with asymmet- ric paths solves weight transport problem without bidirectional connection. In Advances Neural Information Processing Systems 34, NeurlPS, 2021.
Yuxin Wu and Kaiming He. Group normalization. In 15th European Conference Computer Vision, ECCV, 2018.
Will Xiao, Honglin Chen, Qianli Liao, and Tomaso A. Poggio. biologically-plausiie learning algorithms can scale to large datasets. In Proceedings of the 7th International Conference on Learning Representations, ICLR, 2019.
Saining Xie, Ross B. Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He. Aggregated residual transformations for deep neural networks. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR, 2017.
Xiaohui Xie and H Sebastian Seung. Spike-based learning rules and stabilization persistent neural activity. Advances in Neural Information Processing Systems 12, NIPS, 1999.
Yuwen Xiong, Mengye Ren, and Raquel Urtasun. Loco: Local contrastive representation learning. In Advances in Neural Information Processing Systems 33, NeurlPS, 2020.
14

---

Published as a conference paper at ICLR 2023
8 PROOFS OF UNBIASEDNESS
In this section, show the unbiasedness of gw (wy) and ga(wi)). The first proof given by Baydin et al. (2022).
Proposition 1. gw is unbiased gradient estimator if {Usy} are independent zero-mean uni-variance random variables (Baydin et al. 2022).
Proof. We can rewrite the weight perturbation estimator as
Ihe (wer)
Vweyruy
Vaj Vwiyung +
Σ Viyy f
(7)
Note that since each dimension of u is an independent zero-m mean uni-variance random variable, Eluty] Elufy] Varle;,] + E|uy] +1+0- 1, and Bluy,0yy] = 0 Y''
E E(gn(Wwy)) } E V Wig
+E
(8)
djtij
Vwij E [1
Vury E [1
(9)
Vwiy 1+
Σ
Vwggr.
(10)
-Wii.
(11)
Proposition 2. (wig)) is an unbiased gradient estimator if {uz} are independent zero-mean uni- variance random variables.
Proof. The true gradient to the weights Vusy is the product between 2, and Vak. Therefore, rewrite the weight perturbation estimator as
(w(w))) = Iz
Σ 5'
Vzruy
= = V2ju
Vzyuz
uj
(12)
= -z;V
Σ
v. Vzgrujj
(13)
8'
uw +
Uy 8
Uj.
(14)
Since each dimension of is an independent zero-mean variinei random variable, Eluy] Elu²] = Var[u,] + E[uy] = 1 0 = 1, and =
0.
Of
(wgg( = E
Vwiyuy +
Uwiy
Ui
(15)
Vwiy E [u]]
Σ
Vwi E [uzuyl
(16)
j'
= -wii 1+ Σ
Vwiy
(17)
Vwij.
(18)
15

---

Published as a conference paper at ICLR 2023
PROOFS OF VARIANCES
We followed Wen et al. 2018 and show that the variance the gradient estimators can be decom- posed.
Lemma 1. The variance of the gradient estimator can be decomposed into three parts:
Var (g(w);))=)) = 71 Z2 + Z3, where Z1 = 1 Vi Var, (Vuigl>), 2/ N Re Var, (g(wn)))))
Z3 = - EB EiweB
Barins Cov, (ov,(g(w.;) xl)) g(w,;)) ((w,)))((mm))
Proof.
By the law of total variance,
Var (g(wi))) Var (E (o(w))))
E Var(o(wx))) [B)]
(19)
B (1
The first term comes from the gradient variance from data sampling, and it vanishes as batch size grows:
Var
E v
[o(w);)) BI)
(20)
B
Var B
E
Σ
g(wrs) ((wss)larm)
(21)
N
xin't EB
Var
E
Σ
g(wi, (zx(
(22)
N
B
v
EB EB
Var
Σ ZIWEB
E
(23)
N2
B
1
Var
Σ
Vwislatn
(24)
N2
B
Var (Vugl>>)
Var (Vwyyla) = Zz.
(25)
N2
N
The second term comes from the gradient estimation variance:
E
Var (g( [B
(26)
B
E
Var
Σ
g(wis)
x(n)
(27)
B
N
ati) B
g(wir) xm()
E
Var
(28)
B
N2
E.
Var (o(ws))l<(m))
Σ
Σ
Cor(s(u())) g(wi;)) ) (m()))
B
EB
Z m(m) (B(((())
(29)
1
E
Var [(()))))))
1 N2
E
Σ
Cov(g(mg))) s(wij) s((m))
N
( ((m) (B(((())
(30)
=Zg + Zz.
(31)
Remark. Zg is the variance of the gradient estimator in the deterministic case, and Zy measures the correlation between different gradient estimation within the batch. The Z3 is zero if the perturbations independent, and non-zero the perturbations shared within the mini-batchs
16

---

Published as a conference paper at ICLR 2023
Proposition 3. Letp q be the size of the weight matrix, the element-wise average variance of the weight perturbed gradient estimator with batch size N is 2 + (Pq + 1)8 if the perturbations are shared across the batch, and 2 V + 2 S they are independent, where V is the element-wise average variance of the true gradient, and is the element-win wise average squared gradient.
&
Proof. We first derive Zz.
Zz
E
Var (ww((we)) mx)
(32)
N
E
Var
Vuggyeey
Uij
(33)
N
8'52
E
Var
Vwija
+
(34)
N
E
Var (Wwssty)) + Var
+ 2Cov
Vwijuhg.
Σ
N
(35)
Var (Vwssezz) + Var
Σ
Vwoyeejuvy
+
(36)
N
2
E
Vuij Vwy
2E
Vwije E
y
(37)
E
Var V Var (v2))
(38)
2
Vii Vwpz E. [ugpey]] wiy E
[
V Tway E [1ayvrrr]
(39)
8
E
V V
Var (12) + Var
(40)
2
Σ
V Vwi Vway 0 - 2V 1
Vway 0
(41)
E
V Wij
Var (12) + Var
(42)
N
E
Vwiz (E[vg) - 1
Var (Vuny
(43)
N
2
E
Vw (3 (3 Blu
Σ
v Var (1ayeas))
(44)
E
[2
+ N
E
V (Var (1ys) + Blu ( (Var(Dryy) + E(uaggr]
E(e.,]]
El
12
(45)
17

---

Published as a conference paper at ICLR 2023
2
2 N
Vwij
+
Var(
+ N
E
Vuiy Var(us)) Var(wey))
(46)
N
=
2
Vuy
+
2 Var Var(Vwiylzt)
+ N
Σ
E Vwiyy
y
(47)
N
N
1
Vwij + Var (Vwis (x) +
Σ
(Welg + Var Var( (uwylal) (Voceylo))
(48)
V'y'
Zg is nonzero if the perturbations are shared within a batch. Assuming that the perturbations are shared,
Z3
1
E
Σ
Σ
Cov(go(wss)) xm)) Iw gw((w))) al(mn)
(49)
2
(B
I
Σ
Σ
E Iw (wxrs)) (d gu( xm(m)]
EE' gu (w)r)) xl
E gw (wzz) )1 mm(
2
(xf EB glml (e((n))
(50)
F N B
Σ
Σ
E
gm (wxs))] (1)) (u( xm/)
Vunsl (o) V Veney (m)
[xin)B B ammi (i 8
(51)
Σ
E
E
Σ Vueyy Moleeyyy vij
Σ 8'3
Viney la (m)) oeg
Oij
N B
[x()) x (mme Bi
(52)
Vwis (x Veisy (m))
(53)
Σ
E
E o
+
Vwey (ar byg
(54)
N2
Lafw) B Eamme My (alni)
Vuig (a (m)) m
Vuigy /o (o Virgells
(55)
E
Σ
Σ
Vwiz /x (ussla((m)
(56)
N2
[a(n) B slm Biitlw
E
Σ
Σ
E
Vwigle Pwis (m)
(57)
N2
a(()) EB x zlmee
(wigle(")
Σ
Vwez (mm)
Crper + V Wis 1z (m)
Vwey upgrrgg (58)
Σ Pf
Vwigy for (2)
Σ
Vway (ar( (m)) upggos
(59)
E E z(n) mm)
Σ Σ
Vwiy (x(
Vwiy (m3)
(60)
N2
m pl
18

---

Published as a conference paper at ICLR 2023
EEE N
Σ
Σ
E
Vuiy /o Vwig (mr( vigJ woi
(61)
B
[o())
x(())
Vwey la (eer) )
Σ Pittis
Vuyy Mattyyl Dep
vig
1
ΣΣ Vuy mofin
(62)
N²
PERE N2 B
Σ Σ
B
Vurg (olello(r)) Vuige Visglelt ui/+
(63)
(rin) EB
Vugylle Vuery lom +
Σ
Vuegy (a))) Vurey (m)) o Upa gr gon
(64)
1
Σ
Vwijj
(65)
N2
12 mfk
E
Σ
Σ
E
Vwoilae! (n) Peosslalmm) +
Σ
Vwny (x
(2) Vuig /z (m)
2
N2
B
(ninn E amme Martss
(66)
ΣΣ
Σ
Viii
(67)
N
12 mynn
E N B
Σ
Σ
E
[mg]]
+
(68)
( B z(m) Bi/tti
Vweyl (a.
Vweey (m)
E [uggrr] E
[u]j]
ΣΣ Vengs
(69)
N2
E N2 B
Σ Σ
Vuis (1)) Vwij (m)) +
Σ
Vuiyy (co( Vway (m)
(a(i)) B zl (2l
(70)
Vwis
(71)
N2
mfn
1
3 3E 3E| W +
Σ
E
V
1
ΣΣ
(72)
N
N2
mfn
mfn
Σ
3V +
Σ
Vuzy
2
Σ
V
Veij
(73)
N
NZ
n mfn
mtn
Σ 3Σ
2 2 + Σ
Vwigy
2
N(N -1) 1)
Vwis 2
+ >
Vugz
2
(74)
N2
N2
22 mfn
19

---

Published as a conference paper at ICLR 2023
Lastly, we average the variance across all weight dimensions: m Var(ow(())s)
(ar()w(ws))
(75)
pq
fj
1
{Z1 + Z2 + Zz}
(76)
Pq
ij
Var (Vwsyle)+ +
(77)
pq
N
ij
1
Vwly + Var( Vwiy(z) [z) +
Σ
Vuzg + Var (ewy 1z)
+
(78)
N
28
N(N
- 1)
Vuijz +
V 3f
(79)
N2
x'3"
Var (Vusyle) +
(80)
PI
N
ij
1
Var Var(Vuiylo))
Σ V'
Var( Var(Vweyyl)) + (vii' Σ peg'
(81)
N
x
2
In Var (Vw) + pq
m Var (Vw) + (pg + 1) MSqNorm Vw)
(82)
N
N
pq + V + (pq + 1)..
(83)
N
If the perturbations are independent, we show that Z3 is O.
Z3
1
E
Σ
Σ
Cov( (wis)) (m),,, gw (wis)) mm))
(84)
N2
2 E ml
1
E
Σ
E [u((w)s) xm)) Iw (wiy) am(m))]
- EE Mw (wx;)) [e
E Iu (we)) )1 mmm]
N2
aa E mnt) (ati))
(85)
1
E
Σ
Σ
E gu(wess) (() gw (w(s)) )l [mm]
Vwss (x) Vutss (m)
N
B
[((m) EB g(m) EB Baa(nn)
(86)
Σ Σ
E
V Wiy (cstm) (g89))
(())
Σ
Vwisr (os laltm opy
(m)
N B
[(()] B mlm) BB((n(n))
(87)
Vuig 2 (1) V Wits (m)
(88)
1
Σ
Σ
V Vuayl (1) Veigon la (ma) (n) (n)) (0) (m)
1 eingh 1ij Bij
N
B
I E mm()
EB) (B(tn(
5'
(89)
E
Σ
Σ
Vwiyla /wir (m))
(90)
2
20

---

Published as a conference paper at ICLR 2023
Σ E
E
V Wers
Wes.la/ (n)2 vig 1/m11
Veis (oms)l(mm)/
(91)
V Wwiy (x(n) big
(n)2 (m)
E
Vweyy (2lk)
Dzgy Vwiylzlm of
(m)) (m)
Vway (201)) epy
(n)
Vis
(92)
Vwigla la (m) Vuiy (x()) way (n) (1m) ( +
(93)
Σ
Vwyyla oag V Vuigy (m)) balmislgi (m) (n) (m) Vinghn Uij Vi
(94)
fy "
Vesjl (x Veosslatmm) 1)
E
Σ
(95)
N2
2(n) pmm)
n monn
Veiy (x)) Feis V (n)2 ( "i (/m)
(96)
N
Vwigy (m( Vwey (a
(n) (n)
(97)
1
E
Σ
Vwiy x (n) Vwii ( (z(m))
(98)
N
2(n) gm)m)
mynn
E
Σ
E
E
Wuish (nr Vwis (m (n)e
(m))2
(99)
N2
x(t) B n inm (Bi(a
1
Σ
Vuij
(100)
N
m
E
Σ
Σ
Vuiyla (x V Way (m)
Vwiz
E
Uig
[n)2]
E
12
(101)
N2
EB ro Bi(ala)
1
Σ Σ
Vulj
(102)
N2
1
Σ
V Vwij (x)
V Way (m)
Σ mani
ΣΣviE
(103)
N
N²
((in) EB alink (omb)BBittn (aala)))
Σ E [eillll
Σ
ΣΣ
Vwi
2
(104)
N2
N2
mfn
R mfnn
-0.
(105)
21

---

Published as a conference paper at ICLR 2023
Then the average variance becomes: m Var(gu (w)z))
1
= pq
Var(ow(wxs)
(106)
Σ 121 + L2 + Zz}
(107)
pq
Var (Vwiyla)+
(108)
pq
N
yj
1
Vwi + Var Var(VWiyyla (x) Σ
Vuiy Var( (ar(Vusrrlz))
(109)
N
d'y'
pq +
Var (w)) + pq +
mSqNorm Vw)
(110)
N
N
pq +
pq +
V
S.
(111)
N
N
Proposition 4. Let p be the size of the weight matrix, the element-wise average variance of the activity perturbed gradient estimator with a batch size N is V (g + 1)S if the perturbations are shared across the batch, and 2+2 + N S if they are independent, where V is the element-wise average variance of the true gradient, and S is the element-wiee average squared gradient.
Proof.
7:2
E.
Var (ga(wis))
(112)
N
1
E
Var a
Wyg us
(113)
N
E
Var
Vwijun +
Vuyuguy
(114)
N
j'H1
E
Var (Vwiyu?)
+ Var
Σ
Vuigruzau
+
(115)
N
J'Ms
2 Cov
V
2
(116)
3'#3
E
Var u
(Vwiyuz)
+ Var
Σ
Vuiyu
(117)
N
2
2E
Σ
Vwi, Vwi
- 2 Vwiyuz
(118)
زز
E
W Var (u?) + Var
Σ
Vuigruyy
+
(119)
j'
2
Vuiy V Wiji E. [u]uy]
IF. [1]]
Σ
V
Vwiy E [1guz]]
(120)
22

---

Published as a conference paper at ICLR 2023
Wig
Var (uz) + Var
V
Vuyy
+
(121)
N
x
2
Vwig Vwiy 1 2Vwij
1
Vwey
(122)
1
E
Var (u})
Var
V Twiyuuuy
(123)
N
1
E
F 2
(Blua)
Elug13)
+
Var (Wyyy
(124)
N
1 N
E
Vuf (3 Var(u,)²
Var(u,)*EElujf
Elug13)
+
V Vuz Var (1)24))
(125)
[2 251 +
(126)
V by' [Varluy] + Eluy. (Varfuu) + Eluy - -Elu,) E| Euu
(127)
N
2 N
Vwij
oj
+
N
2
Var(Vusylxx)
x
E
Vuiy Var(u,) Var(uy)) a Vw
(128)
2
2
W
+
Var (ussl<)))
E
(129)
N
N
N ド
1
2 Vwi + (uwiyl))) +
Σ
Vwiy
2 + Var V way Viayyla)) (a))
(130)
x
Z3 is nonzero if the perturbations are shared within a batch. Assuming that the perturbations are shared,
Zg
1
E
Σ
Σ
Cov (ov(ga()))) ao ga (a(ww))) )l mm)))
(131)
N2
B
x EB ml BB((in))
1 2
E
Σ
E ga(ww))) (())) Iu (was)) )1 (c)))
ga (al(w))) )l 2(00)) E
E gu (wiz) )1 (m))
[x()) EB ml EB [attn)
(132)
1
B
Σ
Σ
E [a(w)))) (1) Ba(ws))) mmm]
Vuss Jx. (1) (.nW)an Venay (m)
N2
B
((()) BB m(m)
EB [a(n))
(133)
Σ
Σ
E
Σ V Way lot")yy us
Σ 5'
Vwig Jar (m) ug uj
N2 B
EB mlm) BB((
(a)
(134)
Vuig (x( VWille [m
(135)
23

---

Published as a conference paper at ICLR 2023
Σ
E
E
ΣΣ
Vwiy (m)
(m) V Vaiya (x( um uy 2
(136)
N
u
[ alm)
EB) (arim)
8'
5"
Σ
E
Vwiylæ Vusylz
(137)
N
L a(m)
Σ
E
E
Vuiy (m)) Vwiy for (m))
(138)
N
EB alm) EBM
u
Fusgle(m))
Viiiy (or (o(m))yy
3 Vwrey (m
Vwig (2
(139)
j' Vwig a (m))
j'
Wey (x(
(140)
s'Hj
Σ
Vwiyla uy Veay (oe wey
(141)
""i'")
E
E
Σ
V Wail (m)
Veeey (c)(m))
(142)
N2
z(1) pr(m)
n
Σ
E
E
Vuigle ( (m)) og +
Σ
Vuigg (ar (m))
Viegy (x)
N2 B
[o(n)eB B alm) E By (z(n))
j'
(143)
Vwiy (a (m)) 1)
E
Σ
Σ
Vwiylar (1)
(144)
N
2(n) ((m)
E
Σ
Σ
E
Vwig (iocylz Vioeyloo!n oy
Vungg (mm) V Vinay /x:
uzn 2
N2
(x(N) alm) BMaiia
5'
(145)
Σ
Vwij
(146)
N2
Σ
Σ
Vusglz
Vuiyl E 14g1
V Tueyl poo (x) Vuig las (m)
E
[u3]]
E
[1]]]
N2
B
D x(m) (a(oo)))
(147)
Σ Σ
Vuiy
(148)
N
1 mæn
Σ
Σ
3 (x()) Veiiyll (m))
+
V Vwig (x) Vwiy (m())
(149)
N2
B
[nn)en B almil (
3'
ΣΣ
Vwl
(150)
N2
m
24

---

Published as a conference paper at ICLR 2023
3 3 V E
V
wiy |oo)
ΣΣ
Vuijg
2
(151)
2
x
x
N2
m
j'
m
Veigy j'
Σ
ΣΣvE'
Veij
(152)
N
N2
mef
1
Σ
Σ
Viss
Σ
Vuiy
(153)
N
mşn
8'#1
N(N- - 1)
Vwij
V
(154)
N
j'tj
Then compute the average variance across all weight dimensions (for shared perturbation):
m (ar (wi;))
Σ
(w:;))
(155)
pq
ij
1
Σ
(11 + Zz + [Za
(156)
pq
ij
Var (Vwigle)+
(157)
pq
N
ij
1
Vwiy + Var Var(Veey)zz)) + Σ
Vwiy + Var Vuny. (z)
+
(158)
N
N(N
- 1)
Vwic + Vele
(159)
N2
J'
1
Var (wigge))
(160)
pq
N
ij
1
Var (eWiillz)) (z) + Σ Var Var(WWsyy)))) +
Vwnz +
Σ Vilgi
(161)
N
x
x
2
m Var (Vw) + N
var (Vw) + (a + +1) NSQNor((VW))
(162)
N
q + 2
N + + 11S.
(163)
If the perturbations are independent, we show that Zz is O.
Z
1
R
Σ
Σ
Cov (ov(g)())))) (n)",, Da (wis) )1 mm))
(164)
N2
x EB mm)
1
E
Σ
E ga (www)) 2(1)) Iu ( (wby)) )1 mmmm
E
ga (aawws)) 2200))
E
gu (wiz) )l mm()
N
EB) ((t(n)
(165)
25

---

Published as a conference paper at ICLR 2023
Σ
Σ
E
Da (wi;))) (m) Or (u(ws;)) ) mm
Vweey /cr (1)) Veis (e
(166)
2 B
a
E B 2(m) EB ((((1))
E N B
Σ
E
Σ
Vessy (anyylafnuly ugg (2)
(m)) Σ
Vwiy (zf (mm)
ur (m)
a
( B xm EB (wr(n))
(167)
Vwigl x
V
Wiyl (m)
(168)
E
Σ
E
E
Way /xc
Veigon 1o (m)
ui Wgr
(mn) fad (m))
uj
(169)
N2
B
x. xi B ((m) EB (a((n]
5'
g
E
Σ
Σ
Vwigle /wig (wiglzzl))
(170)
N
EB m(m) (Blnn
Veiil (00rglz("///ege Vaielzlima (n)2 a u (m))2 +
E
E
(171)
N2
xi x(m) EB (x(nn)
Vuzy (2y/ (n)
(2)2 (m))
Vwiy | (m)) 2 + Veiy Seweylzl1n (mm u (m)) (n)
V Vunyy (u (m)
u
+
j'7j
5'7s
(172)
Vuigle ( Vwiyi lzc (1
uj ug
(m)) (1)) (m)) (m) (n)
(173)
E Wa uy kwiy (x (m)) uz (m) u (m)
(174)
E
Σ m=F
Vwig |x( Vuss )slms)
(175)
N2
x(n) (m))
E
Σ
E
Vioig (m) Vurey Vereyle (m) ma (r)2 u fm12++
(176)
N2
Jr EL B ((m)
EB (ar(n))
Vuiy |x
Veily (m)
(m)) (n) (m)) (n) u u u Wi
(177)
8
E E x (m) (r(m)
ΣΣ mfn
V Vwigla (o
Vwis (m())))
(178)
N
E
E ( Vwig aol (mm) (n)2
un smmp]]
(179)
N2
B
EB (mm) (a(())
Σ
W
(180)
N2
my
Σ
Vuny (n) Veng (x (m))
E
u
(m)2
E
(m))
(181)
N2
B
B B g(m) EB ((66))
Σ
Viij
(182)
N2
TI m
26

---

Published as a conference paper at ICLR 2023
Σ
Σ
Viei (x
V Wiyl (mm)
ΣΣ
2
(183)
N2
La(n) EB zim) El
E ΣΣ maten
(elsslas
1
Σ ΣΣ Vwij
(184)
N
2
N²
=0.
(185)
Then the average variance becomes:
m Var (wij))
Var(ga
(186)
Pq
Σ
(Z) Zz + Zz}
(187)
pq
Var (Vwiyl)) +
(188)
Pq
N
is
1
Vwig + Var Var(Vi Σ6 (Vi ior + Var (Veryy
(189)
N
q + 2
m Var (w)) + g+1
mSqNo Vw)
(190)
N
N
q 2
9+1 S.
+
(191)
N
10a
10a
10a
Gradlent Veriance
10a
104
perturb
10a
Ind wt perturb
10a
10
act perturb
backprop
wi perturb (heeryy)
10-1
100
10:
ind perturb Itheory)
10-2
10a
act (thcory)
bacippoop (theory)
10a
102
102
109
10
10"
N (Batch Sizel (pw4, Q=
P (Ran In) (N-4.qu
(Pan Outi IN=4, p=4)
Figure 8: Numerical verification of the theoretical variance properties
10 NUMERICAL SIMULATION OF VARIANCES
In Figure 8, we ran numerical simulation experiments verify our analytical variance properties. We used a multi-layer network with 4 input units, hidden units, 1 output unit, a (anh activation function, and the mean squared error loss. We varied the batch size (N) between and 4096. We tested the gradient estimator of the first layer weights using 5000 random samples. We also calculated the theoretical variance by applying the gradient norm and gradient variance constants found by backprop, from 5000 mini-batch true gradients. We then fixed the batch size to be 4 and vary the number of input units (p, fan in) and the number of hidden units (g, fan out) between I and 256. The theoretical variance for backprop was only computed for the batch size experiment since it is an inverse relationship ), but for fan in and fan out, we do not aim to analyze the theoretical variances here. "wt perturb" stands for weight perturbation with shared noise; "ind wt perturb" stands for weight perturbation with independent noise; and "act perturb" stands for activity perturbation with independent noise. Note that indepedent weight perturbation is much more costly compute in
27

---

Published as a conference paper at ICLR 2023
neural networks. As shown in the figure, the empirical variances match very well with our theoretical predictions.
11 TRAINING DETAILS
Here we provide more training details.
MNIST. We use a batch size of 1 128, and the SGD optimizer with learning rate 0.01 and momentum 0.9 for a total of 1000 epochs with no data augmentation and a linear learning rate decay schedule.
CIFAR-10. For the supervised experiments, we use batch size of 128 and the SGD optimizer with learning rate 0.01 and momentum 0.9 for a total of 200 epochs with no data augmentation and a linear learning rate decay schedule. For the contrastive M/8 experiments, and the SGD optimizer with learning rate 1.0 and momentum 0.9 for a total of 1000 epochs with BYO data augmentation using area crop lower bound to be 0.5 and cosine decay schedule with a warm-up period of 10 epochs. For the contrastive L/8 experiments, we use a batch size of 2048 and the SGD optimizer with learning rate 4.0 and momentum 0.9 for a total of 1000 epochs with BOL data augmentation (Grill et al. 2020) using area crop lower bound to be 0.3 and a cosine decay schedule with a warm-up period of 10 epochs.
batch size of 512
ImageNet. For the supervised experiments, we use batch size of 256 and the SGD optimizer with learning rate 0.05 and momentum 0.9 for total of 120 epochs with BYOL data augmentation (Grill) et al. 2020) using area crop lower bound to be 0.3 and cosine learning rate decay schedule with a warm-up period of 10 epochs. For the contrastive experiments, we use a batch size of 2048 and the LARS optimizer with learning rate 0.1 and momentum 0.9 for total of 800 epochs with B BOL data augmentation (Grill et al. 2020)) using area crop lower bound to be 0.08 and cosine learning rate decay schedule with a warm-up period of 10 epochs.
12 FUSED JVP/VJP DETAILS
In Algorithm 1 we provide JAX code snippet implementing fused operators for the supervised cross entropy loss. "Fused" here means that we package several operations into one function. In the supervised cross entropy loss, we combine average pooling, channel concatenation, linear classifier layer, and cross entropy all together. Key steps and expected tensor shapes are annotated in the comments. The fused InfoNCE loss implementation will be included in our full code release.
13 LOCALMIXER ARCHITECTURE
In Algorithm 2 we provide code in JAX style that implements our proposed LocalMixer architecture.
14 ADDITIONAL RESULTS
In this section we provide additional experimental results.
Normalization scheme. Table 5 compares different normalization schemes. Layer normalization (LN) is often better than batch normalization (BN) on our mixer architecture. Local LN is better on contrastive learning experiments and achieves lower error rates using forward gradient learning. Although in our main paper, backprop were used in normalization layers, backprop is not necessary for Local LN. c.f. "NG" (No Gradient) columns in Table [
Place of normalization. We investigate the places where we add normalization layers. Traditionally, normalization is added after linear layers. MLPMixer, LN added at the beginning each block. With our forward gradient learning. it is now a question of which location is the optimal design. Adding after the linear layer has the advantage of shaping the activations to more well behaved, which can make perturbation learning more effective. Adding it before the linear layer can also help reduce the variance since the inputs always get multiplied with the gradient of the output activity.
The results are reported in Table 6 Adding normalization both before and after the linear layer helps forward gradient to achieve lower training errors. While this could result in some overfitting on supervised learning, it is good for contrastive learning which needs more model capacity. This is reasonable as forward gradient introduce lot of variances, and more normalization layers help achieve better training performance.
28

---

Published as a conference paper at ICLR 2023
Algorithm 1 Naive and fused local cross entropy, with custom JVP and VJP operators.
4 N: batch, size; P: num petches; G: num gaps? C: num chancela; D: chennels gEpr K: num cls 4 encoder features [N, P,G, G, G.C/C] classitier weighta KJ; classitier blas [K]
labele: class labele [N,
import jax
import numpy as imp
from www.sccipy.ppec... import logaumexp
nat we_
lipar__ieee__eet (s, v, labels):
N, P, G, x.shape
Average pooling, with stop gradients, 6,8,G, -> [8,1,0,0/8]
avg_peoi__ mnppp.cmmmntt (x, axig=l, koepdims-Truei X X_d19_> - float
[N,
x_div_
ett.ggietoo_eeddeeee gradient pool_ x_div
Concatenate everything, with stop, gradients. 2, => 18,2, = jnp.tiles iippn.nne reshape ret.hhoppp (x, M, P, 1, G, -I]1, -I]I, (1, 1, G, 1]1
jnp. (None, stwp_srdiiitetttt x mask + x jax.1ax.stop C,C,P,C -> [,, -1j)
i, ir None]
mask) x)
reppreshape IN, logits jnp. ('ngummnnpe, npge, od-oogdd', x, #)
logits logits logaumexp (Logits, axis=-l, 1,, krepding-frue) jup.. (logitts labelal;, return 19as
None, 1), axis=-1)
def fused_avg_9t This
labels]:
forward pass. The numerical value each local should same. So we compute one and replicate it many times, N, P, G, x.shape
[/GG
x x_avg = jnp. (x, ax1a=1)
G,
[N,C]
g_p jap. imp.reshape (x_eg, (xLavg.shape(01 -11) [N -> [N,
logits jnp.einaum onkk, X_grp,
logits logits 1ogsumexp ( axis--1, keepdine-True)
loss -jnp.
labala, axig=-1)
Key step: atter computing the loss, replicate it for PXG times. +> [N,P
return jnp.tile(jap. (iosee [N, 11), [1, 2, G])
def fused_avg fusei_avg_groop Lit.earragrouppieeeaaaaaeea linear jent___e (primals, tangents):
This JVP operator parforms both regulas forward and the forward autadi X, H, b, labela primals
dx, de, db, dlabels tangeats
P, 6, D x.chope ax_avg Reshape the classifier weights, since only one group paases gradient at time. jnp. reshape (w, D, -i])
(P)
inp. reshape (b, [-1]]-
Regular forward pass
IN,P, P,C,,,C
[K,6,
XLavg inp.rean axis=1)
C,CC] -> [K
__grp rephres (X_av,, (X_avg, (haavg,) -1])
[N,C] [8,8]
logits inp.einsum('ne dk- x_gEp, w) +
logits logits loggumexp ogits, axig=-1, keepding=True)
jap.sum summsumloogg (logits labels, axis- 1)
We can compute the gradieat through cross entrepy sirst.
clogits_bie bad = cowftmxx wofttox(logitttt logitss axiss-1) labals [N,K] dioas/dlogit
step:
dloss/clogit
Do the einsum zogether to aveid replicsting outputs. éless = jmp jap.cinsum/ apg.,, gak, rk-onpg', dx dx_ev3, w alagits_ + jap.cinsum I' nk nk-on', simp.einsum ne, ck->nk', x_grp, dwl eb), dlogita_bedi [t, None, None] [N, Return loss and loss gradients [N, return jnp.tile (tt.. reshape (Loss, [N, 1, 11), I1, 2, G]), dloss
fused_avg Lood__ag__ee _og_ug linear xent vin (res, gl: This is a tuaed backprop (WPP) operator. legits, labels x.shope
F, G,
XLavg jnp.mean axta=1)
x_grp reshape (x_ag, (x_avg. -1])
Key step: only
patch/group gradients everything same.
g_ gli, 0:1, 5)
elogits q=
(logits, axis- labels) 4 [8,8]
Remember multiply gradients by På
weighs sharing.
db jrp. reshape ("p. aumdddl
(xxa=((0),, [-1]) floas IF Gl dlogits),
dw reshape (3np. einsum
[G (.n.. akz,,kk->ck"
-IJ1 tloat
Key step: use grouped weights to perform backprop. dx app.einsum né, ged dlagits, jep.reshape 16, C, -11)) / #loat aroadcast gradients dx -ap. sipp:ileeds (dxl;, None, i. t), [1, P, 11)
patches.
return dx, dM, db, None
29

---

Published as a conference paper at ICLR 2023
Algorithm 2 A LocalMixer architecture implemented with JAX style code.
inpert jax import jax.numpy
def Linear (x, K, b):
mellinear Layer. han return jnp-einaum opo, X, W]
def group_linear (x, bl:
.ansinear layer
groups, .He
return jop.einsum!' npgeinnmmmneesgooeo god >xpgd', X, 4)
normalize axis=-1, eps=le-5); "**oorma "rPoormalizatiin Lizttiion layer,
mean = m axis=axis, keepdi ma=True)
mean_of_square
jap.mean (Jnp. square (X), axis=axia, Keepdims-T:ue)
var mean ofuuuress inv = lax.lax.rsqe (var + aps) mean) return y
"ap.squase (mean)
(lockC parama):
rengnitisl block with only channel mixing.""
N, P, x.shape
num_groups
X aczmalize
linesr (x, params [0] [0], params [0] (11)
ncrmalize(x)
jax.nn.relu(x)
X [np, reshape [N, P, 6, -1]1
aormalize(x)
group sroup_lineazte (x, params (1) [01, params [1] [1]]
X (ormal17e(x)
(ax.nn.relu(x)
def mip_blcox parama):
"* "RRegular block
channel mixing.
N, P, G, x.shape
inputs
Token mixing.
X inp.reshape (x, [N, P. 1:)
acrmalize (x)
jjpp imp.swoaaxos
X lirear 1x, params [0] [0], params [1])
x ewp.awessxes (x, 1, 2)
ncrmalize
X jax.nn.relu(M)
Chanzel mixing.
1ormal17e
linear 1X, params [1] [1]
normalize_layee (X]
jax.nn.relu(M)
x jnp. www.restpppeee (x, [N, P, S, -1])
normalize(M)
group_lineas (x, pazams [01, parama [1]]
x normalize(x)
= x inputs
(ax.nn.relu(x)
return x
local_mixes parama): renLocalMixer.
preprosess (x, image_mest mean, image_std, num_patches)
pred_ [] Local predictions.
Bulid network blecks,
for bik in zange (nun_blocks):
is bik --
block0 (x,
else:
mlock params[['ttocc_(tbX\)
Projector connects to local losses.
x_rej normalize (x) pred low.loooeee.eddecccoooo
(x_proj, (arama [0], parems pf" ((1)))
Disconnect gradients.
x inp. reshape (x, (x.shapa[0), x.shape(1), -1])
inp.meant akis-1) 4 (N,C)
normallze (x)
pred linear (x, (araass/cLasssiiiee"" [00), params ciaseiiieee [1])
return pred, local
30

---

Published as a conference paper at ICLR 2023
Supervised M/8/16
Contrstive M/8/66
BP
LG-BP
LG-FG-A
LG-RG-A (NG)
ВP
LG-BP
LG-RG-A
LG-EG-A (NG)
BN
30.38 0.00
33,41/5.4
32.84/ 23.09
33,48/ 20.80
27.56/24.33
30.27 28.03
35.47
37,41/31.52
LN
0.000
33,17/ 7.09
29.03/ 17.63
29.33/ 19.26
23.52/ 20.71
27.41 / 24.73
/31.8
36.24/34.12
23.84/ 0.05
17.12
23.24/21.03
32.89/3101
32.25/30.17
Table Comparing different normalization schemes. NG=No normalization gradient. CIFAR-10 test / train error (%)
Supervised M/8/16 LG-BP
Contrastive M/8/16
LN Begin Block
BP
LG-FG-A
BP
LG-BP
LG-FG-A
31.43 / 0.00
34.73 / 1.45
31.1111.11
23.27 / 20.69
25,82 / 22.96
89.93 90.71
Before Linear
32.50/ 0.00
34.15 / 0.05
33.88
27.94
22.62 / 20,38
NaN / NaN
35.01 32.91
After Linear
30.38/ / 0.00
29.83 / 0.44
29.35
23.19
26,50 23.98
28.97 / 26.67
14.10/33..1
Before + After Lincar
33.62 / 0.00
33.84 0.05
30.68 19.39
23.24 / 21.03
28.42 / 25.20
32.89 /1101
Table 6: Place of LayerNorm CIFAR-10, test train emor. (%
40
BP Train
46
BP Train
LG-BP Train
LG-BP Tran
50
LG-FC-A Tran
LC-FC-A Trin
Eroe Ratu (h)
Test
LG-BP
Erce Patte
Texi
15
16
32
Number Groupe
Number of Groupa
(a) CIFAR- 10 Supervised M/8/*
(b) CIFAR-I0 Contrastive M/8/*
Figure 9: Effect of groups. For BP algorithms, groups has minor effect on the final performance, but for local forward gradient, it significantly reduces the variance and achieves lower error rate both training and test sets.
Effect of groups. We provide additional results summarizing the training and test performance of adding more groups in Figure(9 Backprop and local greedy backprop always achieve zero training error with increasing number groups on CIFAR-10 supervised, but adding groups has significant benefit lowering training errors for forward gradient. This suggests that the main opponent here is still the gradient estimation variance, and lowering training errors can generally make test errors lower too; on the other hand adding groups have negligible effect on backprop. For contrastive learning, here the task requires higher model capacity, and adding groups effectively reduce the model capacity by introducing sparsity in the weight matrix. As result, observe a slight drop of less than 5 5% performance on both backprop and local greedy backprop. By contrast, forward gradient gains over 10% of performance by adding 16 groups
31


Reviews:
[{"review_id": "tDTuXFcRNnD", "reviewer": "Reviewer_RQfz", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper shows that an algorithm using only forward-mode differentiation can learn well even on ImageNet-scale problems. The key to achieve this result is to use a carefully crafted neural-network architecture with many local loss functions. ", "summary_of_the_review": "While the paper may be lacking in novelty (e.g., activity perturbation has been proposed before), this is one of the first works demonstrating that backprop-free methods can learn well on ImageNet-scale problems. Such results are interesting for the community, and I therefore recommend acceptance.\n", "strength_and_weaknesses": "Strengths:\n---\nThe paper is well-written and was easy to follow, except some small details (see clarity section). \n\nTo best of my knowledge, this is the first paper which gives reasonable results on ImageNet-scale problems using a backprop-free algorithm.\n\nThe main advantage of the activity perturbation seems to be the reduced variance. Another way of variance reduction is to consider multiple perturbations and averaging them. Perhaps on a smaller dataset, an experiment could be added which shows how the results change under an increasing number of perturbations. \n\nWeaknesses:\n---\nThe method is not evaluated on standard neural architectures (or perhaps did not perform well on them?). But this does not seem like a big weakness to me, since innovations in training algorithms may require also new architectures as current DL models may be \"tuned\" to work well with backpropagation. \n\nWhile the proposed backprop-free methods can learn well, it seems still much worse than standard training on standard architectures. I am not too familiar with the proposed architecture, but note that a small MLP trained with back-prop on MNIST reaches ~1.5% test error, which can be further reduced using tricks like drop-out, etc.   I am a bit surprised that the methods have ~2.5% test error on MNIST despite 0% training error. Could this be due to overfitting? Maybe adding a weight-decay regularizer could help. \n\nIn Table 3 and 4, training error is also written in \"bold face\". Ideally, a good training algorithm has similar training and test error indicating that there is no over-fitting to the training set.  This appears to be the case for \"FG-W\" method, which however does not perform very well. \n\nMy suggestion here is simply not bold the training error, as it is not clear to me whether having a small training error is desirable.  If the computational resources are available, it would also be good to re-run the experiments over multiple random seeds and report mean and standard error.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lgFaQbzFAMR", "reviewer": "Reviewer_g2TA", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new method for training neural networks, combining forward mode auto-differentiation with directional gradients (forward gradient), weight permutation, and a novel way of using local losses. When combined with a new, complimentary architecture they propose (Local-Mixer), they show that their method is superior to other backprop-free algorithms, on a variety of classification tasks.", "summary_of_the_review": "I found the paper to be of high quality-- it is well-written, and provides novel methods that will improve the applicability of backdrop-free training.\n\nMy main concern is that section 4 was not clear enough, and I am not confident I understood how local losses (a major part of the contribution) were implemented. I intend to increase the score from 6 to 8 once all the details are clarified. (Edit: Updated to 8)", "strength_and_weaknesses": "Strengths:\n1. The paper does a good job reviewing past literature and building up to the proposed solution.\n2. The solution itself (both LG-FG-A and the arch, LocalMixer), appears to be non-trivial and interesting.\n3. For a backprop (BP) free algorithm, the results for large tasks presented in the paper are impressive. The comparison also includes BP-like methods, which helps at illustrating where the gaps in performance are coming from. \n4. The paper includes a large variety of ablation tests to showcase the benefit of different components- this is very helpful since the solution includes both a new method and a new neural network architecture, and it would be otherwise difficult to asses the empirical results. \n5. The paper closely monitors the biological plausibility of the algorithm, in line with the motivation of the paper of producing a biologically plausible method for training a neural network.\n\n\nWeaknesses:\n1. Despite its key importance to the paper, I found section 4 (Local Losses), to be too lacking in details. As is, the implementation of the local losses is not clear to me. \n    * I found the extensive use of the StopGrad operator to be confusing, and I wonder if the section would be clearer, had the paper used partial-derivative notations instead (With the StopGrad included in the implementation details). All in all, I am not sure why replicating the features with different gradient masks for different losses is better than just aggregating the losses (It's makes the algorithm more complex, but appears to be quite similar).\n   * The local losses (supervised and contrastive) use a shared linear layer. Is it shared across all local losses in a given layer? How is it trained (for either variation)? Unless I missed something, it is very not intuitive to me that a single shared weight is sufficient to give helpful information for all the different losses in the layer. I would appreciate this point being elaborated upon.\n   * The fused implementation seems to be a key to reducing the memory footprint to be scalable, but only the general concept is explained. The paper empirically compares a naive implementation with the fused implementation, without giving sufficient details as to what either of these implementations means. As such, the results in figure 5 are not reproducible.\n   * The contrastive method uses two different \"views\" for the InfoNCE loss, but I did not understand what these views are. Both views share permutation, and neither MNIST/CIFAR include data augmentation, so how are these views different from each other?\n\n2. The submission did not include code. When combined with the previous points, this raises some reproducibility concerns.\n\n3. The trade-offs for using the algorithm, and different variations of it, are not explained. E.g., Figure 5 illustrates the benefit of the \"fused\" implementation in terms of both memory and runtime, but it is not clear what the cost of adding more groups is when the fused implementation is used. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "T4bDieWsBU", "reviewer": "Reviewer_jdBC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper address the scalability issue of forward gradient learning by employing  many different local greedy loss functions (blockwise, patch-wise, and group-wise local losses, and a combination of all three). The paper shows good performance on MNIST and CIFAR-10, and also outperforms other backprop free algorithms on ImageNet. ", "summary_of_the_review": "The paper tackles an interesting problem i.e., thinking about more biologically plausible learning rules as compared to backprop. The reviewer likes the results, as well as clarity of the paper. ", "strength_and_weaknesses": "Strengths:\n\n- The paper is easy to read. \n- The paper tackles a very interesting problem.\n- The way the local losses are instantiated is interesting, as naive application of local losses does not seem to result in improvement of the results. \n\nWeaknesses:\n\n\"Silver et al. (2022) proposed to update the weights based on the directional gradient along a random perturbation direction\"\n\nThis seems a bit mis-leading. Silver et. al (2022) proposed to update the weights based on the directional gradient along a random as well as along different learned directions (as a result of truncated backprop or as a result of learned critic), which further helps to reduce the variance of the gradient. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster