Title: Block-local learning with probabilistic latent representations
Venue: NeurIPS 2023
Accepted: no
Keywords: alternative to backprop

locking problem

probabilistic models

weight transport problem
Abstract: The ubiquitous backpropagation algorithm requires sequential updates across blocks of a network, introducing a locking problem.
Moreover, backpropagation relies on the transpose of weight matrices to calculate updates, introducing a weight transport problem across blocks. Both these issues prevent efficient parallelisation and horizontal scaling of models across devices. We propose a new method that introduces a twin network that propagates information backwards from the targets to the input to provide auxiliary local losses. Forward and backward propagation can work in parallel and with different sets of weights, addressing the problems of weight transport and locking. Our approach derives from a statistical interpretation of end-to-end training which treats activations of network layers as parameters of probability distributions. The resulting learning framework uses these parameters locally to assess the matching between forward and backward information. Error backpropagation is then performed locally within each block, leading to `block-local' learning. Several previously proposed alternatives to error backpropagation emerge as special cases of our model. We present results on various tasks and architectures, including transformers, demonstrating state-of-the-art performance using block-local learning. These results provide a new principled framework to train very large networks in a distributed setting and can also be applied in neuromorphic systems.

Abstract

The ubiquitous backpropagation algorithm requires sequential updates across blocks of a network, introducing a locking problem. Moreover, backpropagation relies on the transpose of weight matrices to calculate updates, introducing a weight transport problem across blocks. Both these issues prevent efficient parallelisation and horizontal scaling of models across devices. We propose a new method that introduces a twin network that propagates information backwards from the targets to the input to provide auxiliary local losses. Forward and backward propagation can work in parallel and with different sets of weights, addressing the problems of weight transport and locking. Our approach derives from a statistical interpretation of end-to-end training which treats activations of network layers as parameters of probability distributions. The resulting learning framework uses these parameters locally to assess the matching between forward and backward information. Error backpropagation is then performed locally within each block, leading to 'block-local' learning. Several previously proposed alternatives to error backpropagation emerge as special cases of our model. We present results on various tasks and architectures, including transformers, demonstrating state-of-the-art performance using block-local learning. These results provide a new principled framework to train very large networks in a distributed setting and can also be applied in neuromorphic systems.

1 Introduction

Recent developments in machine learning have seen deep neural network architectures scaling to 21

billions of parameters [Touvron et al., 2023, Brown et al., 2020]. This development has boosted 22

the capabilities of these models to unprecedented levels but simultaneously pushed the computing 23

hardware on which large network models are running to its limits. It is therefore becoming increas24

ingly important to distribute learning algorithms over a large number of independent compute nodes. 25

However, today's machine learning algorithms are ill-suited for distributed computing. The error 26

backpropagation (backprop) algorithm requires an alternation of inter-depended forward and back27

ward phases, introducing a locking problem (the two phases have to wait for each other) [Jaderberg 28

et al., 2016a]. Furthermore, the two phases rely on the same weight matrices to calculate updates, 29

introducing a weight transport problem across blocks [Grossberg, 1987, Lillicrap et al., 2014a]. These 30

two issues make efficient parallelisation and horizontal scaling of large machine learning models 31

across compute nodes extremely difficult. 32

We propose a new method to address these problems by distributing a globally defined optimisation 33

algorithm across a large network of nodes that use only local learning. Our approach uses a message34

passing approach that uses results from probabilistic models and communicates uncertainty messages 35

forward and backwards between compute nodes in parallel. To do so, we augment a network 36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

architecture with a twin network that propagates information backwards from the targets to the input to provide uncertainty measures and auxiliary targets for local losses. Forward and backward messages comprise information about extracted features and feature uncertainties and are matched against each other using local probabilistic losses. Importantly, forward and backward propagation can work in parallel, reducing the locking problem. Inside each block, conventional error backpropagation is performed locally ('block-local'). These local updates can be used in the forward network and its backward twin for adapting parameters during training. The developed theoretical learning provides a new principled method to distribute very large networks over multiple compute nodes. The solutions emerging from this framework show striking similarities to earlier models that used random feedback weights as local targets [Lillicrap et al., 2020, Frenkel et al., 2021] but also provide a principled way to train these feedback weights.

In summary, the contribution of this paper is threefold:

We provide a theoretical framework on how interpreting the representations of deep neural networks as probability distributions provides a principled approach for block-local training of these networks. This can be used to distribute learning and inference over many interacting neural network blocks for various neural network architectures.

We demonstrate an instance of this probabilistic learning model on several benchmark classification tasks, where classifiers are split into multiple blocks and trained without end-to-end gradient computation.

We demonstrate how this framework can be used to allow deep networks to produce uncertainty estimates over their predictions. This principle is showcased on an autoencoder network that automatically predicts uncertainties alongside pixel intensity values after training.

2 Related work

Anumber of methods for using local learning in DNNs had been introduced previously. Lomnitz et al. [2022] introduced Target Projection Stochastic Gradient Descent (tpSGD), which uses layer-wise SGD and local targets generated via random projections of the labels, but does not adapt the backward weights. LocoProp [Amid et al., 2022] uses a layer-wise loss that consists of a target term and a regularizer, which is used however to enable 2nd order learning and does not focus on distributing the gradient optimization. Jimenez Rezende et al. [2016] used a generative model and a KL-loss for local unsupervised learning of 3D structures.

Some previous methods are based on probabilistic or energy-based cost functions and use a contrastive approach with positive and negative data samples. Contrastive learning Chen et al. [2020], Oord et al. [2019] can be used to construct block-local losses Xiong et al. [2020], Illing et al. [2021]. Equilibrium propagation replaces target clamping with a target nudging phase [Scellier and Bengio, 2017]. Another interesting contrastive approach was recently introduced [Hinton, 2022, Ororbia and Mali, 2023, Zhao et al., 2023]. However, it needs task-specific negative examples. [Han et al., 2018] uses a local predictive loss to improve recurrent networks' performance. In contrast to these methods, our approach does not need separate positive and negative data samples and focuses on block-local learning.

Feedback alignment [Lillicrap et al., 2020, Sanfiz and Akrout, 2021] uses random projections to propagate gradient information backwards. Jaderberg et al. [2016b] used pseudo-reward functions which are optimized simultaneously by reinforcement learning to improve performance. Random feedback alignment [Amid et al., 2022, Refinetti et al., 2021] and related approaches [Clark et al., 2021, Nøkland, 2016, Launay et al., 2020], use fixed random feedback weights to back-propagate errors. [Jaderberg et al., 2017] used decoupled synthetic gradients for local training. Target propagation demonstrates non-trivial performance with random projections for target labels instead of errors [Frenkel et al., 2021]. In contrast to these methods, we provide a principled way to adapt feedback weights.

Other methods [Belilovsky et al., 2019, Löwe et al., 2019] used greedy local, block- or layer-wise 86 optimization. Notably, Nøkland and Eidnes [2019] achieved good results by combining a matching 87 and a local cross-entropy loss. [Siddiqui et al., 2023] recently used block-local learning based on a 88 cross-correlation metric over feature embeddings [Zbontar et al., 2021], demonstrating promising 89

forward block 1

forward block 2

forward block 3

NA

a

local loss

local loss

inloss

local errorpropagation

soi nndino

local loss

local loss

NB

backwardblock1

backwardblock2

performance. [Wu et al., 2021] used greedy layer-wise optimization of hierarchical autoencoders for 90 video prediction. [Wu et al., 2022] used an encoder-decoder stage for pretraining. In contrast to these 91 methods, we do not rely solely on local greedy optimization but provide a principled way to combine 92 local losses with feedback information without locking and weight transport across blocks. 93

3 A probabilistic formulation of distributed learning 94

At a high level, our method interprets the activations of a neural network as the parameters of 95 probability distributions of latent variables. We use these intermediate representations at each block 96 to derive block local losses. These latent variables over multiple blocks implicitly define a Markov 97 chain, which allows us to tractably minimize the block's local loss. We show that the derived block 98 local losses and the resulting block local learning (BLL) are a general form of various existing local 99 losses and provide an upper bound to a global loss. 100

101

3.1 Using latent representations to construct probabilistic block-local losses

Learning in deep neural networks can be formulated probabilistically [Ghahramani, 2015] in 102 terms of maximum likelihood, i.e. the problem is to minimize the negative log-likelihood 103 L = -log p ( x , y ) = -log p ( y | x ) -log p ( x ) with respect to the network parameters θ . 104 For many practical cases where we may not be interested in the prior distribution p ( x ) , we would 105 like to directly minimize L = -log p ( y | x ) . 106

This probabilistic interpretation of deep learning can be used to define block-local losses and distribute 107 the learning over multiple blocks of networks by introducing intermediate latent representations. The 108 idea is illustrated in Fig. 1. A neural network that computes the distribution log p ( y | x ) takes x as 109 input and outputs the statistical parameters to the conditional distribution. The deep neural network 110 is split at an intermediate layer k (in Fig. 1 we used k ∈ (1 , 2) ) and end-to-end estimation of the 111 gradient is replaced by two estimators that optimize the sub-networks x → z k and z k → y separately. 112

To do this, consider the gradient of the log-likelihood loss function 113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

For any deep network, it is possible to choose any intermediate activation at layer k as latent representations z k , such that log p ( y | x ) = 〈 p ( y | z k ) p ( z k | x ) 〉 p ( z k | x , y ) , where 〈 〉 p denotes expectation with respect to p . Therefore, the representations of y depend on x only through z k as expected for a feed-forward network. Using this conditional independence property, the log-likelihood (1) expands to

This well-known result is the foundation of the Expectation-Maximization (EM) algorithm [Dempster et al., 1977]. Computing the marginal with respect to p ( z k | x , y ) corresponds to the E-step and calculating the gradients corresponds to the M-step. The sum inside the expectation separates the gradient estimators into two parts: x → z k and z k → y .

However, the E-step is impractical to compute for most interesting applications because of the combinatorial explosion in the state space of z k . To get around this, we use a variational lower bound to EM, based on the ELBO loss L V = -log p ( y | x ) + D KL ( q | p ) [Mnih and Gregor, 2014] and demonstrate that this yields a practical solution to split gradients in a similar fashion to Eq. (2). In the next section, we describe how we construct the variational distribution q .

3.2 Auxiliary latent representations

As described earlier, the output of any layer of a DNN can be interpreted as parameters to a distribution over latent random variable z k . The sequence of blocks across a network therefore implicitly defines a Markov chain x → z 1 → z 2 → . . . (see Fig. 2A). This probabilistic interpretation of hidden layer activity is valid under relatively mild assumptions, studied in more detail in the Supplement. It is important to note that the network at no point produces samples from the implicit random variables z k , but they are introduced here only to conceptualize the mathematical framework. Instead the network outputs the parameters to α k ( z k ) which is the probability distribution over z k (e.g. means and variances if α k is Gaussian). The network thus translates α k -1 → α k → . . . by outputting the statistical parameters of the conditional distribution α k ( z k ) and taking α k ( z k -1 ) parameters as input. More precisely, the network implicitly computes a marginal distribution

where 〈 〉 p denotes expectation with respect to the probability distribution p . Consequently, the network realizes a conditional probability distribution p ( y | x ) (where x and y are network inputs and outputs, respectively). And by the universal approximator property of deep neural networks, an accurate representation of this distribution can be learnt in the network weights through error back-propagation (as demonstrated for the example in Fig. 2). Eq. (3) is an instance of the belief propagation algorithm to efficiently compute conditional probability distributions.

To construct the variational distribution q we introduce the backward network N B that propagates messages β k backwards according to Eq. 4 (see Fig. 1 for an illustration). Inference over the posterior distribution p ( z k | x , y ) for any latent variable z k can be made using the belief propagation algorithm, propagating messages α k ( z k ) forward through the network using Eq. (3). In addition messages β k ( z k ) need to be propagated backward according to

such that the posterior p ( z k | x , y ) can be computed up to normalization

We make use of the fact that, through Eq. (3), the parameters of a probability distribution p ( z k | x ) are a function of the parameters to p ( z i | x ) , for 0 < i < k , e.g. if α is assumed to be Gaussian we have ( µ ( α k ) , σ 2 ( α k ) ) = f ( µ ( α i ) , σ 2 ( α i ) ) , where µ ( . ) and σ 2 ( . ) are the mean and variance of the distribution respectively. Thus, if a network outputs ( µ ( α i ) , σ 2 ( α i ) ) on layer i and ( µ ( α k ) , σ 2 ( α k ) ) on layer k , a suitable probabilistic loss function will allow the network to learn

A

4.66…...

X

>←²z←N←

↑v↑~↑

x'

μ

input

ndino

√

encoder

decoder

B

original image

X10-3

validation MSE loss

Ley

10

validationKL loss

MSE matching loss

loss

MSEestimatorloss

predicted image

5

0

error(MSEbetweenoriginalandpredicted)

0

5

10

15

20

D

predicteduncertainty(scaled)

X

f from examples. Therefore, the conditional distributions p k ( z k | z k -1 ) and the expectation in 156 Eq. (3) are only implicitly encoded in the network weights. We will study the exponential family of 157 probability distributions for which this observation can be formalized more thoroughly. 158

159

160

161

Exponential family distributions: To derive concrete losses and update rules for the forward and backward networks, we assume that α k are from the exponential family (EF) of probability distributions, given by

with base measure h , sufficient statistics T , log-partition function A , and natural parameters ϕ kj . 162 This rich class contains the most common distributions, such as Gaussian, Poisson or Bernoulli, as 163 special cases. For the example of a Bernoulli random variable we have z kj ∈ { 0 , 1 } , T ( z kj ) = z kj 164 and A ( ϕ kj ) = log ( 1 + e ϕ kj ) [Koller and Friedman, 2009]. A network directly implements an EF 165 distribution if the activations a kj encode the natural parameters, a kj = ϕ kj . Using this result, a 166 feed-forward DNN N A : x → y , can be split into N blocks by introducing implicit latent variables 167 z k : x → z k → y , and generating the respective natural parameters. In principle, blocks can be 168 separated after any arbitrary layer, but some splits may turn out more natural for a particular network 169 architecture. 170

Conveniently, if both α kj and β kj are members of the EF with natural parameters a kj and b kj , then 171 ρ kj is also EF with parameters a kj + b kj . We will use this property to deconstruct a single global 172 loss into multiple block-local losses. 173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

3.3 Illustrative example: forward-backward networks as an autoencoder

Probability representations in DNNs are useful since they provide a principled way to represent uncertainties in the network. Before we establish our main result to show how a DNN can be deconstructed into local blocks, we first demonstrate how representations of Bayesian uncertainty can emerge in DNNs by using appropriate probabilistic losses. We consider the autoencoder network illustrated in Fig. 2A and use it to learn representations for the Fashion-MNIST dataset [Xiao et al., 2017]. The CNN comprises a bottleneck layer y that implicitly splits the architecture into a decoder and encoder part (Fig. 2A). It is well known that such a network is able to learn compact representations and features that allow it to reconstruct the gray scale pixel intensities of a given input [Kingma and Welling, 2013]. Here we demonstrate that autoencoders are also able to learn representations of uncertainties, i.e. to automatically output high uncertainties for pixel values that are poorly represented in the learnt features.

To show this, we augmented the pixel representations on the inputs and outputs with additional channels that represented the logarithms of the variances of a Gaussian distribution (see Supplement for details). The input and outputs now represent the parameters of probability distributions, where the variances are proxies for the uncertainties. An appropriate loss function for this architecture is one that measures the distance between probability distributions. We used the Kullback-Leibler (KL) divergence between Gaussian distributions. This augmentation to conventional deep auto-encoders requires us to also provide uncertainty values for training data samples. Since the Fashion-MNIST dataset does not contain this information, we set the variances of pixels for all training samples to the same small constant values, reflecting high confidence (low variance) in the training set. Thus, during training, the network has only seen the same constant inputs (and outputs) for the variance channels.

Fig. 2B shows representative sample outputs for the test dataset after training. As expected, the network is able to represent the means of gray scale values in the dataset well and generalize to new images. Interestingly, the network also learned meaningful representations of the variances. Although the network has only seen constant values for the variances during training, it is able to infer information about its own uncertainty during testing. The true MSE errors between inputs and predictions qualitatively match the pixel-level variance predictions across a wide variety of inputs. For example, the network poorly represents the logo on the shirt (leftmost example) and predicts high variance in the output for these pixels. Other samples like the trousers (second from left) that are well represented correctly predict low variance. To further quantify this result, we developed additional metrics that measure the mismatch between estimated and true prediction errors (Fig. 2C, see Supplement for details). These metrics consistently decrease throughout training even though they were not directly minimized. These results suggest that DNNs are able to represent uncertainties well enough that they show zero-shot generalizations to unseen data from very limited training data.

3.4 Modularized learning using local variational losses

The autoencoder example described in Section 3.3 shows that DNNs can represent probability distributions well in principle, and also provides an idea of how probabilistic losses could be constructed locally at any layer. By 'folding' the network along the bottleneck layer y we are able to construct a sequence of pairs of auxiliary targets ( z 1 , z ′ 1 ) , ( z 2 , z ′ 2 ) , . . . (see Fig. 2D). Finally, by introducing suitable loss functions L 0 , L 1 , . . . , the mismatch between the encoder and decoder parts of the network can be minimized on a per-layer basis.

The forward and backward networks N A and N B can be used to construct local loss functions L ( k ) V at blocks k . In the Supplement, we show in detail that minimizing L ( k ) V locally and in parallel optimizes a lower bound to the log-likelihood loss L (Eq. 1), without propagating gradients end-to-end. To arrive at this result, we take the forward α k and posterior messages ρ k to be given by EF distributions with natural parameters ϕ kj and γ kj . Using this we show in the Supplement that the local loss can be optimized using the modularized gradient estimator

where µ ( · ) and σ 2 ( · ) are means and variances of EF distribution. Note that the gradients of the natural parameters ϕ kj and γ kj are computed independently and modulated by the forward and posterior weight , respectively.

The result in Eq. (7) holds for general EF distributions. For the special case of Bernoulli random 225 variables we get 226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

where a kj = f j ( a k -1 ) and b kj = g j ( b k +1 ) , are the outputs of the forward and backward network at block k ,

where m is a mixing parameter described below and S ( x ) = 1 / 1 + e -x is the sigmoid/logistic function.

The Bernoulli solution in Eq. (8) is convenient because it is a single parameter distribution (mean and variance share one parameter) such that all channels in z can be treated independently. Also the structure of Eq. 9 is well suited for a DNN implementation. In our experiments, we focus on this Bernoulli variant of the general result in Eq. (7). In the Supplement, we study a number of other relevant members of the EF. Furthermore, it is interesting to study the structure of Eq. (8) more carefully. The first term minimizes the mismatch between the forward and the posterior distribution with respect to the forward blocks. The second term is the uncertainty-weighted backward activation b kj which modulates local gradients (see Supplement). Therefore, the backward activations b kj act directly as learning signals for local updates. The BLL method is therefore related to feedback alignment [Lillicrap et al., 2020] and target propagation [Frenkel et al., 2021] where backward information is provided through random weights. However, since the gradients of the backward blocks appear in the second term, our model also provides a principled way to optimize the backward flow of information from the targets.

Data mixing schedule: The equation for the posterior distribution Eq. 9 contains a data mixing parameter m , with 0 ≤ m ≤ 1 , that scales the influence of the backward messages in the posterior distribution. This parameter serves two important functions, (1) It scales the balance between forward and backward messages in the posterior distribution ρ and (2) it scales the first term in the parameter updates Eq. 8. We found that a annealing schedule for this parameter that decreases m slowly during learning works well in practice. If not stated otherwise, we used m = (1 + τ M ) -1 in our experiments, where M is the index of the current epoch and τ is a scaling parameter (see the Supplement for further details).

4 Experimental results

We evaluated the BLL model on a number of vision and sequence learning tasks. All models used the Bernoulli BLL gradients described in Eq. (8) for local optimization. Additional details of the network models can be found in the Supplement.

4.1 Block-local learning of vision benchmark tasks

We compare the performance of our block local learning (BLL) algorithm with that of end-to-end backprop (BP) and Feedback Alignment (FA) Lillicrap et al. [2014b]. Three datasets are considered: MNIST, Fashion MNIST and CIFAR10 together with two residual network architectures [He et al., 2016]: ResNet-18 and ResNet-50, each trained with one of the three methods (BP, FA, BLL).

The BLL architectures were split into 4 blocks that were trained locally using the Bernoulli loss in Eq. (8). Splits were introduced after residual layers of the ResNet architecture by grouping subsequent layers into blocks. Group sizes were (4,5,4,5) for ResNet-18 and (12,13,12,13) for ResNet-50. Backward twin networks were here constructed simply by using the same network architecture (ResNet-18 or ResNet-50) in reverse order, introducing appropriate splits to provide intermediate targets. For CIFAR-10 gradients were propagated between two neighboring blocks (see Supplement for details and a comparison with purely local gradients). The kernels of ResNet18/ResNet-50 + FA architectures used during backpropagation are fixed and uniformly initialised following the Kaiming He et al. [2015] initialisation method. The bias is set to one.

The results are summarized in Table 1. Test top-1, top-3 and train top-1 accuracies are shown. Top-3 accuracies count the number of test samples for which the correct class was among the network's

A

B

100

add&norm

test accuracy [%]

BLL

feed

BP

forward

dense

backward

block k+1

add&norm

layer

block k+1

multi-head

attention

2

3

4

local loss

training epoch

add&norm

test accuracy [%]

100

feed

forward

dense

backward

blockk

add&norm

layer

blockk

multi-head

attention

99

5

10

15

20

number of blocks

3 highest output activations. See Supplement for results over multiple runs. BLL achieved good 272 performance on MNIST and Fashion-MNIST, closely matching end-to-end training and outperforming 273 FA networks. Note that in contrast to FA and BP, BLL does not need to compute error gradients at 274 the output but can work directly with the target labels. Performance on CIFAR-10 was significantly 275 lower than BP but outperformed FA. Interestingly the performance on the training set was close to 276 perfect for ResNet-50 suggesting over-fitting the task. 277

278

4.2 Block-local transformer architecture for sequence-to-sequence learning

279

280

281

282

283

284

285

286

287

288

289

290

291

292

Transformer architectures are in principle well suited for distributed computing due to their modular network structure that comprises a repetition of homogeneous blocks. We demonstrate a proof-ofconcept result on training a transformer with BLL. We used a transformer model with 20 self-attention blocks with a single attention head each. Block local losses were added after each layer and blocks were trained locally. A backward twin network was constructed by projecting targets through dense layers and used the Bernoulli loss Eq. (8) for local training (see Fig. 3 A for an illustration). The transformer was trained on a sequence-to-sequence task, where a random permutation of numbers 0..9 was presented on the input and had to be re-generated at the output in reverse order. We trained the network for 5 epochs.

BLL achieves convergence speed that is comparable to that of end-to-end BP on this task. Fig. 3 B shows learning curves of BLL and BP. Both algorithms converge after around 3 epochs to nearly perfect performance. BLL also achieved good performance for a wide range of network depths. Fig. 3 C shows the performance after 5 epochs for different transformer architectures. Using only 5 transformer blocks yields performance of around 99.9% (average over five independent runs). The

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

test accuracy on this task for the 20 block transformer was 99.6%. These results suggest that the BLL method is equally applicable to transformer architectures.

5 Discussion

In this work, we have demonstrated a general purpose probabilistic framework for rigorously defining block-local losses for deep architectures. This not only provides a novel way of performing distributed training of large models but also hints at new paradigms of self-supervised training that are biologically plausible. We have also shown that our block-local training approach outperforms existing local training approaches while still getting around the locking and weight transport problems. Our method introduces a twin network that propagates information backwards from the targets to the input and automatically estimates uncertainties on intermediate layers. This is achieved by representing probability distributions in the network activations. The forward network and its backward twin can work in parallel and with different sets of weights.

The proposed method may also help further blur the boundary between deep learning and probabilistic models. A number of previous models have shown that DNNs are capable of representing probability distribution [Abdar et al., 2021, Pawlowski et al., 2017, Tran et al., 2019, Malinin and Gales, 2019]. Unlike these previous methods, our method does not require Monte Carlo sampling or contrastive training, but instead exploits the log-linear structure of exponential family distributions to efficiently propagate uncertainty-aware messages through a network using a belief-propagation strategy. We have demonstrated that implicit uncertainty messages can be learnt from sparse data and accurately represent the network's performance.

Greedy block-local learning has recently shown compelling performance on a number of tasks [Nøkland and Eidnes, 2019, Siddiqui et al., 2023]. These methods use local losses with an informationtheoretic motivation but are agnostic to global back-propagating information. In future work, it may be interesting to combine these approaches with the proposed model to get the best of both worlds. Being able to produce block-level uncertainty predictions can also be useful for enhancing the sparsity of the network and using optimal amount of compute for predictions. The uncertainty predictions can also be used to handle missing labels, and for evaluating the model's confidence about its predictions. Since the framework is flexible enough to apply to self-supervised training, it can be used on unlabelled and multi-modal datasets as well. Due to the local nature of the training process, our method is particularly attractive for application on neuromorphic systems that co-locate memory and compute and use orders of magnitude less energy if the computation is local.

This work addresses potential problems of modern ML: The estimation of uncertainties in neural networks is an important open problem and understanding the underlying mechanisms better will likely help to make ML models safer and more reliable. Also the main focus of this work, which is on distributing large ML models over many compute nodes may make these model more energy efficient in the future. The energy consumption and resulting carbon footprint of ML is a major concern and the proposed model may provide a new direction to approach this problem. This method may enable training of larger models which also come with associated risks in terms of biases and inappropriate use in the real world. It is also not known what biases using this method itself and extensions with sparsity may introduce in the models predictions.

References

Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U Rajendra Acharya, et al. A review of uncertainty quantification in deep learning: Techniques, applications and challenges. Information Fusion , 76:243-297, 2021.

Ehsan Amid, Rohan Anil, and Manfred Warmuth. LocoProp: Enhancing BackProp via local loss optimization. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics , pages 9626-9642. PMLR, 2022. URL https://proceedings.mlr.press/v151/ amid22a.html .

Eugene Belilovsky, Michael Eickenberg, and Edouard Oyallon. Greedy layerwise learning can scale to ImageNet. In Proceedings of the 36th International Conference on Machine Learning , pages 583593. PMLR, 2019. URL https://proceedings.mlr.press/v97/belilovsky19a.html .

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. arXiv:2005.14165 [cs] , July 2020. URL http://arxiv.org/abs/2005.14165 .

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

David Clark, L F Abbott, and Sueyeon Chung. Credit assignment through broadcasting a global error vector. In Advances in Neural Information Processing Systems , volume 34, pages 10053-10066. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/hash/ 532b81fa223a1b1ec74139a5b8151d12-Abstract.html .

P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incomplete data via the EM algorithm. 39(1):1-22, 1977. ISSN 00359246. doi: 10.1111/j.2517-6161.1977.tb01600.x. URL https://onlinelibrary.wiley.com/doi/10.1111/j.2517-6161.1977.tb01600.x .

Charlotte Frenkel, Martin Lefebvre, and David Bol. Learning without feedback: Fixed random learning signals allow for feedforward training of deep neural networks. 15, 2021. ISSN 1662453X. URL https://www.frontiersin.org/articles/10.3389/fnins.2021.629892 .

Zoubin Ghahramani. Probabilistic machine learning and artificial intelligence. Nature , 521(7553): 452-459, 2015.

Stephen Grossberg. Competitive learning: From interactive activation to adaptive resonance. Cognitive science , 11(1):23-63, 1987.

Kuan Han, Haiguang Wen, Yizhen Zhang, Di Fu, Eugenio Culurciello, and Zhongming Liu. Deep predictive coding network with local recurrent processing for object recognition, 2018. URL http://arxiv.org/abs/1805.07526 .

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

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

Bernd Illing, Jean Ventura, Guillaume Bellec, and Wulfram Gerstner. Local plasticity rules can learn deep representations using self-supervised contrastive predictions. In Advances in Neural Information Processing Systems , volume 34, pages 30365-30379. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/paper/2021/hash/ feade1d2047977cd0cefdafc40175a99-Abstract.html .

Max Jaderberg, Wojciech Marian Czarnecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled Neural Interfaces using Synthetic Gradients. arXiv:1608.05343 [cs] , August 2016a. URL http://arxiv.org/abs/1608.05343 .

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks, 2016b. URL http://arxiv.org/abs/1611.05397 .

Max Jaderberg, Wojciech Marian Czarnecki, Simon Osindero, Oriol Vinyals, Alex Graves, David Silver, and Koray Kavukcuoglu. Decoupled neural interfaces using synthetic gradients. In Proceedings of the 34th International Conference on Machine Learning , pages 1627-1635. PMLR, 2017. URL https://proceedings.mlr.press/v70/jaderberg17a.html .

394

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

Danilo Jimenez Rezende, S. M. Ali Eslami, Shakir Mohamed, Peter Battaglia, Max Jaderberg, and Nicolas Heess. Unsupervised learning of 3d structure from images. In Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016. URL https://proceedings. neurips.cc/paper/2016/hash/1d94108e907bb8311d8802b48fd54b4a-Abstract.html .

Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. arXiv:1312.6114 [cs, stat] , December 2013. URL http://arxiv.org/abs/1312.6114 .

Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques . MIT press, 2009.

Julien Launay, Iacopo Poli, François Boniface, and Florent Krzakala. Direct feedback alignment scales to modern deep learning tasks and architectures. In Advances in Neural Information Processing Systems , volume 33, pages 9346-9360. Curran Associates, Inc., 2020. URL https://proceedings. neurips.cc/paper/2020/hash/69d1fc78dbda242c43ad6590368912d4-Abstract.html .

Timothy P. Lillicrap, Daniel Cownden, Douglas B. Tweed, and Colin J. Akerman. Random feedback weights support learning in deep neural networks. arXiv:1411.0247 [cs, q-bio] , November 2014a. URL http://arxiv.org/abs/1411.0247 .

Timothy P. Lillicrap, Daniel Cownden, Douglas B. Tweed, and Colin J. Akerman. Random feedback weights support learning in deep neural networks, 2014b.

Timothy P. Lillicrap, Adam Santoro, Luke Marris, Colin J. Akerman, and Geoffrey Hinton. Backpropagation and the brain. 21(6):335-346, 2020. ISSN 1471-0048. doi: 10.1038/s41583-020-0277-3. URL https://www.nature.com/articles/s41583-020-0277-3 .

Michael Lomnitz, Zachary Daniels, David Zhang, and Michael Piacentino. Learning with local gradients at the edge, 2022. URL http://arxiv.org/abs/2208.08503 .

Sindy Löwe, Peter O' Connor, and Bastiaan Veeling. Putting an end to end-to-end: Gradient-isolated learning of representations. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ 851300ee84c2b80ed40f51ed26d866fc-Abstract.html .

Andrey Malinin and Mark Gales. Reverse kl-divergence training of prior networks: Improved uncertainty and adversarial robustness. Advances in Neural Information Processing Systems , 32, 2019.

Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning , pages 1791-1799. PMLR, 2014.

Arild Nøkland and Lars Hiller Eidnes. Training neural networks with local error signals. In International conference on machine learning , pages 4839-4850. PMLR, 2019.

Arild Nøkland. Direct feedback alignment provides learning in deep neural networks. In Advances in Neural Information Processing Systems , volume 29. Curran Associates, Inc., 2016. URL https://proceedings.neurips.cc/paper/2016/hash/ d490d7b4576290fa60eb31b5fc917ad1-Abstract.html .

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019. URL http://arxiv.org/abs/1807.03748 .

Alexander Ororbia and Ankur Mali. The predictive forward-forward algorithm. arXiv preprint arXiv:2301.01452 , 2023.

Nick Pawlowski, Andrew Brock, Matthew CH Lee, Martin Rajchl, and Ben Glocker. Implicit weight uncertainty in neural networks. arXiv preprint arXiv:1711.01297 , 2017.

Maria Refinetti, Stéphane d'Ascoli, Ruben Ohana, and Sebastian Goldt. Align, then memorise: the dynamics of learning with feedback alignment, 2021. URL http://arxiv.org/abs/2011. 12428 .

Albert Jiménez Sanfiz and Mohamed Akrout. Benchmarking the accuracy and robustness of feedback alignment algorithms, 2021. URL http://arxiv.org/abs/2108.13446 .

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

Benjamin Scellier and Yoshua Bengio. Equilibrium propagation: Bridging the gap between energybased models and backpropagation. Frontiers in computational neuroscience , 11:24, 2017.

Shoaib Ahmed Siddiqui, David Krueger, Yann LeCun, and Stéphane Deny. Blockwise self-supervised learning at scale, 2023. URL http://arxiv.org/abs/2302.01647 .

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and Efficient Foundation Language Models, February 2023. URL http://arxiv.org/abs/2302.13971 .

Dustin Tran, Mike Dusenberry, Mark Van Der Wilk, and Danijar Hafner. Bayesian layers: A module for neural network uncertainty. Advances in neural information processing systems , 32, 2019.

Bohan Wu, Suraj Nair, Roberto Martin-Martin, Li Fei-Fei, and Chelsea Finn. Greedy hierarchical variational autoencoders for large-scale video prediction. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2318-2328, 2021.

Kan Wu, Jinnian Zhang, Houwen Peng, Mengchen Liu, Bin Xiao, Jianlong Fu, and Lu Yuan. TinyViT: Fast pretraining distillation for small vision transformers, 2022. URL http://arxiv.org/abs/ 2207.10666 .

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747 , 2017.

Yuwen Xiong, Mengye Ren, and Raquel Urtasun. LoCo: Local contrastive representation learning. In Advances in Neural Information Processing Systems , volume 33, pages 11142-11153. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/hash/ 7fa215c9efebb3811a7ef58409907899-Abstract.html .

Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction, 2021. URL http://arxiv.org/abs/2103.03230 .

Gongpei Zhao, Tao Wang, Yidong Li, Yi Jin, Congyan Lang, and Haibin Ling. The cascaded forward algorithm for neural network training. arXiv preprint arXiv:2303.09728 , 2023.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe present work proposes a block-wise learning strategy, whereby the architecture is split into several blocks, with each block receiving an error signal stemming from a local (block-wise) loss. As this technique makes use of a parametrized twin network to compute these error signals, it also bypasses the so-called \"weight transport problem\" of the standard backprop algorithm. The motivation of this work is the hardware-friendliness of the resulting algorithm.\n\nMore precisely:\n\n- Section 3.1 formulates the learning objective in a probabilistic fashion (likelihood maximization) and how to do it (variational inference)\n- Section 3.2 explains the parametrization of the base model $p(z|x) := \\alpha(z)$ to be trained and that of the variational posterior $q(z|x, y)$. The variational posterior itself decomposes into the product of  bottom-up top-down messages: $q(z|x, y) \\propto p(z|x) p(y|z) := \\alpha(z) \\times \\beta(z)$. The base model and the variational posterior are parametrized within the *exponential family* and the neural networks are made to output the natural parameters associated with these distributions. No stochastic quantity is ever propagated through the nets -- e.g. it does not operate as a VAE for instance.\n- Section 3.3 grounds the previous idea in an encoder-decoder setting where the last layer of the decoder is augmented with extra channels to output pixel-wise variances in the image space. Training this architecture on F-MNIST with the proposed approach (the learning objective is not yet introduced at this stage), it is shown that the resulting uncertainties are good proxy to the reconstruction error. \n- Section 3.4 introduces the proposed local (block-wise) learning objectives -- the derivation is in Appendices 1.2 -> 1.3.4. Heuristically: the learning rule reads as the forward/posterior mismatch backpropagated into the feedforward block (first term and second term of Eq. 7) and its feedback counterpart (second term of Eq. 7) through the natural parameters gradients. Importantly, an heuristic called \"data mixing\" is introduced to midly take into account the residual terms of the upper bound of the loss (details in Appendix).\n- Section 4.1 presents results obtained on MNIST, F-MNIST and CIFAR-10 with the proposed technique on ResNet architectures, benchmarked against feedback alignment and standard backprop. The proposed approach performs comparably with backprop on MNIST and F-MNIST, slightly outperforms FA on CIFAR-10 but is considerably degraded with respect to backprop.\n- Section 4.2 finally presents results on a 20-layers deep transformer architecture on a toy task (reverting a random permutation of numbers ranging between 0 and 9) where the proposed approach is shown to perform comparably with backprop.\n\nSTRENGTHS:\n- The paper tackles an interesting problem: block-wise local training casted into a probabilistic setting.\n- The idea of the paper is interesting: the starting point is the same as Predictive Coding (Whittington & Bogacz, 2017), but: i/ picking a different variational family ii/ amortizing inference with a single forward pass (rather than minimizing energy for the E-step at each batch iteration) iii/ having a whole block of layers (where backprop applies) to compute the parameters of the distributions. I like the idea of stitching several algorithms together to solve a problem.\n\nWEAKNESSES:\n- The writing and the structure of the paper make it extremely difficult to deeply understand the proposed approach. To my eyes, ideas do not appear in the right order in the main, important ideas are in the appendix and the notations are confusing, important details are missing.\n- It is *not* true that the algorithm parallelizes the forward and backward pass (L.40: \"forward and backward propagation can work in parallel\"). The underlying algorithm is just variational inference applied to a generative model conditioned on a given input $x$. I do not see any valid (theoretically grounded) argument as to why the first block could start processing a novel input $x'$ while the upper blocks process an input $x$. **The block-wise locality of the learning rule does not suffice as an argument here**.\n- In the same vein, it is *neither* true that the algorithm allows for the parallelization the backward pass *across different blocks*. Here again, as the underlying algorithm simply is variational inference, I see no valid argument as to why each (feedforward) block could do without having a top-down error signal. However, it is true that the training of the *feedback* parameters can be parallelized.\n- The derivation of the learning rule -- which should be, to my eyes, the *central piece* of the paper -- is unfortunately not presented in a sufficiently clear and detailed fashion. \n- Section 3.3 on uncertainty estimation is weak and orthogonal to the scope of the paper in my eyes. The uncertainty estimation is carried out on a single simple task, and does not abide by the standard of the uncertainty estimation literature: are the uncertainties well calibrated? Can they be used for anomaly detection? What is the quantitative performance of the anomaly detection in terms of binary classification metrics? A good task to consider (and not too difficult) is the MVTech dataset (https://www.mvtec.com/company/research/datasets/mvtec-ad).\n- The baselines chosen in the experimental section are not very relevant. The proposed technique applies to block-wise training, but not a single block-wise training baseline (e.g. Belilovsky (2019)) is considered. Why? A relevant choice would have been to consider a VGG-11 architecture with 3 layers per block and try to reach $\\approx 67.6 \\\\%$ top-1 performance on ImageNet (Belilovsky et al, 2019: https://arxiv.org/pdf/1812.11446.pdf).\n- Given that the use of backprop is allowed within each block and that ResNet architectures are considered, the experimental results are disappointing: the performance obtained on CIFAR-10 is very poor ($\\approx 70 \\\\%$ accuracy on CIFAR-10 is achievable by a 4 layers-deep convnet), ImageNet32 (or even ImageNet) has not been considered, and some results are surprising (see one of the questions below).\n\nQUESTIONS:\n- It is more a suggestion for clarity improvement than a question: for the sake of clarifying the derivation of the learning rule, could you consider presenting things **in the main** closer to the following:\n  + $-\\log p(y|x) \\leq L_b^k  (^*)$ with $L_b^k:=-\\mathbb{E}_q[\\log p(z^k, y|x)] - H(q(z^k |x,y)) = -\\log p(y|x) + KL(q(z^k |x,y) || p(z^k|x,y))$.\n  + Sum $(^*)$ over k, divide per $N$ to obtain the upper bound $L_b := -\\log p(y|x)+ \\frac{1}{N}\\sum_{k=1}^N KL(q(z^k |x,y) || p(z^k|x,y))$.\n  + Re-write $L_b$ as: $L_b=\\frac{1}{N}\\sum_{k=1}^N\\ell(x, y, z^k)$ with $\\ell(x, y, z^k) := \\mathbb{E}_{q(z^k|x,y)}\\left[\\log \\frac{q(z^k|x,y)}{p(y|z^k)p(y|z^k)}\\right]$\n  + Finally rewrite $\\ell(x, y, z^k) = KL(q(z^k|x,y) || p(z^k|x)) - \\mathbb{E}_{q(z^k|x,y)}\\left[\\log p(y|z^k)\\right]$\n\nFor next questions onward, we denote $\\ell^{(1)}(x, y, z^k) := KL(q(z^k|x,y) || p(z^k|x))$ and $\\ell^{(2)}(x, y, z^k) := - \\mathbb{E}_{q(z^k|x,y)}\\left[\\log p(y|z^k)\\right]$ such that $\\ell(x, y, z^k) = \\ell^{(1)}(x, y, z^k) + \\ell^{(2)}(x, y, z^k)$.\n\n- In the light of the previous remark, I don't think it is optimal for clarity purposes to directly give the formula in Eq. (7). Also, Eq. 7 is essentially skewed as you simply write $\\partial_\\theta \\ell^{(1)}(x, y, z^k)$ and *not* $\\partial_\\theta \\ell (x, y, z^k)$ (i.e. the total ELBO).\nAnother thing which is very unclear is that you make no distinction between the parameter of the *base network* and that of the *target* / feedback network. In L. 104, you define $\\theta$ as the \"network parameters\", but in Eq. 8, you are explicitly taking gradient of the output of the feedback network with respect to $\\theta$. There is an ambiguity here that hinders clarity. Please distinguish between the two sets of parameters.\n\n- Following up on the previous bullet: there is *no* explanation in the main of why you are discarding $\\ell^{(2)}(x, y, z^k)$ from the ELBO gradient. It only when looking at 1.3.4 inside the supplementary material that we understand that: 1/$\\ell^{(2)}(x, y, z^k)$ is intractable 2/ \"data mixing\",which appears to be a heuristic optimization trick in the main, is in fact intended to approximate/emulate the intractable contribution of $\\ell^{(2)}(x, y, z^k)$. Even after reading multiple times L. 72-80, I still misunderstand this trick. If I missed something important, please clarify it inside the main. \n\n- I'm doubtful about the experimental results: why does the proposed technique overfits so well CIFAR-10 while being so poor at generalization, in spite of using a ResNet? Do you train all layers (4 layers would be enough to overfit CIFAR-10)? Do you not apply optimization tricks to avoid overfitting (i.e. weight decay, dropout, data augmentation)? This is very surprising.\n\n- One other question which is absent in the paper (perhaps hinted by Fig. 1) is how **the last block**, with the classification error signal, is trained? My assumption is that it is trained by mere backprop.\n\n- **A potential interpretation for your CIFAR-10 results**. If the previous point holds true, my interpretation of your surprising CIFAR-10 results is that you might end up training *only* the last block (4 layers for ResNet-18), which is sufficient to overfit CIFAR-10, but generalizes as a 4 layers-deep architecture (consistently with my remark above) because the error signal received by previous layers might be irrelevant. To sanity-check this, I would suggest performing a low-dimension project (e.g. using t-SNE) of the activations of the last layer of the penultimate block (e.g. $a_2$ on Fig. 1) and visualize whether the classes are well separated. I assume (but perhaps I'm wrong) they are not.\n\n- **Why the previous blocks might not learn?** This is something to be investigated further. I see several possibilities:\n   + The \"data mixing\" trick is too heuristic and does not suffice to \"emulate\"/take into account the intractable part of your ELBO gradient.\n   + If you indeed parallelize gradient computation *across different blocks*, it might be that some blocks never receive top-down error information.\n\n- A really minor point at this stage would be to mention Predictive Coding (Whittington & Bogacz, 2017), whose starting point is the same as yours (e.g. variational inference on a generative model), but using a much simpler variational family.\n\n- Coming back on uncertainty estimation: I'm not sure it is a desirable property that the model exhibit high pixel-wise variance for *in-distribution* features. A more desirable property rather is to have high uncertainty on *out-of-distribution* features -- namely: segmentation anomalies/defects. Although I still think this direction is orthogonal to the scope of the paper and should be removed, if this is of interest to you, consider the MVTech dataset and see if the pixel uncertainties can be leveraged to detect object anomalies. \n\n- Another minor point: Frenkel et al (2021) does not suffice as a reference to Target Prop algorithms. Could you please add Lee et al (2015) -- the seminal target prop paper --, Meulemans et al (2020), Ernoult et al (2022).\n\nLIMITATIONS:\nTo summarize my points of advice above:\n- In terms of presentation, I would recommend:\n   + you clarify the derivation of your learning rule along the lines suggested above,\n   + state clearly the intractability of $\\ell^{(2)}$ for the ELBO gradient and better explain the heuristic used (\"data mixing\"),\n   + distinguish between the feedback and feedforward block parameters, e.g. $\\theta_f$ and $\\theta_b$,\n   + remove that your approach allows for forward/backward pass parallelization and backward pass parallelization across layers. While you *can* do this in practice, the theoretical approach itself (e.g. variational inference) does not prescribe doing this. This might explain why the penultimate and upstream blocks don't learn.\n+ **Please write a detailed pseudo-algorithm** for the proposed procedure for a given training batch. That would be extremely helpful.\n\n- Try to check if the previous blocks are really learning. My hypothesis is that it is not and that it requires fixing the algorithm itself. \n\n- If uncertainty estimation matters to you, I would suggest considering the MVTech dataset.", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a block-local learning rule as an alternative to end-to-end gradient backpropagation to train neural networks. They present a probabilistic view of neural network representations and assuming an exponential family of distributions, derive a learning rule that can be understood as forward and backward message passing between blocks. Notably their message passing interpretation allows them to formulate auxilliary local losses that can be then optimized using gradient descent at the block-local level. Furthermore, they claim that their algorithm is a more principled way of performing algorithms like Feedback Alignment and Target Propagation. Finally, they demonstrate that their algorithm can be used to train ResNets (ResNet-18 & ResNet-50) on certain vision datasets and Transformer architecture on a sequence prediction task. Overall, the proposed algorithm seems a promising alternative to backpropagation with local learning properties that enable better memory footprint and neuroscientific realism than backpropagation.\n\nSTRENGTHS:\n1. The paper does a good job in explaining the probabilstic interpretation and message passing view of forward and backward phase in neural networks.\n2. The authors also empirically demonstrate the efficacy in training neural networks in Autoencoder setting as well as image classification and sequence prediction settings. \n3. The proposed probabilistic interpretation of activations allows uncertainty estimation in neural network predictions. Although these uncertainty estimates are not benchmarked in the paper, I feel this is a strength of the proposed method.\n4. Owing to its block-local nature, the proposed algorithm can be thought to be a biologically-plausible credit assignment algorithm for hierarchical neural networks. In doing so, this paper also offers a potential solution to memory-efficient distributed training of neural networks. \n\nOverall, I think it's a strong submission and is very relevant for the NeurIPS community.\n\nWEAKNESSES:\n1. The writing and presentation in the paper is sometimes hard to read and understand, thereby leading to lack of an in-depth understanding of the proposed algorithm.\n2. The current version of the paper does a great job in introducing the probabilistic interpretation of the neural network activations and the message passing view of forward and backward passes. However, their main learning rule is described in Section 3.4 which is not as clearly described. Unfortunately, the reader is deferred to the Appendix for key details of the derived learning rule, which adds to the lack of clarity regarding their proposed learning rule. \n3. In contrast to the Feedback alignment algorithm, the block-local learning algorithm seems to overfit significantly to the Cifar-10 dataset. Although this is an interesting finding in itself, the paper doesn't offer an explanation into this phenomenon or which components of their algorithm contribute to this. \n4. The discussion section offers more conjectures and falls short of discussing the implications of their results. For instance, the discussion section doesn't revisit the issue of overfitting or dive deeper into strategies around choosing the blocks & their backward counterparts and how these choices could affect the performance of the algorithm. The discussion section could also potentially highlight the potential biological plausibility of the proposed algorithm.\n\nQUESTIONS:\n1. In Like 115 (top of page 4), there seems to be a $log$ missing after the equals sign. Also, it seems that the right hand side of the expression would be a log-sum term. It is not clear how you write it as a sum of log terms in right hand side of Eq. 2. \n2. In Eq. 4, you describe $\\beta_k(z_k)$ to be the backward messages and $\\rho_k(z_k)$ to be the estimated posterior. However, in Eq. 7, you describe passing posterior messages $\\rho_k(z_k)$ using backward network activations. Could you please clarify this apparent change of notation and/or interpretation of backward messages?\n3. In Eq. 7, the partial derivative wrt $\\theta$ is computed only for the block local parameters right? But in the true formulation, gradients should be computed for all parameters in the computational graph. Is that correct? Or does the formulation of the probalistic graph enable inferring the true gradients by just computing the block-local parameter gradients? The variational local loss also probably plays a role here. Could you kindly clarify this part as it seems central to the entire proposal?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper introduces a novel framework for block-local training of deep networks. It proposes a twin network design that propagates information backwards from targets to the input to provide auxiliary local losses. This design allows forward and backward propagation to occur in parallel preventing the problems of weight transport and locking across blocks. This design is applied to training ResNets and transformers on several tasks.

STRENGTHS:
- Overall this paper is clearly written and proposes a novel interesting idea. I think this is an exciting direction that others will be able to build upon.
- The proposed twin architecture and the treatment of block outputs as uncertainty is novel
- The empirical results are strong and show a clear advantage of the block-local learning method particularly on CIFAR-10

WEAKNESSES:
- One selling point of this work is improved training efficiency. I would like to see an analysis even theoretically of what type of speedups can achieved with the proposed block-local training method.
- Although a few different architectures were evaluated the effect of block size on performance was not discussed. This seems like a important parameter to address as it effects both how parallelized/distributed training can be and biological plausibility.

QUESTIONS:
How would this method scale with block size?

What sort of practical performance speedups can be expected from using your block-local learning? (for example on a system with 1 gpu vs on a system with several gpus)

LIMITATIONS:
I think the limits of the biological plausibility of the twin architecture where the backwards network requires the same number of parameters as the forwards network should be discussed more.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors address the problem of weight transport and weight locking issue in backprop by introducing a new bio-plausible algorithm known as block-learning to train NNs. The model uses different forward and backward weights, creating a twin network-like scheme to learn efficient signals via local losses. The proposed learning algorithm is tested on convolution and transformer-based architectures to show that the proposed framework can scale to complex architectures.\n\nSTRENGTHS:\n1. Well-written paper\n2. Experiments on convolution and transformers show that the proposed work can scale to complex architectures.\n\nWEAKNESSES:\n1. Novelty is limited, given current framework shares several similarities with other approaches, such as Local representation alignment.\n2. Related work is missing several key citations.\n3. Experimental setup is restricted, as several SoTa bio-plausible approaches, including PC based approaches, are not compared\n\nQUESTIONS:\nOther examples of Block-learning framework is Local representation alignment (LRA-E[4], Rec-LRA [5]) , Difference target Propagation (DTP [2], DTP-sigma [4], DTP with backward targets [1], DTP with fixed weights [3]), weight mirroring [11] and Neural Generative Coding (Conv-NGC [6], NGC [8], Act-NGC [7]). The authors should compare and contrast against these existing lines of works, given all these frameworks have shown scaling results on various domains and architectures.\n\nSecond several works on predictive coding is not cited [9,10], given they have shown to approximate BP and have shown to achieve similar performance on various benchmarks. Can authors compare against PC based approaches? FA is known to struggle on complex architectures, hence to have \n\nBackprop is known to struggle whenever Bias is set to high value, it is not clear, whether experiments with FA and BP are performed with this settings. If yes, then how does model perform when you set biases to some low number such as 0.01? Can authors report these numbers?\n\n\nWhat is the benefit of current approach? Do you observe faster convergence? Better features (for instance one can use T-sne plots to visualize separations between classes, or visualize features of intermediate conv layers)? It would be beneficial if authors can report advantages of proposed approach.\n\nAs report by Lillicrap and [4], can authors show update angle compared to BP for the proposed framework? Does model update lies within 90 degree or even 45 degree compared to BP? Such analysis would further strengthen proposed work.\n\nHow robust is the model? Can authors report model performance across various settings/hyperparameter settings?\n\n1.\tErnoult, M.M., Normandin, F., Moudgil, A., Spinney, S., Belilovsky, E., Rish, I., Richards, B. and Bengio, Y., 2022, June. Towards scaling difference target propagation by learning backprop targets. In International Conference on Machine Learning (pp. 5968-5987). PMLR.\n2.\tLee, D.H., Zhang, S., Fischer, A. and Bengio, Y., 2015. Difference target propagation. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2015, Porto, Portugal, September 7-11, 2015, Proceedings, Part I 15 (pp. 498-515). Springer International Publishing.\n3.\tShibuya, T., Inoue, N., Kawakami, R. and Sato, I., 2023, June. Fixed-Weight Difference Target Propagation. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 37, No. 8, pp. 9811-9819).\n4.\tOrorbia, A.G. and Mali, A., 2019, July. Biologically motivated algorithms for propagating local target representations. In Proceedings of the aaai conference on artificial intelligence (Vol. 33, No. 01, pp. 4651-4658).\n5.\thttps://ojs.aaai.org/index.php/AAAI/article/view/26118\n6.\tOrorbia, A. and Mali, A., 2022. Convolutional Neural Generative Coding: Scaling Predictive Coding to Natural Images. arXiv preprint arXiv:2211.12047.\n7.\tOrorbia, A.G. and Mali, A., 2022, June. Backprop-free reinforcement learning with active neural generative coding. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 1, pp. 29-37).\n8.\tOrorbia, A. and Kifer, D., 2022. The neural coding framework for learning generative models. Nature communications, 13(1), p.2064.\n9.\tMillidge, B., Salvatori, T., Song, Y., Bogacz, R. and Lukasiewicz, T., 2022. Predictive coding: towards a future of deep learning beyond backpropagation?. arXiv preprint arXiv:2202.09467.\n10.\tSalvatori, T., Pinchetti, L., Millidge, B., Song, Y., Bao, T., Bogacz, R. and Lukasiewicz, T., 2022. Learning on arbitrary graph topologies via predictive coding. Advances in neural information processing systems, 35, pp.38232-38244.\n11.\tAkrout, M., Wilson, C., Humphreys, P., Lillicrap, T. and Tweed, D.B., 2019. Deep learning without weight transport. Advances in neural information processing systems, 32.\n\nLIMITATIONS:\n1. Comparison is needed with relevant methods.\n2. Analysis and ablation study missing.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel approach to the estimation of deep neural network parameters using block-localized backpropagation in conjunction with belief propagation. This approach is much more parallelizable, thus should help for distributed training, enabling horizontal scaling across devices.\n\nThe proposed approach works using a twin *backward* network, and by incorporating the belief messages into the block-local losses which are optimized using gradient descent.\n\nSTRENGTHS:\n- The proposed approach is an interesting combination between belief propagation and back-propagation. I particularly appreciated the view of a neural network as a Markov chain.\n- The proposed approach is significant, especially considering the current state of deep learning research. Large neural networks are steadily becoming the norm, and the development of specific learning algorithms for this kind of models is a valid and important research direction.\n\nWEAKNESSES:\n- The paper is sometimes not clear. I personally had difficulties understanding the following sections;\n    - 3.1, especially after equation (2)\n    - 3.4, it seems it requires the supplementary material to be correctly comprehended.\n- I think with the current state of the paper it might be difficult to reproduce the reported results. The algorithm itself is not completely clear, and I think the submission would improve from a clear explanation (maybe provided in the supplementary material).\n- A more in-depth analysis of the BLL algorithm convergence would improve the paper's strength.\n\nQUESTIONS:\n- How big should the twin network be? As much as the forward network, or have you found that it can be smaller?\n- This is a suggestion, but I think the submission will benefit from a clear description of the algorithm (possibly written in pseudo-code). This can be appended in the supplementary material and then referenced in the main paper.\n- Can a Gaussian distribution with non-constant variance be used to model the layer probabilities, or would the algorithm not work with this assumption?\n- I believe there is an error in the equations S9, S10, S11, S13. I imagine index *k* should not be one of the variables of the summation.\n- I believe there is an error in line [115]. I imagine it should be *p(y|x)* instead of *log p(y|x)*.\n- This is a simple suggestion, but in my experience (in the ML community) it seems to be more common to use 𝔼[…] and ∇ to indicate the expected value and gradient respectively. For a more standardized notation, I personally would favor that nomenclature. If you think that would be detrimental to the paper presentation in any way, you are free to ignore this suggestion.\n\nLIMITATIONS:\n- The algorithm requires to train an additional twin network, thus the number of total parameters optimized is greater than with standard gradient descent.\n- The message-passing operation could affect convergence time, Fig. 3 seems to suggest otherwise, but it may not be always holtrue for other dataset or hyper-parameters.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-26"}