Path-Level Network Transformation for Efficient Architecture Search

In plain words

Finding a good neural network design is usually done by trying many designs from scratch, which is expensive. A cheaper route is to start from a network that already works and edit it in ways that leave its behaviour unchanged, so the edited version starts out as good as the original and can be trained further.

Earlier editing tools could only change how wide or how deep a layer is. This paper adds edits that change how the signal travels through the network: one layer can be split into several parallel branches, similar to Inception-style blocks. A learned controller picks which edits to make, the resulting networks are trained briefly, and their accuracy is used as the reward that teaches the controller. The authors use this to search for a small CIFAR-10 cell that reaches strong accuracy with much less compute than earlier searches needed.

TL;DR

• Path-level edits let a search change the branch structure of a network while reusing trained weights §3.1.
• The searched space is a tree of branches (depth at most 3), controlled by a bidirectional tree-LSTM reinforcement-learning agent §3.2 §3.3.
• Author-reported: TreeCell-A reaches 2.30% CIFAR-10 test error with 14.3M parameters (with DropPath and Cutout) and 25.4% top-1 error on ImageNet in the mobile setting, using about 200 GPU-hours §4.3 §4.4 tab. 1 tab. 2.
• Several training details are not in the paper but are set in the code; see Paper vs code.

Problem

Architecture search from scratch needs a training run per candidate, which is costly on large datasets. The paper cites a NASNet-style search that used 48,000 GPU-hours §1 §4.3. Earlier network-transformation search (EAS) reuses trained weights, but its operations only change layer width, depth, or compression. They cannot change the connection topology, so a chain-shaped starting network stays chain-shaped §1 §2.3. The authors assume the starting network is a reasonable teacher and that transformed networks can be checked by a short training run on CIFAR-10 §4.1 §4.2.

Method

Path-level transformations §3.1. A layer is replaced by a multi-branch motif, merged either by addition or by concatenation. For a convolution layer with N branches, the paper gives two constructions fig. 1:

1. Addition: replicate the layer in every branch and divide each branch output by N, so the sum equals the original output.
2. Concatenation: replicate the input, split the output filters into N groups, one group per branch, and concatenate the results.

For an identity layer fig. 2: addition uses N identity branches; concatenation splits the input channels, with identity mappings in each branch. Net2Net-style operations (Net2DeeperNet, Net2WiderNet) are then applied inside the branches so that the branches differ in depth and kernel type, as in Inception models fig. 3 §2.3.

Tree-structured architecture space §3.2. A node has an allocation scheme (replication or split) and a merge scheme (addition or concatenation). Each node-to-child edge is a primitive operation. With m children, the node output is defined recursively as in eq. 1:

\mathbf{z}_i = \text{allocation}(\mathbf{x}, i), \quad \mathbf{y}_i = N^c_i(E_i(\mathbf{z}_i)), \quad 1 \le i \le m, \quad N(\mathbf{x}) = \text{merge}(\mathbf{y}_1, \dots, \mathbf{y}_m) eq. 1

A leaf returns its input. Multi-branch networks are depth-1 trees §3.2. The seven primitive operations are: 1×1 convolution, identity, 3×3 / 5×5 / 7×7 depthwise-separable convolution, 3×3 average pooling and 3×3 max pooling §3.2. Pooling cannot be initialised as an identity, so after a pooling edge is inserted the student's weights are fitted to the teacher's output logits by knowledge distillation. The paper calls this cost negligible §3.2.

Controller encoder §3.3. The tree is encoded bottom-up and top-down with tree-LSTM units. Bottom-up, a child-sum LSTM handles addition nodes and an N-ary LSTM handles concatenation nodes; a separate ordinary LSTM handles each edge eq. 2. Top-down states use analogous N-ary units eq. 3 fig. 4. The hidden state of each node combines both directions.

