Block-local learning with probabilistic latent representations

Q: How to the principle or theory to select split blocks for each architecture?

In general, the number of blocks will be determined based on the requirements of distributing the model (model size, no. of devices, memory per device etc.). There is a trade-off in number of splits and sacrifice in performance which depends on learning task and network architecture (See Figure S1).

Q: How to deal with models with multiple branches? Is it possible to split the network for each branch?

We did not test models with multiple branches yet. However, the theory would support such models. Will be interesting to test that in future work.

Q: What is the reason for the large performance gap on Cifar10 and ImageNet tasks, compare with BP method? Are there using same data preprocessing and augmentations?

Yes, they use the same preprocessing and augmentation. We think there are two main reasons for the gap: 1) due to time and resource limitations we were not able to do any hyper-parameter search for ImageNet. 2) For tasks with number of classes in the range of the batch size, the greedy update of the backward network may not be suitable (backward messages are estimated by 1-2 samples on expectation which doesn't provide rich enough statistics). Non-greedy updates of backward nets would likely fix that, which we are currently developing. We still thought it is beneficial for the paper to include the ImageNet experiments in the current, since the performance is at least not trivial. We will update the discussion and clarify these shortcomings of the current form of the algorithm.

Q: What is the training time and convergence speed compared with BP method?

We implemented 2 versions of our algorithm, one a naive implementation which is as expensive as backprop (with additional linear layers for the backward network) and second , as the reviewer rightly suggested to put different blocks on multiple devices. However the updates are still locked due to pytorch’s limitations. We consider these implementations as proof of concept and to fully utilize the speed-up of unlocked updates will require a c++ (libtorch) implementation.

The main contribution of our model is to show that using separated forward and backward networks, in a variational learning setup, local learning rules can be derived (Equations 6 and 7), that do not suffer from the locking and weight transport problem. Updates are calculated and applied on a per-block basis after a forward pass through the network. The backward network can have any architecture including sharing parameters with the forward network, however we didn't consider that case yet. Hence we would consider this to be the most general case. Also Bidirectional backpropagation doesn't interpret the outputs of the network as exp. family distribution parameters. Furthermore, our method doesn't include a network-wide finetuning phase as in Bidirectional backpropagation. The goal of Bidirectional backpropagation also doesn't seem to be solving the locking and weight transport problem. We will be happy to point out these differences in more details in a revised version of the paper.

Comments Thank you for pointing out these style issues. We have corrected them and revise the pdf.

Clarification questions 2. Actually, it should be $\rho_k$ as in Eq. (4). Technically $q_k$ is correct here, however, the definition of $q_k$ comes only late in the supplement. We noticed an ambiguity in our use of the distribution $q$. We will simplify the notation to resolve that issue.

3. Correct. However, a number of previous methods used a contrastive step to arrive at block-local learning rules. Here, no such contrastive step is needed. We will clarify this.

Conceptual questions

1. As mentioned in the response to ACie, we implemented 2 versions of our algorithm, one a naive implementation which is as expensive as backprop (with additional linear layers for the backward network) and second , as the reviewer rightly suggested to put different blocks on multiple devices. However the updates are still locked due to pytorch’s limitations. We consider these implementations as proof of concept and to fully utilize the speed-up of unlocked updates will require a c++ (libtorch) implementation.

2. No there is no backward locking in our method. All updates can be computed in forward mode and can be parallelized. The only required synchronization is in the forward path through the network, which is unavoidable. We explained the precise structure of the parallelism in Fig.3.

1. We tried to keep the initial part of the derivation general, therefore $\beta$ is initially only defined as any distribution $q(y|z_k)$. We use a variational posterior that factorizes into the parts $\alpha$ and $\beta$ according to Eq.(4). $\alpha$ is the true forward message computed by the forward network. $\beta$ is the approximate backward message that determines how well the network activation at block $k$ predicts the labels $y$. Based on this general derivation, to arrive at a concrete implemenation we used the exponential family assumption. The variational posterior that was used in our experiemnts is defined in Eq. S19.

2. $\alpha_{kj}$, $\beta_{kj}$ are the independent components corresponding to activations of individual units of the forward and backward network, respectively. They determine the natural parameters of distributions $\alpha_k$, $\beta_k$.

3. Eq. 6 is the general form for any member of the exp. family. Conveniently, only the first and second order statistics are needed to compute the divergence and $\mu$ and $\sigma$ refer here to these sufficient statistics. We establish this result in section S1.4 in detail. In experiments we used a Gaussian distribution, for which the derivation is in S1.5.

4. The approximation becomes exact when if $\beta_k = p(y|z_k)$, i.e. the true posterior of the targets $y$, given the network output at block $k$, according to the forward network (its true Bayesian inversion). However, with the simplified form of variational distribution that we use, where $\beta_k$’s are given by backward network activity, a gap between the true log-likelihood and the variational surrogate will always remain, as is the case in standard variational methods.

5. We chose the simplest possible feedback network that did not need additional gradient based updates for efficiency. Early tests showed that more sophisticated architectures were not needed in those cases but it is still a future work. However, more complex tasks like ImageNet may require more complex backward network structures, which may explain the gap in performance.