Source-linked AI summary
Discovering Neural Wirings
Mitchell Wortsman, Ali Farhadi, Mohammad Rastegari
TL;DR
Neural architecture search largely explores constrained connectivity patterns, while searching all possible wirings is combinatorial and traditional methods can be prohibitive or approximate. DNW allows channels to connect independently and learns wiring by updating edge weights so selected edges can be replaced during training. DNW yields a 10% improvement on ImageNet over hand engineered MobileNetV1 at ∼41M FLOPs.
Problem
Neural architecture search largely explores constrained connectivity patterns, while searching all possible wirings is combinatorial and traditional methods can be prohibitive or approximate.
Source: §1 · p. 1Method
DNW allows channels to connect independently and learns wiring by updating edge weights so selected edges can be replaced during training.
Source: §1 · pp. 1–2Results
DNW yields a 10% improvement on ImageNet over hand engineered MobileNetV1 at ∼41M FLOPs.
Source: §1 · p. 2Takeaways & Limitations
DNW provides a mechanism for discovering sparse subnetworks of the complete graph in a single training run.
Source: §1 · p. 2The analysis assumes α is small enough that sign(˜w) = sign(w).
Source: §2.2 · p. 4
Abstract
from arXiv · showhide
The success of neural networks has driven a shift in focus from feature engineering to architecture engineering. However, successful networks today are constructed using a small and manually defined set of building blocks. Even in methods of neural architecture search (NAS) the network connectivity patterns are largely constrained. In this work we propose a method for discovering neural wirings. We relax the typical notion of layers and instead enable channels to form connections independent of each other. This allows for a much larger space of possible networks. The wiring of our network is not fixed during training -- as we learn the network parameters we also learn the structure itself. Our experiments demonstrate that our learned connectivity outperforms hand engineered and randomly wired networks. By learning the connectivity of MobileNetV1we boost the ImageNet accuracy by 10% at ~41M FLOPs. Moreover, we show that our method generalizes to recurrent and continuous time networks. Our work may also be regarded as unifying core aspects of the neural architecture search problem with sparse neural network learning. As NAS becomes more fine grained, finding a good architecture is akin to finding a sparse subnetwork of the complete graph. Accordingly, DNW provides an effective mechanism for discovering sparse subnetworks of predefined architectures in a single training run. Though we only ever use a small percentage of the weights during the forward pass, we still play the so-called initialization lottery with a combinatorial number of subnetworks. Code and pretrained models are available at https://github.com/allenai/dnw while additional visualizations may be found at https://mitchellnw.github.io/blog/2019/dnw/.
1 Introduction
DNW addresses constrained architecture search by discovering channel-level wirings, learning network parameters and connectivity together. The introduction presents it as a route to sparse networks and reports improved MobileNetV1 ImageNet accuracy at about 41M FLOPs.
1 Introduction
- DNW lets channels connect independently, expanding the wiring space beyond architectures built from predefined blocks.The paper frames this as fine-grained architecture search over individual channels.Source: §1 · p. 1
- Because possible wirings grow combinatorially, traditional search over thousands of nodes is prohibitive or approximate.The introduction motivates searching at the single-channel level rather than among only predefined blocks.Source: §1 · p. 1
- Unlike randomly wired networks whose wirings are fixed during training, DNW continues to update its connectivity.The paper also considers static and dynamic neural graphs, with dynamic graphs allowing cycles and representing sequential models such as LSTMs.Source: §1 · p. 2
- With a fixed edge budget, DNW selects high-magnitude edges while updating all edge weights so a hallucinated edge can replace a real one.The update rule is tailored so that a swap is beneficial.Source: §1 · p. 2
- DNW links architecture search with sparse-network learning, finding sparse subnetworks in one training run while retaining a combinatorial initialization lottery.Only a small percentage of weights are used in the forward pass, yet the training process explores a combinatorial number of subnetworks.Source: §1 · p. 2
- 10% ImageNet accuracy improvement over hand-engineered MobileNetV1 at ∼41M FLOPs is reported with DNW.The paper also describes evaluations across small and large datasets and feed-forward, recurrent, continuous, and sparse networks.Source: §1 · p. 2
2 Discovering Neural Wirings
DNW jointly learns network parameters and wiring, treating channels as nodes and extending the same approach from static graphs to dynamic neural graphs.
2.1 Static Neural Graph
- Each node in the static neural graph processes a two-dimensional channel, replacing the traditional layer-based view with independent channel connections.Node operations include a nonlinear activation, normalization, and convolution.Source: §2.1 · p. 3
2.2 Discovering a k-Edge neural graph
- The k-edge method keeps the k largest-magnitude edges active while updating weights on active and hallucinated edges, allowing edge membership to change during training.When one edge enters the active set, another leaves, conserving the edge count.Source: §2.2 · pp. 3–4
- Gradients reach hallucinated edges but do not flow through them, so their weights can change without affecting the forward pass until selected.An edge may strengthen when its source state aligns with the gradient direction for its destination input.Source: §2.2 · pp. 3–4, Figure 3 · p. 4
- Under stated conditions, including a Lipschitz-continuous loss and fixed node states, swapping in an edge whose updated weight overtakes an active edge decreases loss.The claim concerns a sufficiently small learning rate and a single edge swap.Source: §2.2 · p. 4
2.3 Dynamic Neural Graph
- Dynamic neural graphs allow node states to vary over time; discrete-time graphs can express static networks and sequential models, while continuous-time dynamics include Augmented Neural ODEs.The authors report better empirical performance than a random graph in the dynamic setting, but do not formally justify applying the algorithm there.Source: §2.3 · pp. 5–6
2.4 Implementation details for Large Scale Experiments
- Large-scale experiments use static graphs because dynamic optimization is too expensive, with graph blocks constraining the longest path to match the corresponding MobileNetV1 layer count.For each spatial resolution, the edge count is matched to the corresponding MobileNetV1 model; the node operation is batch normalization, ReLU, and a 3 × 3 convolution.Source: §2.4 · p. 6
3 Experiments
The experiments evaluate DNW on small-scale static and dynamic graphs, ImageNet classification, methods for discovering wirings, and sparse-network training. Across these settings, the paper reports improved connectivity results and describes how DNW maintains sparsity during training.
3.1 Small Scale Experiments For Static and Dynamic Neural Graphs
- On CIFAR-10, DNW outperforms a random graph by a large margin in experiments spanning static, discrete-time, and continuous-time settings.The classifiers have 41k parameters; continuous-time models backpropagate through an adaptive ODE solver.Source: §3.1 · p. 7, Table 1 · p. 7, §3 · p. 7
3.2 ImageNet Classification
- 10% improvement over hand-engineered MobileNetV1 at approximately 41M FLOPs is reported for ImageNet; DNW also yields a 7% performance boost over random graphs.Results are divided among models with similar FLOPs; the authors note random-graph models have considerably more FLOPs, while DNW can skip operations when neurons die.Source: §3.2 · p. 7, Table 3 · p. 8
3.3 Related Methods
- One-shot pruning and Lottery Ticket retraining are described as comparable to DNW, but require training the full graph first and are more expensive in training FLOPs.Fine-tuning the selected subnetwork is also evaluated; the comparisons use CIFAR-10 and the MobileNetV1-DNW structure.Source: Table 2 · p. 7, §3.3 · p. 8
- The no-update ablation demonstrates the importance of updating hallucinated edges, while accuracy is reported as much better when magnitude pruning adds L1 regularization.The paper also compares DNW with Targeted Dropout and reports comparable accuracy for lottery-ticket approaches.Source: §3.3 · pp. 8–9
3.4 Training Sparse Neural Networks
- The ResNet50 sparse-training comparison uses a tuned model with modern optimization techniques and distinguishes fixed sparsity in every layer from leaving the first convolution dense.The first convolution's parameters constitute 0.04% of the total network.Source: Table 4 · p. 9, §3.4 · p. 9
3.3 Related Methods
- A simple MobileNetV1-like DNW network achieves accuracy comparable to an expensive reinforcement-learning neural architecture search method.The NAS comparison is illustrated in Table 3.Source: §3.3 · p. 9
3.4 Training Sparse Neural Networks
- For sparse-network training, DNW uses the top k% of convolution parameters by magnitude in the forward pass while allowing gradients to reach weights zeroed there.The method requires no retraining or fine-tuning, maintains sparsity throughout training, and leaves biases and batch normalization dense.Source: §3.4 · p. 9
4 Conclusion
The conclusion presents DNW as a method for discovering neural wirings that improves accuracy over randomly wired networks. It also emphasizes overparameterization during training and the method's freedom from typical NAS constraints.
4 Conclusion
- DNW demonstrates a significant accuracy boost over randomly wired networks.The conclusion describes the method as novel and points to further directions for more complex and efficient wiring discovery.Source: §4 · p. 9
- The resulting model can be sparse while still benefiting from overparameterization during training.Source: §4 · p. 9
- The authors characterize DNW networks as free from typical NAS constraints and suggest further work on more complex and efficient wiring discovery.Source: §4 · p. 9
A Architecture
Table 5 presents the general MobileNetV1 and MobileNetV1-DNW structures used for ImageNet experiments. Its caption defines the convolution abbreviations and specifies stride and post-convolution conventions.
A Architecture
- Table 5 compares the general structures of MobileNetV1 and MobileNetV1-DNW for ImageNet experiments.Source: Table 5 · p. 12
- The caption defines dwconv as depthwise convolution, pwconv as pointwise convolution, and FC as a fully connected layer.Source: Table 5 · p. 12
- Unless specified otherwise, stride is 1, and batch normalization and ReLU follow each MobileNetV1 convolution.For CIFAR-10 training, the first convolution uses stride 1.Source: Table 5 · p. 12
B Targeted Dropout: Details, Regular vs. Unconstrained and Additional Hyperparameters
This section describes targeted weight dropout variants and their comparison. Unconstrained dropout selects edges by spatial resolution rather than by each neuron's incoming weights.
B Targeted Dropout: Details, Regular vs. Unconstrained and Additional Hyperparameters
- The comparison uses the MobileNetV1-DNW architecture and fixes γ by stage so the post-pruning edge count matches that architecture.Source: B Targeted Dropout: Details, Regular… · p. 12
- Unconstrained targeted weight dropout outperforms the regular variant in the setting considered.The paper reports Table 2 results for the unconstrained variant and compares the two variants in Table 6.Source: B Targeted Dropout: Details, Regular… · p. 12, Table 6 · p. 13
- Unconstrained dropout applies dropout to the bottom γ fraction of edges at a spatial resolution, whereas regular dropout selects the bottom γ fraction of each neuron's incoming weights.The unconstrained variant may leave neurons with no incoming or outgoing edges.Source: B Targeted Dropout: Details, Regular… · p. 12
C A More General Case
The section extends the edge-swapping argument to a more general case, showing that a weight-magnitude crossing can make a swap decrease loss under specified assumptions. It also notes that the proof idealizes node states, which vary in practice.
C A More General Case
- The derivation assumes α is small enough that updated and original weights retain the same sign.Source: C A More General Case · p. 13
- For sufficiently small α, swapping (i, ℓ) for (j, k) decreases loss under Lipschitz loss, fixed node states, no path from i to j, and the stated weight-magnitude crossing.The theorem requires |wiℓ| < |wjk| but |˜wiℓ| > |˜wjk|.Source: C A More General Case · p. 13
- The proof compares swapped and unswapped node states, reduces the argument to Lemma 1, then completes it using the argument for equation 6.Source: C A More General Case · p. 13
- In practice, stochastic updates prevent fixed node states, although large graphs make one node's state have little effect on another's.The authors characterize the proof as idealized and say empirical results demonstrate the method works in practice.Source: C A More General Case · p. 13
D Lemma 1
This section uses a Taylor expansion in the lemma argument and notes that an equivalent argument applies in two dimensions.
D Lemma 1
- The preceding argument completes the lemma.Source: D Lemma 1 · p. 14
- An equivalent argument holds for two dimensions.Source: D Lemma 1 · p. 14
- A Taylor expansion is used to establish the displayed loss inequality.Source: D Lemma 1 · p. 14
E Effect of θv on wuv
The section raises a concern about choosing edges by magnitude when transformations affect incoming-weight magnitudes, then argues batch normalization makes this concern unlikely.
E Effect of θv on wuv
- Arbitrarily scaling a convolutional kernel can make important incoming edges appear small by magnitude, so they may not be selected.Source: E Effect of θv on wuv · p. 14
- With batch normalization, incoming-weight energy is conserved, and the authors argue the magnitude-selection concern is unlikely to occur.A formal treatment is attributed to Theorem 3 of.Source: E Effect of θv on wuv · p. 14
F Reformulation as a Straight-Through Estimator
The section reformulates the update rule as a straight-through estimator: weights are thresholded in the forward pass, while gradients pass through the threshold in the backward pass.
F Reformulation as a Straight-Through Estimator
- The update rule is reformulated as a straight-through estimator.Source: F Reformulation as a Straight-Through… · p. 14
- The forward pass uses h(wuv) = wuv1{|wuv|>τ}, while the backward pass treats h as the identity.Source: F Reformulation as a Straight-Through… · p. 14
- Applying the chain rule with a standard SGD update aligns this formulation with the update rule in Algorithm 1, as implemented in PyTorch.Source: F Reformulation as a Straight-Through… · p. 14