\begin{aligned} (h', c') &= \begin{cases} \text{ChildSumLSTM}^{\uparrow}(s, [h^c_1, c^c_1], \dots, [h^c_m, c^c_m]) & \text{if add} \\ \text{NaryLSTM}^{\uparrow}(s, [h^c_1, c^c_1], \dots, [h^c_m, c^c_m]) & \text{if concat} \end{cases} \\ (h, c) &= \text{LSTM}^{\uparrow}(e, [h', c']) \end{aligned} eq. 2

\begin{aligned} (\tilde{h}'_i, \tilde{c}'_i) &= \text{NaryLSTM}^{\downarrow}\big(s, [\tilde{h}^p, \tilde{c}^p], [h_1, c_1], \dots, [\mathbf{0}, \mathbf{0}]_i, \dots\big) \\ (\tilde{h}_i, \tilde{c}_i) &= \text{LSTM}^{\downarrow}(e_i, [\tilde{h}'_i, \tilde{c}'_i]) \end{aligned} eq. 3

Here s is the node's allocation and merge scheme, e is the edge to the parent, and e_i is the edge to child i §3.3.

Decisions §3.3 fig. 5 alg. 1. Softmax heads read the node state and make three kinds of decision:

1. A node with one leaf child chooses a merge scheme from {none, add, concatenation}. For add or concat it also chooses the branch count, and the node then gets child nodes. For none, no further merge decision is made there.
2. A leaf chooses whether to expand: inserting a new leaf joined by an identity edge deepens the cell.
3. An identity edge chooses one of the seven primitive operations (identity is allowed). This happens once per edge.

Training uses the REINFORCE policy gradient with an entropy bonus. The validation accuracy of each sampled net gives the reward, and a baseline reduces variance §3.3 §4.1. Algorithm 1 in Appendix D re-uses a result when the same cell is sampled again, and updates the controller after every B new evaluations alg. 1.

Search loop §4.2. The base network is a small DenseNet-BC. At each step the controller samples a cell, which replaces every 3×3 convolution in the dense blocks. The result inherits the base weights and is trained for 20 epochs on CIFAR-10. The reward is tan(acc × π/2), where acc is the validation accuracy as a fraction §4.1. The search ran 500 networks §4.2.

The search loop follows this pattern alg. 1:

# Path-level efficient architecture search  [alg. 1] [§4.2]
# Input:  base network, training set, validation set, batch size B, budget M
# Output: the cells found and their validation accuracies

def path_level_search(base_net, train_set, val_set, B, M):
    # 1. Start with a fresh controller and an empty gradient buffer  [alg. 1]
    controller = init_controller()
    grad_buffer = []
    results = {}
    # 2. Stop once M distinct networks have been trained [alg. 1]
    while len(results) < M:
        # 3. Sample a tree-structured cell with the controller  [§3.3]
        cell = controller.sample_cell(base_net)
        # 4. Reuse the stored accuracy for a cell seen before  [alg. 1]
        if cell in results:
            acc = results[cell]
        else:
            # 5. Transform the base net, train 20 epochs, validate
            #    [§4.1] [§4.2]
            net = apply_cell(base_net, cell)
            acc = train_and_validate(net, train_set, val_set)
            results[cell] = acc
        # 6. Queue the policy gradient for this cell  [§3.3] [alg. 1]
        grad_buffer.append(cell_gradient(controller, cell, acc))
        # 7. After B queued cells, update the controller  [§4.1] [alg. 1]
        if len(grad_buffer) == B:
            controller.update(grad_buffer)
            grad_buffer = []
    return results

Final training §4.2 §4.3. The top 10 cells are embedded into larger base networks (DenseNet-BC and PyramidNet). Each base network is trained for 300 epochs with an initial learning rate of 0.1, and the other settings stay as in the search. Regularised runs (DropPath, Cutout) train for 600 epochs.

ImageNet §4.4. The learned cells are combined with CondenseNet in the mobile setting: 224×224 input and fewer than 600M multiply-adds.

