InfoPro: Locally Supervised Deep Learning by Maximizing Information Propagation

Abstract

End-to-end (E2E) training has become the de-facto standard for training modern deep networks, e.g., ConvNets and vision Transformers (ViTs). Typically, a global error signal is generated at the end of a model and back-propagated layer-by-layer to update the parameters. This paper shows that the reliance on back-propagating global errors may not be necessary for deep learning. More precisely, deep networks with a competitive or even better performance can be obtained by purely leveraging locally supervised learning, i.e., splitting a network into gradient-isolated modules and training them with local supervision signals. However, such an extension is non-trivial. Our experimental and theoretical analysis demonstrates that simply training local modules with an E2E objective tends to be short-sighted, collapsing task-relevant information at early layers, and hurting the performance of the full model. To avoid this issue, we propose an information propagation (InfoPro) loss, which encourages local modules to preserve as much useful information as possible, while progressively discarding task-irrelevant information. As InfoPro loss is difficult to compute in its original form, we derive a feasible upper bound as a surrogate optimization objective, yielding a simple but effective algorithm. We evaluate InfoPro extensively with ConvNets and ViTs, based on twelve computer vision benchmarks organized into five tasks (i.e., image/video recognition, semantic/instance segmentation, and object detection). InfoPro exhibits superior efficiency over E2E training in terms of GPU memory footprints, convergence speed, and training data scale. Moreover, InfoPro enables the effective training of more parameter- and computation-efficient models (e.g., much deeper networks), which suffer from inferior performance when trained in E2E. Code: https://github.com/blackfeather-wang/InfoPro-Pytorch.

Appendices

Appendix: Architecture of Auxiliary Networks

In this section, we introduce the network architectures of $\varvec{w}$, $\varvec{\psi }$ and $\varvec{\phi }$ we use in our experiments. Note that, $\varvec{w}$ is a decoder that aims to reconstruct the input images from deep features, while $\varvec{\psi }$ and $\varvec{\phi }$ share the same architecture except for the last layer. The architectures used on CIFAR, SVHN and STL-10 are shown in Table 18 and Table 19. Architectures on ImageNet for ResNet/ResNeXt and EfficientNet are shown in Tables 20, 21 and Tables 22, 23, respectively. We apply the same mobile inverted bottleneck convolutional layers (MobileConv) (Sandler et al., 2018; Tan and Quoc, 2019) and Swish activations (Elfwing et al., 2018; Tan and Quoc, 2019) as EfficientNet to its corresponding auxiliary networks. For DeiT-S, we simply use an MLP as $\varvec{w}$ to reconstruct the model inputs, while $\varvec{\psi }$ consists of a 3x3 convolutional layer with stride=2, followed by 3 standard Transformer layers. The architectures of $\varvec{w}$ and $\varvec{\psi }$ for semantic segmentation are shown in Table 24 and Table 25, where the same networks are used with varying input sizes. Herein, “s”, “p” and “oc” refer to stride, padding and output channels, respectively. “conv.” refers to a single convolutional layer.

Table 18 Architecture of $\varvec{w}$ on CIFAR, SVHN and STL-10

Table 19 Architecture of $\varvec{\psi }$ and $\varvec{\phi }$ on CIFAR, SVHN and STL-10

In particular, in all of our experiments, DGL (Belilovsky et al., 2020) uses exactly the same architecture of auxiliary networks as us.

Appendix: Details of Mutual Information Estimation

In this section, we describe the details on obtaining the estimates of $I(\varvec{h}, \varvec{x})$ and $I(\varvec{h}, y)$ we present in Figs. 2, 3 and 5.

Table 20 Architecture of $\varvec{w}$ on ImageNet

Table 21 Architecture of $\varvec{\psi }$ on ImageNet

Table 22 Architecture of $\varvec{w}$ for EfficientNet-B2

Table 23 Architecture of $\varvec{\psi }$ for EfficientNet-B2

