PaperlayerView exact revision on arXiv

Source-linked AI summary

Supermasks in Superposition

Mitchell Wortsman, Vivek Ramanujan, Rosanne Liu, Aniruddha Kembhavi, Mohammad Rastegari, Jason Yosinski, Ali Farhadi

arXiv:2006.14769v3cs.LGcs.AIstat.ML

TL;DR

Catastrophic forgetting limits sequential learning across many tasks. SupSup fixes a random base network, learns task-specific supermasks, and infers unknown task identity by entropy-minimizing gradient optimization. It performs well across task settings, including 2500 MNIST permutations, while extending to task-free mask allocation and constant-sized Hopfield storage.

  • Problem

    Sequentially learning many tasks without catastrophic forgetting remains a notable challenge for neural networks.

  • Method

    SupSup fixes a randomly initialized base network, learns a supermask for each task, and uses gradient-based entropy minimization to infer unknown task identity.

  • Results

    SupSup achieves state-of-the-art SplitImageNet performance with task identity and performs well on thousands of MNIST permutations without task information.

  • Takeaways & Limitations

    SupSup supports continual learning across task-identity settings while using task-specific subnetworks and extensions for task-free allocation and fixed-size storage.

  • Takeaways & Limitations

    Task inference fails on non-uniform, more challenging problems when models are poorly calibrated and overly confident for the wrong task.

Abstract

from arXiv · show

We present the Supermasks in Superposition (SupSup) model, capable of sequentially learning thousands of tasks without catastrophic forgetting. Our approach uses a randomly initialized, fixed base network and for each task finds a subnetwork (supermask) that achieves good performance. If task identity is given at test time, the correct subnetwork can be retrieved with minimal memory usage. If not provided, SupSup can infer the task using gradient-based optimization to find a linear superposition of learned supermasks which minimizes the output entropy. In practice we find that a single gradient step is often sufficient to identify the correct mask, even among 2500 tasks. We also showcase two promising extensions. First, SupSup models can be trained entirely without task identity information, as they may detect when they are uncertain about new data and allocate an additional supermask for the new training distribution. Finally the entire, growing set of supermasks can be stored in a constant-sized reservoir by implicitly storing them as attractors in a fixed-sized Hopfield network.

1 Introduction

SupSup addresses catastrophic forgetting by keeping a random base network fixed and learning task-specific supermasks. It supports task-identity inference, task-boundary detection, and fixed-size implicit mask storage.

  • Core approach: SupSup keeps a randomly initialized base network fixed and learns a separate supermask, or subnetwork, for each task.This prevents forgetting caused by changing shared network weights.
  • Core approach: Unknown task identity is inferred by optimizing a convex combination of learned supermasks to minimize output entropy.The combination weights α_i are adjusted using gradients toward the mask producing the most confident output.
  • Results: SupSup outperforms recent task-ID methods and scales to 2500 MNIST permutations without forgetting when task ID is available during training but not testing.For uniform tasks, task identity can be inferred with a single gradient computation.
  • Extensions: With task identities unavailable during training, SupSup can detect task boundaries and allocate new supermasks for new training distributions.This extends the approach beyond settings where task identity is supplied during training.
  • Extensions: A fixed-size Hopfield network can store supermasks implicitly as attractors, keeping reservoir size constant as tasks accumulate.The extension targets growing collections of learned supermasks.
  • Extensions: SupSup also reports more accurate task inference from adding superfluous neurons.This is presented as an empirical extension to the basic algorithm.

2 Continual Learning Scenarios and Related Work

