Greedy Layerwise Learning Can Scale to ImageNet

In plain words

Deep image classifiers are usually trained all at once: the error at the output is sent backwards through every layer. That makes it hard to say what any one middle layer is supposed to do, and it needs a lot of memory.

This paper trains the network one layer at a time instead. Each new layer gets a small temporary classifier sitting on top of it, and that classifier is trained to predict the labels. Once a layer is finished it is frozen, its output becomes the input for the next layer, and the process repeats. The temporary classifiers are simple and shallow, which makes each step easier to analyse.

The surprise is that this works on ImageNet, a large dataset where earlier layer-by-layer methods had failed. Stacking such blocks gives networks that are competitive with standard end-to-end CNNs of the same kind. The authors also look inside the trained layers and find that the data becomes easier to separate with a straight line as depth grows.

TL;DR

• Train a CNN block by block: each block is fit with a shallow auxiliary classifier, and earlier blocks are frozen.
• A block is a 3x3 convolution with bias and ReLU, with an invertible 2x2 space-to-depth downsampling at some depths; no pooling is used §3.1.
• One-hidden-layer blocks (k=1) exceed AlexNet on ImageNet; 3-hidden-layer blocks (k=3) reach VGG-13-level accuracy and match end-to-end VGG-11 tab. 3.
• All numbers are author-reported. No official code is linked.

Problem

End-to-end backpropagation specifies only the final output objective, so the role of intermediate layers is only indirectly defined §1. Backpropagation also stores most intermediate activations and computes most intermediate gradients, which costs memory and is considered biologically implausible §1.

Earlier layer-wise supervised methods were mostly shown on small datasets. The authors report that alternative training methods such as Feedback Alignment reach only 6.6% top-1 on ImageNet, against 50.9% with backprop §4 intro tab. 3. The assumption under test is that greedy layer-wise training can reach end-to-end accuracy on ImageNet-scale CNNs §1.

Method

The network has J blocks, trained in order. Starting from x_0 equal to the input image, block j computes x_{j+1} by applying its 3x3 convolution W_{θ_j} (with bias) to the downsampled input P_j x_j and then ReLU. P_j is the invertible downsampling, which reorganises each 2x2 spatial neighbourhood into channels and halves the resolution; it is applied at some depths and is the identity at others §3.1 §3.1 discussion of pooling.

x_{j+1} = \rho\, W_{\theta_j} P_j x_j, \qquad z_{j+1} = C_{\gamma_j} x_{j+1} \in \mathbb{R}^c eq. 1

Here ρ is ReLU, c is the number of classes, and C_{γ_j} is the auxiliary classifier of block j §3.1.

The auxiliary classifier has a fixed structure. Its spatial average A pools the features into four quadrants, giving a 2x2 map; L is a linear map to c outputs. For k hidden layers there are k−1 extra convolutions of constant width before A §3.1 §4 Naming Conventions.

C_{\gamma_j} x_j = \begin{cases} L A x_j & k = 1 \\ L A \rho \tilde W_{k-2} \cdots \rho \tilde W_0 x_j & k > 1 \end{cases} eq. 2

Training is layer-wise. While block j is trained, all earlier parameters are fixed. Block j and its classifier are fit together by minimising the empirical risk on the block-j inputs; the loss is cross-entropy, given as an example §3.2.

\hat{\mathcal{R}}(z; \theta, \gamma) = \frac{1}{N} \sum_{n} l\big(z(x^n; \theta, \gamma), y^n\big) §3.2

The procedure is Algorithm 1 alg. 1. The optimiser is SGD in the experiments §3.2 §4. The paper states a technical requirement: the result must not be worse than the starting point, which is met by keeping the best iterate along the optimisation trajectory §3.2.

# Layer-wise training of a CNN  [alg. 1]
# Input:  training samples (x^n, y^n), n = 1..N
# Output: block parameters theta_0 .. theta_{J-1}

def train_layerwise(samples, num_blocks):
    # x_0 is the raw input image  [§3.2] [eq. 1]
    inputs = [x for x, y in samples]
    labels = [y for x, y in samples]
    blocks = []
    for j in range(num_blocks):
        # Fit block j and its auxiliary head; earlier blocks stay frozen  [§3.2]
        theta_j, head_j = minimise_risk(inputs, labels)
        # Keep the best iterate of the optimiser, not the last one  [§3.2]
        blocks.append(theta_j)
        # Pass the inputs through the frozen block to get x_{j+1}  [eq. 1]
        inputs = relu(conv(theta_j, downsample_j(inputs)))
    return blocks

# Auxiliary head for a block with k hidden layers  [eq. 2]
# Input:  features (B, channels, H, W)
# Output: logits (B, c)

def auxiliary_head(features, hidden_convs, linear):
    # k - 1 extra convolutions of constant width, each followed by ReLU  [eq. 2]
    for conv_layer in hidden_convs:
        features = relu(conv_layer(features))
    # Average over the four spatial quadrants  [§4 Naming Conventions]
    pooled = quadrant_average(features)  # (B, aux_width, 2, 2)
    # Linear map to the c class logits  [eq. 2]
    return linear(flatten(pooled))