Hyperparameters

Results

Author-reported. Setup: CIFAR-10 test error (%) with DenseNet and PyramidNet base networks; ImageNet top-1 and top-5 error (%) in the mobile setting (≤600M multiply-adds). Baselines are human-designed and automatically designed networks as cited in the paper. Search cost: about 200 GPU-hours (500 networks), versus 48,000 GPU-hours for NASNet-A §4.3.

CIFAR-10 tab. 1:

Transcribed from tab. 1. Empty parameter cells mean the table gives "–".

Search progress fig. 7. The RL controller's average validation accuracy rises as more networks are sampled; random search does not show the same trend. The best RL-found model, after 20 epochs, is 0.16% above the best random-search model §4.2. The best discovered cells reach 3.64% (TreeCell-A) and 3.71% (TreeCell-B) error after 300-epoch training with DenseNet, against 3.98% for the best random-search cell §4.2.

CIFAR-10 text results §4.3. TreeCell-A with DenseNet reaches 3.64% error at 3.2M parameters. With PyramidNet it reaches 3.14% at 5.7M parameters, against 3.31% at 26.0M for the best PyramidNet in tab. 1. With DropPath (600 epochs) it reaches 2.99%; with Cutout, 2.49% at 5.7M and 2.30% at 14.3M.

ImageNet, mobile setting tab. 2:

The paper reports TreeCell-A at 25.5% top-1 and 8.0% top-5 error, which is better than the CondenseNet (G1=G3=4) baseline at a similar cost §4.4. Compared with NASNet-A, the paper reports slightly better results at similar multiply-adds §4.4.

Appendix A: search from scratch fig. 8 tab. 3. The start point is a chain of identity mappings. The paper reports the best RL-found cell at 3.93% error with 11.5M parameters, the best random cell at 4.44% with 10.0M parameters, and TreeCell-A at 3.14% with 5.7M parameters. RL beats random search by 1.54% on validation accuracy in this setting app. A.

Paper vs code

Code: han-cai/PathLevel-EAS at commit 7a6a34fb8da042a96b717424572253709b13f56a, linked from the paper's experiments section §4. The repository belongs to the first author, and its README links the paper. It has no LICENSE file, so it is cite-only: no verbatim snippets. I searched the paper's links only; the paper gives the repository directly, so no GitHub search was needed.