The paper introduces a taxonomy for continual-learning scenarios based on task identity, label sharing, and evaluation conditions. It situates SupSup among related methods and emphasizes its broader scenario coverage.

  • Scenario taxonomy: Continual-learning scenarios vary by whether task identity is given during training and inference, whether labels are shared, and whether task space is discrete or continuous.These identifiers define the dimensions used to organize prior work.
  • Scenario taxonomy: The proposed three-letter taxonomy uses G or N for task identity during training and inference, and s or u for shared or unshared labels.In unshared settings, the model predicts both task identity and class; shared settings require only the shared label.
  • Scenario taxonomy: GG cases are treated as equivalent because a model predicting shared labels can also predict unshared labels when task identity is provided at test time.NNu is excluded because unseen labels reveal a new task, making it effectively GNu.
  • Related work: Regularization methods penalize changes to parameters important for previous tasks, but they mitigate rather than eliminate catastrophic forgetting.These methods operate in the GNs scenario.
  • Related work: Replay and exemplar methods preserve information about previous tasks by retaining or regenerating past data.Examples include memory buffers, generative replay, and exemplar-based classification.
  • Related work: Task-specific component methods allocate different nodes, subnetworks, or expanded modules to different tasks, while SupSup extends this approach beyond GG settings.The paper compares primarily with BatchE and Parameter Superposition.

3 Methods

SupSup learns task-specific supermasks over a frozen random network, then retrieves or infers the appropriate mask as task identity becomes less available. Its extensions support task discovery and implicit mask storage in a fixed-size Hopfield network.

  • 3.1 Preliminaries: SupSup keeps a randomly weighted backbone fixed and learns a binary supermask for each task, producing task-specific subnetworks.Outputs use p = f(x, W ⊙ M), while the masks are the learned parameters.
  • 3.2 Scenario GG: When task identity is known, SupSup selects the corresponding mask and achieves higher accuracy with fewer stored bytes on SplitCIFAR100 and SplitImageNet.A new mask can be initialized randomly or from the running mean of previously learned masks.
  • 3.3 Scenarios GNs & GNu: When task identity is unknown at inference, SupSup forms a convex combination of learned masks and minimizes output entropy to identify the correct task.The coefficients represent beliefs over tasks and are renormalized after gradient updates; One-Shot uses a single gradient.
  • 3.3 Scenarios GNs & GNu: SupSup supports sublinear task inference through One-Shot or Binary algorithms, with Binary eliminating half the candidate tasks at each step.After inference, the selected mask produces the class probabilities for the inferred task.
  • 3.4 Scenario NNs: Without task identities during training, SupSup detects uncertainty in the inferred-task distribution and allocates a new supermask for unseen data.A new mask is allocated when the One-Shot distribution is approximately uniform.
  • 3.5 Beyond Linear Memory Dependence: HopSupSup stores learned masks as attractors in a fixed-size Hopfield network and recovers a task mask by jointly minimizing Hopfield energy and output entropy.It masks layer outputs rather than base-network parameters because the Hopfield network is quadratic in mask size.

4 Experiments

Experiments evaluate SupSup across continual-learning scenarios, architectures, datasets, and task-inference settings. SupSup learns thousands of tasks, handles similar tasks, and offers low storage and computation overhead.

  • Datasets, Models & Training: SupSup is evaluated on SplitCIFAR100 and SplitImageNet using reduced-channel ResNet-18 and standard ResNet-50 architectures.SplitCIFAR100 contains 20 five-way tasks, while SplitImageNet contains 100 ten-way tasks.
  • Computation: SupSup stores only the random seed for its fixed base model, while mask application costs approximately 1% of a ResNet-50 forward pass on a 1080 Ti.Sparse supermasks are stored using 16-bit integers.
  • Results: 2500 permutations of MNIST are learned sequentially, yielding 25,000-way classification while inferring task identity from single images with one gradient computation.This result uses the One-Shot algorithm.
  • Results: SupSup distinguishes RotatedMNIST tasks differing by 10 degrees using full-batch Binary inference.The experiment uses the FC 1024-1024 model.
  • Baselines & Ablations: SupSup outperforms PSP and BatchE in the harder GNu setting despite comparisons with methods evaluated using task identity.After 250 permutations, SupSup reaches 94.91% accuracy, compared with 33.88% for Online EWC and 29.31% for SI after 10 permutations.
  • Scenario NNs: Without task identity during training or inference, SupSup learns thousands of tasks, although performance dips when a 2500-supermask budget is exhausted.The model allocates a new mask or selects an existing one every 100 batches.

5 Conclusion