Table 24 Architecture of $\varvec{w}$ for semantic segmentation

Table 25 Architecture of $\varvec{\psi }$ for semantic segmentation

As we have discussed in Sect. 4.3, the expected reconstruction error $\mathcal {R}(\varvec{x}|\varvec{h})$ follows $I(\varvec{h}, \varvec{x})\!=\!H(\varvec{x})\!-\!H(\varvec{x}|\varvec{h})\!\ge \!H(\varvec{x})\!-\!\mathcal {R}(\varvec{x}|\varvec{h})$ (Vincent et al., 2008; Rifai et al., 2012; Kingma and Welling, 2013; Makhzani et al., 2015; Hjelm et al., 2019). Therefore, similar to Sect. 4.3, we estimate $I(\varvec{h}, \varvec{x})$ by training a decoder parameterized by $\varvec{w}$ to obtain the minimal reconstruction loss, namely $I(\varvec{h}, \varvec{x}) \approx \mathop {\max }_{\varvec{w}}[H(\varvec{x}) - \mathcal {R}_{\varvec{w}}(\varvec{x}|\varvec{h})]$. Note that, ideally, this bound can be arbitrarily tight provided that $\varvec{w}$ has sufficient capacity. In specific, we use the same network architecture as Table 18, and train it for 10 epochs to minimize the averaged binary cross-entropy reconstruction loss of each pixel. An adam (Kingma et al., 2014) optimizer with default hyper-parameters (lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) is adopted. Naive as this procedure might seem, for one thing, we find that it is sufficient to reconstruct the input images well given enough information, and meanwhile distinguish different values of $I(\varvec{h}, \varvec{x})$ via the quality of reconstructed images. For another, we are primarily concerned with the comparisons of $I(\varvec{h}, \varvec{x})$ between end-to-end training and various cases of greedy supervised learning rather than obtaining the exact values of $I(\varvec{h}, \varvec{x})$. The same training process is applied to all the intermediate features $\varvec{h}$, and hence the comparisons are fair. Finally, since $H(\varvec{x})$ is a constant, for the ease of understanding, we simply present $1-{AverageBinaryCrossEntropyLoss}(\varvec{x}|\varvec{h})$ as the estimates of $I(\varvec{h}, \varvec{x})$, equivalent to adding the constant $1 - H(\varvec{x})$ to the real estimates of $I(\varvec{h}, \varvec{x})$.

For $I(\varvec{h}, y)$, as we have discussed in Sect. 4.3 as well, since $I(\varvec{h}, y) = H(y) - H(y|\varvec{h}) = H(y) - \mathbb {E}_{(\varvec{h}, y)}[-\text {log}\ p(y|\varvec{h})]$, we train an auxiliary classifier $q_{\varvec{\psi }}(y|\varvec{h})$ with parameters $\varvec{\psi }$ to approximate $p(y|\varvec{h})$, such that we have $I(\varvec{h}, y)\!\approx \!\mathop {\max }_{\varvec{\psi }} \{H(y)-\frac{1}{N}[\sum _{i=1}^N\!-\text {log}\ q_{\varvec{\psi }}(y_i|\varvec{h}_i)]\}$. Here we simply adopt the test accuracy of $q_{\varvec{\psi }}(y|\varvec{h})$ as the estimates of $I(\varvec{h}, y)$, which is highly correlated to the value of $-\frac{1}{N}[\sum _{i=1}^N\!-\text {log}\ q_{\varvec{\psi }}(y_i|\varvec{h}_i)]\}$ (or say, the cross-entropy loss). This can be viewed as the highest generalization performance that a classifier based on $\varvec{h}$ is able to reach. Notably, we use a ResNet-32 as $q_{\varvec{\psi }}$. For the inputs of $q_{\varvec{\psi }}$, we up-sample $\varvec{h}$ to $32\times 32$ and map $\varvec{h}$ to 16 channels at the first layer. All the training hyper-parameters of the ResNet-32 are the same as Sect. 5.2.