• Reward baseline update. The paper describes an exponential moving average with decay 0.95 §4.1. The code computes baseline += 0.95 * (mean_reward - baseline), so 0.95 is the weight on the newest batch mean, not on the old baseline arch_manager.py. The update happens before the advantage is computed arch_manager.py.
• Gradient clipping and learning rate. The paper says Adam and REINFORCE §4.1 but gives no learning rate or clipping. The code uses learning rate 6e-3 and clipping at norm 40 arch_search_tree_cell.py rl_trainer.py.
• Cell resampling. The code rejects any sampled cell whose height is not greater than 1 and samples again arch_search_tree_cell.py. The paper does not mention this.
• Transformation noise. The code adds Gaussian noise (std = 6e-3 × the std of the weights) to the copied and deepened edges arch_search_tree_cell.py layers.py. Copies therefore differ from the original by a small amount, so function preservation is approximate. The paper states exact function preservation §3.1.
• Addition merge. The paper divides each replicated branch's output by N §3.1. The code divides the merged output by the child count in the forward pass tree_node.py, and its copies carry noise (see above). The two agree only if the noise is ignored.
• Allocation for concatenation. The paper uses replication for convolutions and split for identity layers §3.1. The code uses split for identity and pooling edges and for grouped convolutions, and replication for other convolutions tree_node.py. The grouped-convolution case is not in the paper.
• Pooling reconstruction. The paper says the student is fitted to the teacher's logits by knowledge distillation §3.2. The code runs an MSE fit of the whole student to the teacher's final logits, on 1,000 samples for 30 epochs with SGD at a constant LR of 0.02 expdir_monitor.py utils.py. This runs only when the cell contains pooling or the network uses BN-activation-weight ordering arch_search_tree_cell.py.
• Search-time weight decay and momentum. The paper gives weight decay 1e-4 and Nesterov momentum 0.9 for search-time training §4.1. The search script's validation config overrides only epochs, batch size, initial LR, validation size and schedule type arch_search_tree_cell.py. Weight decay and momentum are therefore inherited from the start network's run config. The DenseNet run config sets them to these values run_densenet.py. I did not trace the start-network config used by the search itself.
• Chain-start search. The appendix says all settings match Section 4 app. A. The code's chain-start setting uses 300 networks, not 500, and allows an extra 3×3 convolution edge when random initialisation is forced arch_search.py arch_search_tree_cell.py.
• Edge candidates. Seven operations in the default search, as the paper says §3.2 arch_search_tree_cell.py. The code's vocabulary also contains grouped and normal 3×3 convolution tokens for encoding; these are not sampled by default arch_search_tree_cell.py.
• ImageNet. The mobile model code is in code/ImageNet. A hand-written TreeCell-A and TreeCell-B are in code/ImageNet/tree_cells.py tree_cells.py. I did not check these definitions against Figures 6 and 9, and I did not confirm which file the CondenseNet model imports.
• Final training schedule. The config for final training lists both cosine decay and step reductions at 50% and 75% of training, with a factor of 10 each run_densenet.py. The paper says cosine decay §4.2. I did not trace how the schedule code uses both settings.

Pitfalls

• Baseline ambiguity. "Decay 0.95" can mean the weight on the old baseline or on the new mean. The code puts 0.95 on the new batch mean. Using the other convention changes the advantage signal considerably.
• Advantage uses the updated baseline. Because the baseline is updated before it is subtracted, the advantage of the current batch depends on that batch's own mean arch_manager.py.
• Path-level transformations are not exactly identity. Noise on copied edges and BN statistics estimated from about 2,000 images (10 batches of 200 from the validation split) mean that the transformed network only roughly matches the parent arch_search_tree_cell.py layers.py. The paper does not measure how far the transformed network's initial accuracy departs from the parent's.
• Reward scale. tan(acc × π/2) is steep near 100%. Computed from the formula, accuracies 0.90 and 0.91 give rewards of about 6.31 and 7.04, while 0.60 and 0.61 give about 1.38 and 1.44. The paper motivates this choice in the same way §4.1.
• Duplicate cells. The same cell can be sampled many times. The code returns the cached result and queues the gradient for the next update; the cache key is the tree string arch_search_tree_cell.py.
• Depth limit applies to node insertion only. Nodes are inserted only if their depth is below 3 arch_search_tree_cell.py. Check this when you port the space.
• Cost is not fully stated. The 200 GPU-hours count the search; the final training runs (300 epochs, DropPath 600, Cutout) add more. The paper does not give a total for those.
• No variance reported. The paper gives single numbers without seed or run variance §4.2 tab. 1. The 0.34% gap between TreeCell-A (3.64%) and the best random-search cell (3.98%) therefore has no stated uncertainty §4.2.

Lineage

• Builds on: EAS, network transformation search with layer-level Net2Net-style operations and a recurrent controller (Cai et al., 2018), arXiv:1707.04873 README.md. This paper keeps the controller and reward scheme and adds path-level edits and a tree-LSTM encoder §2.3 §3.3.
• Builds on: Net2Net function-preserving operations (Chen et al., 2015), cited in the paper §2.3. Path-level edits are built on these.
• Compared against: NASNet-A, NASNet-B and NASNet-C tab. 2; Hierarchical architecture search, Block-QNN and NAS v3 tab. 1; DenseNet, PyramidNet and CondenseNet as base networks tab. 1 tab. 2.
• Superseded by: not stated in this card.