[Paper Review] Putting An End to End-to-End: Gradient-Isolated Learning of Representations

Sindy Löwe, Peter O’Connor+1 more|arXiv (Cornell University)|May 28, 2019

Domain Adaptation and Few-Shot Learning55 references49 citations

TL;DR

We introduce Greedy InfoMax (GIM), a self-supervised, gradient-isolated, layer-wise training method that optimizes mutual information per module rather than end-to-end backpropagation, yielding competitive representations for audio and vision tasks and enabling asynchronous, memory-efficient training.

ABSTRACT

We propose a novel deep learning method for local self-supervised representation learning that does not require labels nor end-to-end backpropagation but exploits the natural order in data instead. Inspired by the observation that biological neural networks appear to learn without backpropagating a global error signal, we split a deep neural network into a stack of gradient-isolated modules. Each module is trained to maximally preserve the information of its inputs using the InfoNCE bound from Oord et al. [2018]. Despite this greedy training, we demonstrate that each module improves upon the output of its predecessor, and that the representations created by the top module yield highly competitive results on downstream classification tasks in the audio and visual domain. The proposal enables optimizing modules asynchronously, allowing large-scale distributed training of very deep neural networks on unlabelled datasets.

Motivation & Objective

• Motivate learning representations without labels or global end-to-end backpropagation.
• Develop a greedy, module-wise training scheme inspired by information preservation principles.
• Demonstrate that locally trained modules can collectively yield strong downstream performance.
• Show practical benefits in memory efficiency and asynchronous training for very deep networks.

Proposed method

• Split a deep network into gradient-isolated modules with a gradient-blocking operator to prevent inter-module backpropagation.
• Train each module with a local InfoNCE loss to maximize information preservation between its input and output.
• Use a module-local contrastive loss f_k^m(z_{t+k}^m,z_t^m) with Eq. (3)–(4) to maximize mutual information.
• Optionally append an autoregressive context module g_ar to aggregate wider temporal or spatial context (Eq. (6)).
• After module training, discard the local scoring functions and use the stacked encoders for downstream features; optionally train a separate context module.
• Argue that iterative mutual information maximization encourages non-degenerate, informative representations across layers.

Experimental results

Research questions

• RQ1 Can mutual information maximization be effectively optimized in a layer-by-layer greedy (gradient-isolated) fashion without end-to-end backpropagation?
• RQ2 Do greedily trained modules progressively improve representations for downstream tasks in vision and audio?
• RQ3 What are the memory and training-efficiency trade-offs of gradient-isolated learning compared to end-to-end methods?
• RQ4 How does the addition of an autoregressive context module affect performance on tasks requiring broader context?

Key findings

• GIM achieves competitive downstream performance on STL-10 image classification without end-to-end backpropagation (81.9% ±0.3 with GIM vs 80.5% ±3.1 for CPC).
• GIM outperforms several self-supervised baselines and comparable supervised methods under limited labeling, illustrating strong representation quality.
• Ablations show the autoregressive module benefits tasks with temporal/contextual dependencies; removing BPTT or the ar module can degrade performance on such tasks.
• Greedy, layer-wise InfoNCE training improves intermediate representations, with each module outperforming its predecessor on downstream speaker classification tasks.
• GIM enables memory-efficient, asynchronous training; training three modules can reduce GPU memory usage by a factor of 2.8 versus end-to-end training.
• In LibriSpeech, GIM achieves 99.4% speaker classification accuracy and 62.5% phone classification accuracy, with ablations showing drops when removing BPTT or the autoregressive component.