Paper: Block-local learning with probabilistic latent representations (ICLR-2024)

Reviewer comment (Theory):
How to the principle or theory to select split blocks for each architecture?

Author response:
*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).

Reviewer comment (Experiments):
What is the reason for the large performance gap on Cifar10 and ImageNet tasks, compared with BP method? Are they using the same data preprocessing and augmentations?

Author response:
*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.

Reviewer comment (Evaluation):
What is the training time and convergence speed compared with BP method?

Author response:
*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.

Reviewer comment (Novelty):
The bidirectional propagation in this work overlaps with the paper titled "Bidirectional Backpropagation," which uses a set of synaptic weights for forward and backward propagation.

Author response:
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.

Reviewer comment (Novelty):
The bidirectional propagation approach generalizes to using two separate networks, which is a special case that uses deterministic dropout along the forward and backward propagation.

Author response:
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.

Reviewer comment (Novelty):
The prior paper "Bidirectional backpropagation for high-capacity blocking networks" combines bidirectional backpropagation with blocking, representing a combination of bidirectional propagation and deep-neural blocking, which appears similar to the methods described in the present paper.

Author response:
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.

Reviewer comment (Evaluation):
It's also worth noting that block/layer-wise learning can perform as well as backprop, and with a similar idea of using small backward error networks. See Fig. 1 and overall results in [Belilovsky et al., 2019] (cited in this paper). It's not clear if the performance difference is due to inherent problems with the probabilistic interpretation in this paper or some other reasons.

Author response:
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.

Reviewer comment (Presentation):
First of all, the authors should use the ICLR style (currently the typeface is wrong and citations are not highlighted) and fit the paper into 9 pages.

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

Reviewer comment (Theory):
Eq. 7 should have $\beta_k$ instead of $q_k$, right?

Author response:
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.

Reviewer comment (Reproducibility):
What’s the actual speed-up compared to backprop? With a naive implementation that doesn't account for unlocked updates, it should be as expensive. I guess it might be possible to put different blocks on different devices so the backward pass in earlier blocks will happen early on, but that wouldn’t work on a single GPU (right?) and also wouldn’t be easy to implement in PyTorch. How was it implemented by the authors?

Author response:
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.

Reviewer comment (Presentation):
Judging by Fig. 2, posterior bootstrapping introduces backward locking. Is that correct?

Author response:
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.

Reviewer comment (Theory):
It is difficult to understand how the variational distribution $q$ is actually defined. Only backward messages $\beta_k = q(y | z_k)$ are provided, and it is unclear how the full posterior $q(z_k | x, y)$ is defined. My best guess is that it is implicitly defined through the Bayes rule by combining the messages $\beta_k$ for all $1 \leq k \leq N$.

Author response:
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.

Reviewer comment (Writing):
Some notations are only implicitly defined, such as $q_k$, or not defined at all such as $a_{kj}$ and $z_{kj}$.

Author response:
$\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$.

Reviewer comment (Theory):
It is difficult to understand in which case the variational posterior could recover the true posterior. In other words, it is unclear which messages $\beta_k$ would lead to a perfect reconstruction of the true posterior, as well as how to compute them.

Author response:
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.

Reviewer comment (Experiments):
There are no guidelines on how to design the different components of the proposed algorithm. For example, there is no ablation study on the expressivity of the EF distribution used as a latent representation, nor on the expressivity of the feedback network.

Author response:
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.

Reviewer comment (Theory):
The distribution $q$ is only defined implicitly through the definition of backward messages $\beta_k = q(y | z_k)$. Could the author either confirm that it is only implicitly defined because we only need to compute messages $\beta_k = q(y | z_k)$ or specify the definition of the full posterior $q(z_k| x, y)$?

Author response:
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.

Reviewer comment (Theory):
Does equation (6) make use of the assumption that $p_k$ and $\alpha_k$ are Gaussian or from a distribution in the exponential family?

Author response:
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.

Reviewer comment (Theory):
Could the author clarify the approximation gap between the true log-likelihood and their variational surrogate? In particular, when does the variational loss recover the true log-likelihood?

Author response:
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.

Reviewer comment (Theory):
Could the authors give some intuition on the design of the latent representation and the expressivity needed in the feedback network?

Author response:
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.