The paper also describes its reading of k. With k=1 each block solves a one-hidden-layer problem, which is the case with the clearest theory §3.3. With k=2 or 3, each block is the first layer of a 2- or 3-hidden-layer auxiliary problem, of which only the first convolution is kept §3. The authors argue that these shallow sub-problems are easier to optimise than deep ones §3.3.

Two results in §3.2 and Appendix A give partial theory. Proposition 3.1 says that with P_j = Id the trained layer's objective is no worse than that of the previous layer, because the previous layer's solution can be represented inside the new block §3.2 app. A. Proposition 3.2 says that if each block solves its sub-problem to accuracy ε and the blocks are 1-Lipschitz with bias bounded by B, the final representation differs from the exact one by O(J²ε) §3.2 app. A.

\lVert x_J^* - \hat x_J \rVert = \mathcal{O}(J^2 \epsilon) eq. 3

For k=1 the authors relate each sub-problem to a least-squares fit of a one-hidden-layer network, which can be analysed with known approximation bounds §3.3. The estimator is written as follows.

(\hat L, \hat\theta) = \arg\inf_{L, \theta} \sum_n \lVert f(x^n) - L \rho W_\theta x^n \rVert^2 eq. 4

The ensemble prediction is a sum over the auxiliary classifier outputs, with weight 2^j on block j §4 Naming Conventions.

Architectures used. ImageNet k=1: J=8, widths 256 to 4096, downsampling before blocks j=2,3,4,6, no batch-norm §4.1 tab. 5. ImageNet k=2,3: J=8, width M=128, auxiliary width 256, downsampling at j=2,4,6, batch-norm in the auxiliary problems §4.3 tab. 4. CIFAR-10 k=1: J=5, M=256, downsampling at j=2,3 §4.1. CIFAR-10 k=2,3: J=4, M=128, auxiliary width 256, downsampling at j=2 §4.3.

Final model. For the ImageNet k=3 model, the final auxiliary network is part of the model; its width is reduced from 2048 to 512 for the last step at a small cost in accuracy §4.3 tab. 1. A layer-wise compression step, which prunes each new block's filters from 128 to 64 and fine-tunes the auxiliary head, reaches the baseline accuracy on CIFAR-10 §4.3 Layerwise Model Compression.

Hyperparameters

Results

Author-reported. Setup: ImageNet-1k ILSVRC, single-crop top-5 on the validation set, ensemble numbers in brackets; baselines are AlexNet, VGG and ResNet-152 trained end-to-end and the layer-wise alternatives §4.3 tab. 3. Batch and epoch counts are in the table above.

ImageNet tab. 3:

Square brackets give the end-to-end version of the same model, where available tab. 3.

CIFAR-10 tab. 2:

The end-to-end version of the k=1 CIFAR model is 89.7% §4.1.

Final auxiliary width, ImageNet k=3 tab. 1: M̃_f = 2048 gives 88.7 top-5; M̃_f = 512 gives 88.5 top-5 §4.3.

Downsampling choice, CIFAR-10 k=1 tab. 7: strided convolution 87.8, invertible downsampling 88.3, average pooling 87.6, max pooling 88.0 app. C.1.

Layer-wise compression, CIFAR-10 §4.3: baseline 5-layer network of width 64, end-to-end, 87.5%; layer-wise with compression 87.6%.

Progressive linear separability, CIFAR-10 k=1 §4.2 fig. 3: linear separability rises monotonically with depth; CNN-p separability also improves, more so for k>1. Average standard deviation over five seeds is 0.58% §4.2.

Caltech-101 transfer, linear model on frozen k=1 ImageNet features, 30 training images per class, 10 splits app. C.3 tab. 8: ConvNet from scratch 46 ± 1.7%; layer 1 45.5 ± 0.9; layer 2 59.9 ± 0.9; layer 3 70.0 ± 0.9; layer 4 75.0 ± 1.0; layer 8 82.6 ± 0.9.

Per-layer ImageNet top-5 rises from about 23% to about 80%, nearly linearly over the first five blocks §4.1 fig. 2.

Paper vs code

No official code.

I searched for an official repository in the order the skill gives. The paper and its LaTeX source link no code or project page; the source only mentions a planned release in commented-out text, which is not a published link. A GitHub repository search for "greedy layerwise", "layerwise cnn imagenet" and "greedy layerwise imagenet belilovsky" found nothing by the authors. The first hits (for example moskomule/greedy-learning and srags/greedy-layerwise-survey) are not by the authors and were not used. The repository lists of the account edouardoyallon (its matches are scattering-transform code) and of the account eugenium (not verified as an author; its only related match is a downsampled-ImageNet dataset repository) contain no layer-wise training code.

Pitfalls