SupSup uses task-specific supermasks and gradient-based task inference to support continual learning across several scenarios. It performs well on SplitImageNet and on thousands of MNIST tasks, but task inference remains limited on non-uniform, difficult problems.

  • Conclusion: SupSup achieves state-of-the-art performance on SplitImageNet when task identity is given.It also performs well on thousands of permutations and nearly indistinguishable MNIST rotations without task information.
  • Limitations: Task inference can fail on non-uniform, challenging problems when models are poorly calibrated and overly confident for the wrong task.The authors identify calibrated models, self-supervision, and energy-based objectives as future directions.

Broader Impact

The paper frames efficient continual-learning models as having environmental, privacy, and access benefits. It also warns that broadly available powerful models may proliferate harmful biases without sufficient study and documentation.

  • Motivation: Learning many tasks with a single model is presented as a concrete goal for efficient continual learning.The paper notes that the meaning of a “single model” is not precisely defined.
  • Positive consequences: Efficient models require less compute than training one model per task, potentially reducing environmental harm.The paper further connects efficiency with faster training when models leverage past tasks.
  • Positive consequences: Running efficient models on end devices can preserve privacy because user data need not be sent to the cloud.
  • Positive consequences: More efficient research can reduce barriers that limit large-scale work to wealthier institutions, whose perspectives may not represent the broader research community.
  • Negative consequences: Widely available powerful models may lack Model Cards and proliferate harmful biases, including disproportionate errors affecting marginalized groups.The paper links this concern to models becoming less subject to regularization and study.

B.1 Training

HopSupSup stores supermasks in a fixed-size Hopfield reservoir and recovers a task mask by jointly optimizing Hopfield energy and output entropy. Recovery converges to the correct mask in roughly 30 gradient steps.

  • B.1 Training: HopSupSup stores the growing set of supermasks implicitly as attractors in a fixed-size reservoir rather than explicitly storing every mask.The Hopfield network uses fixed-size variables Ψ and μ.
  • B.1 Training: The Storkey learning rule updates Hopfield weights so learned binary representations become energy minimizers.The alternative Hebbian rule is also described.
  • B.2 Inference: During inference, HopSupSup initializes a binary string from μ and jointly minimizes Hopfield energy and output entropy to recover a task mask.The recovered mask is converted from z using m = 1/2z + 1.
  • B.2 Inference: Approximately 30 gradient-descent steps make the recovered mask converge to the correct training mask on SplitMNIST.When recovery succeeds, evaluation accuracy matches direct use of the correct mask.

B.3 Network Architecture

The architecture masks layer outputs rather than weights, combines binary masks with fixed weights and non-affine normalization, and uses a 2048-unit two-hidden-layer fully connected network without masking or normalization.

  • Non-affine batch normalization is used without learned parameters, while masking operates on layer outputs and leaves weights fixed.
  • With mask m = (m_1, m_2) and weights W = (W_1, W_2, W_3), the network computes masked, normalized outputs across its layers.
  • Without masking or normalization, f is a fully connected network with two hidden layers of size 2048.
  • HopSupSup uses 10 output neurons for SplitMNIST in Scenario GNu.

C Augmenting BatchE For Scnario GNu

ABatchE extends BatchE task inference by repeating an input batch across candidate tasks, applying task-specific fast weights, and selecting the candidate with minimum entropy.

  • ABatchE infers task identity without a backward pass by combining BatchE with a large batch size.It repeats the data for each task learned so far and uses different fast weights for each repetition.
  • 97.43% average accuracy is reported on SplitMNIST after mask recovery with 20 random seeds.
  • For a batch of b examples and k learned tasks, ABatchE constructs repeated inputs and augmented fast-weight matrices with bk rows.The augmented inputs have shape R^bk×m, while the corresponding matrices have shapes R^bk×m and R^bk×n.
  • ABatchE is evaluated on PermutedMNIST using LeNet 300-100 and FC 1024-1024 architectures with output size 100.
  • The candidate task is the output block whose rows minimize the entropy objective H.Each candidate contributes a contiguous block of b output rows, which is evaluated against the same objective.
  • ABatchE avoids backward computation but requires very large batches for large k, increasing forward-pass compute and memory.Its performance is also limited by BatchE, which SupSup outperforms when BatchE receives task identity information.
  • The entropy objective outperforms M(p) = −max_i p_i, while forgetting is mitigated only with 16 images, corresponding to batch size 4000 for 250 tasks.

