Learning to Flow from Generative Pretext Tasks for Neural Architecture Encoding

TL;DR

This work tackles the efficiency and guidance problems of neural-architecture encoding for performance prediction and NAS. It introduces Flow Generative Pre-Training (FGP), which pre-trains any encoder to capture neural-architecture information flow by reconstructing a flow surrogate vector $\mathbf{s}$ derived from a topologically ordered DAG, without requiring flow-specific architectures. Across NB-101, NB-201, and NB-301, FGP improves performance predictors and NAS outcomes, outperforming strong baselines especially in label-scarce settings, and shows transferability across domains and encoders. The results demonstrate that learning information flow via a synthetic surrogate can yield fast, generalizable representations that enhance downstream neural-architecture search tasks.

Abstract

The performance of a deep learning model on a specific task and dataset depends heavily on its neural architecture, motivating considerable efforts to rapidly and accurately identify architectures suited to the target task and dataset. To achieve this, researchers use machine learning models-typically neural architecture encoders-to predict the performance of a neural architecture. Many state-of-the-art encoders aim to capture information flow within a neural architecture, which reflects how information moves through the forward pass and backpropagation, via a specialized model structure. However, due to their complicated structures, these flow-based encoders are significantly slower to process neural architectures compared to simpler encoders, presenting a notable practical challenge. To address this, we propose FGP, a novel pre-training method for neural architecture encoding that trains an encoder to capture the information flow without requiring specialized model structures. FGP trains an encoder to reconstruct a flow surrogate, our proposed representation of the neural architecture's information flow. Our experiments show that FGP boosts encoder performance by up to 106% in Precision-1%, compared to the same encoder trained solely with supervised learning.

Figures (11)

• Figure 1: Remaining challenges in neural architecture encoding.Challenge 1: Regarding model structures, as shown in (a), flow-based encoders, though more effective, are considerably slower than non-flow-based encoders. Challenge 2: Regarding pretext objectives, in (b), while a deep learning model learns chemical rules by identifying masked atoms in a molecule, it gains no clear training guidance when identifying masked operations in a neural architecture, where most operations are possible options due to the absence of explicit rules.
• Figure 2: Graph modeling of a neural architecture. An example of how a neural architecture is modeled as a directed acyclic graph.
• Figure 3: Generation of the proposed pre-training objective. Red edges mark active computation for the current step; gray edges indicate inactive computation. Starting from 'in' (input), a message (spec., randomly-drawn embedding of 'in') flows through the '1x1' and '3x3' operations, where it is converted depending on each operation type. The two converted messages are summed and flow to 'out' (output) for further conversion. This process is then reversed, and the message flows back through the operations until it returns to 'in'. The resulting message, called the architecture's flow surrogate, serves as the pre-training objective.
• Figure 4: Example of topological order assignment in a neural architecture graph. Assignment results are ($\mathcal{V}^{(1)} = \{v_{1}\}$, $\mathcal{V}^{(2)} = \{v_{2}, v_{3}\}$, $\mathcal{V}^{(3)} = \{v_{4}\}$, and $\mathcal{V}^{(4)} = \{v_{5}\}$), and $T = 4$.
• Figure 5: NAS results. The mean test error of the top-1 architecture across 10 trials is reported. The backbone NAS algorithm NPENAS wei2022npenas adopts pre-trained performance predictors, with each line representing a distinct pre-training approach.