• Indexing conventions differ. Algorithm 1 uses block indices j = 0..J−1 with parameters θ_j, while the text writes the trained parameter as θ̂_{j+1} alg. 1 §3.2. The downsampling depths (j = 2, 3, 4, 6) are 0-based: they match the 1-based table layers 3, 4, 5, 7 in tab. 5 and the layers in tab. 4. Implementers should use 0-based indices throughout.
• The ensemble is described as a weighted average but written as an unnormalised sum. The text gives Z = Σ 2^j z_j with j running from 1 to J, which does not match the 0-based block indexing in Algorithm 1 and is not normalised §4 Naming Conventions. The normalisation and index range used for the reported ensemble numbers are not stated.
• Ensemble on CIFAR-10 is not defined. The paper reports ensemble accuracies in tab. 2 but does not say how the CIFAR ensemble is formed.
• Loss is only an example. Cross-entropy is given as "e.g." in the risk definition §3.2. The experiments do not state the loss explicitly, so cross-entropy is an inference from the wording.
• Training schedule scope is unclear. The CIFAR-10 schedule is explicitly per layer (50 epochs each) §4. The ImageNet schedule (45 epochs, decays every 20) is given without saying whether it is per block or total, and whether k=2,3 models use it too. The k=3 VGG-11 comparison uses "the same 45 epoch training schedule" §4.3, which suggests per block, but this is not confirmed.
• Batch-norm placement is unclear. The k=1 ImageNet model has none §4.1 tab. 5. The k=2,3 models use batch-norm "for these auxiliary problems" §4.3. The paper does not say whether this covers the block convolutions, the auxiliary convolutions or both.
• Widths double at downsampling. The output width is twice the input width when a block downsamples, because the 2x2 space-to-depth step quadruples the channels and the block keeps the same spatial size §4 Naming Conventions tab. 5. Proposition 3.1 assumes P_j = Id, which keeps the width unchanged, but the architecture doubles the width at downsampling blocks.
• The proof of Proposition 3.1 is terse. The proof chooses θ_0 so that W_{θ_0} is the identity, using ρ(ρ(x)) = ρ(x) app. A. A 3x3 convolution with bias can only be the identity when its input and output channel counts match, and the first block maps 3 colour channels to M. The proof does not address this case, and it does not say what "under mild conditions" refers to in the main text §3.2.
• Proposition 3.2 assumes its own conditions. It takes the blocks to be 1-Lipschitz "without loss of generality" and biases bounded by B §3.2 app. A. The paper does not justify this, and the ε-approximation assumption on the solver is not shown to hold for SGD.
• Single-block accuracies in the introduction are not the headline models. The text says a one-hidden-layer CNN reaches 61% on CIFAR-10 and 12.3% on ImageNet §4 intro. These numbers are not tied to a configuration. They are not the 88.3% and 79.7% headline results from the full k=1 models tab. 2 tab. 3.
• Table and text disagree on BoostResNet. The text gives 82.0% §4.1 and tab. 2 gives 82.1%.
• Parameter counts are for other configurations. tab. 6 lists the k=1 model with J=6 (96M), not the reported J=8 model; the row for J=8 is commented out in the source. Appendix B says that dropping the last two blocks cuts the k=1 model size by a quarter and gives 78.8% top-5 at J=6 against 79.7% at J=8 app. B.
• Memory and parameter costs. The ImageNet k=3 model has about 102M parameters, and the reduced-width version has about 46M tab. 6. The large final auxiliary network carries most of the parameters for k=2,3 app. B.
• Memory saving is an author claim. The end-to-end CIFAR-10 version is said to use about five times more GPU memory than the k=1 training §4.1. No memory measurements are given for the ImageNet runs.
• Validation split is the test set. ImageNet results are on the validation set §4.1 tab. 3. The paper does not describe a separate held-out split.
• Ensembles and the final auxiliary head. Which classifier heads are part of the final deployed model for k=1 is not stated in the main text. Figures, tables and the ensemble numbers assume all auxiliary outputs are combined.
• Downsampling choice is minor on CIFAR-10. The four pooling and downsampling options differ by at most 0.7 points, and the invertible operator is only slightly ahead tab. 7. The paper's stated reason for choosing it is to avoid lossy operations app. C.1.
• Several equation and figure references are unnumbered in the source. Equation numbers and figure and table numbers are taken from the order of appearance in the LaTeX source, not from a compiled PDF, which was not available. Check against the PDF before relying on a specific number.

Lineage

• Builds on: greedy layer-wise pretraining of deep generative and supervised networks (Bengio et al. 2007; Hinton et al. 2006), cited as the origin of the idea §2. Shallow-network theory for one-hidden-layer training (Barron 1994; Arora et al.; Malach et al. 2018) is cited for the k=1 analysis §3.3.
• Compared against: BoostResNet (Huang et al. 2017), ProvableNN (Malach et al. 2018), Mosca et al. 2017 and FeedbackAlign and DTargetProp (Bartunov et al. 2018; Xiao et al. 2019) tab. 2 tab. 3. End-to-end AlexNet, VGG and ResNet references tab. 3.
• Superseded by: not stated in this card's source.