D.3 GNu Experiments

The GNu experiments examine implementation choices, transfer, and masked-forward-pass cost. Transfer accelerates SplitCIFAR learning, while masking is estimated to add only a small fraction of forward-pass time under stated hardware assumptions.

  • The experiments use Kaiming-normal initialization and task-specific learning-rate settings for BatchE, while PSP results use its code’s MNISTPerm hyperparameters.
  • Masking is estimated to require approximately 1% of a ResNet-50 forward pass on an NVIDIA GTX 1080 Ti.The estimate uses 480 GB/s memory bandwidth and approximately 2.5 · 10^7 four-byte parameters.
  • Even with real-world memory bandwidth 60–70% below advertised specifications, masking remains within the ≤3% range.
  • Transfer initializes each task’s score matrix from the running mean of earlier supermasks and applies Kaiming fan-in normalization.Without this normalization, accuracy degrades significantly.
  • Transfer improves SplitCIFAR learning while reducing subsequent-task training from 100 to 50 epochs.Task 0 still trains for 250 epochs, and the improvement occurs with about half the total training iterations.

E Supermask Training with Edge-Popup

Edge-Popup trains binary supermasks by ranking score matrices over fixed random weights. The section also examines interpolation between binary and one-shot training and reports tabulated accuracies for several task settings.

  • Edge-Popup computes a binary mask by setting the top k% of score-matrix entries to 1 and updates scores with a straight-through estimator.The masked linear layer uses fixed weights W and mask M to compute y = (W ⊙ M)^T x.
  • Interpolating with γ recovers the binary algorithm at γ = 1/2 and the one-shot algorithm at γ = 1/k.A performance drop is observed on MNISTRotate when moving from binary to one-shot training.
  • The continual-learning scenarios are organized as a tree using whether task labels are shared, rather than only whether output heads are shared.The revised distinction captures cases where tasks share an output head but not labels.
  • Table 2 reports SplitCIFAR100 accuracy and notes that SupSup with Transfer approaches the upper bound.

I Analysis

The analysis explains task inference as entropy minimization over supermask mixtures and develops intuition for why gradient updates favor the correct task. It also connects this behavior to superfluous output neurons and fixed random features under explicit assumptions.

  • Task inference: The correct mask is expected to produce the lowest-entropy output, making differentiable gradient optimization a route to task recovery.The analysis states that entropy is differentiable and uses this property to motivate gradient-based inference.
  • Fixed random features: With fixed random features, the model computes logits using random weights and binary task masks, while mask scores are updated during training.The weights are independently ±√(2/m), with no bias term, and the binary mask is obtained by thresholding learned scores.
  • Task inference: SupSup infers task identity by minimizing output entropy over a convex combination of learned supermasks.The coefficients α_i sum to one and are initialized uniformly; the desired minimum occurs at the one-hot vector for the correct task.
  • Gradient construction: Superfluous neurons enlarge the output from ℓ true labels to n entries and support the construction of an objective whose gradients match supervised loss on those neurons.The gradient matches ∂L/∂y_v for s-neurons and is zero elsewhere; detaching the other outputs prevents their gradient signal.
  • Gradient analysis: Under assumptions that unrelated-task masks are independent of current data and positive and negative weights are equally likely to be masked, expected gradients disfavor incorrect tasks.The lemmas analyze the cases i ≠ j and establish nonpositive gradients for incorrect-task coefficients, while the correct-task direction is treated separately.
  • Objective relationship: Increasing output size moves the entropy-related objective H toward G because the additional term becomes small as each output probability p_v decreases.The analysis decomposes H into two terms, with G as the latter term, and links larger output layers to closer agreement between the objectives.
Loading 2006.14769v3…