Architecture-Aware Learning Curve Extrapolation via Graph Ordinary Differential Equation

We use a sequence encoder to compute the mean and standard deviation of the posterior distribution $q_{\phi}(\boldsymbol{z}_{n+1}|\{y_{i},t_{i}\}_{i=1}^{n})$, which is assumed to be Gaussian:

We implement this using an RNN with GRU units for the sequence encoder. Other encoder options include Self-Attention (Vaswani et al., 2017) and Temporal Convolutional Networks (TCN) (Pandey and Wang, 2019), which adapt well to varying observation lengths. We conduct an ablation study in the experimental section to evaluate the performance of these alternative sequence encoder implementations.

Assuming the graph representation $G$ contains $N$ nodes, the adjacency matrix $A\in\mathbb{R}_{\geq 0}^{N\times N}$ details node connections, where $A_{ij}>0$ represents the edge from node $i$ to node $j$. We analyze two foundational neural network types: MLPs and CNNs. In MLPs, nodes correspond to neurons and edges correspond to the presence of the connections between neurons, whereas in CNNs, nodes represent feature maps and edges depict operations like $1\times 1$ and $3\times 3$ convolutions, or $3\times 3$ average pooling. For CNNs, we adopt the cell-based representation (Dong and Yang, 2020; Liu et al., 2018), which consists of four principal building blocks: stems, normal cells, reduction cells, and classification heads. The stem block is a fixed sequence of convolutional layers to process the input images. This is followed typically by 14-20 cells with reduction cells placed at 1/3 and 2/3 of the total depth. A normal cell contains 4 nodes, each of which belongs to the set of operations: skip connections, identity or zero (indicating the presence or absence of connections between certain layers), $1\times 1$ and $3\times 3$ convolutions, and $3\times 3$ average pooling. Finally, the classification head employs a global pooling layer followed by a single fully connected layer and returns the network’s output. Since the cell is repeated throughout this macro-skeleton, a CNN can be uniquely represented by its cell.

To derive a graph-level embedding, we first implement node-level message passing, followed by global pooling to extract a global embedding. Our method differs from existing architecture-based performance prediction approaches (Liu et al., 2018; Wen et al., 2020; Knyazev et al., 2021) by treating the node as a feature map rather than focusing on operations. Nodal features are calculated from the in-degree and out-degree of each node, normalized by the total number of edges. For MLPs, edge features are binary, while for CNN cells, they are integers representing operation types: $\{\text{zeroize}:0,\text{$1\times 1$ conv}:1,\text{$3\times 3$ conv}:2,\text{$3\times 3$ avg pooling}:3\}$. The nodal feature matrix is denoted as $X\in\mathbb{R}^{N\times d}$. The graph encoding process returns a vector representation of the architecture.

For node-level message passing, we employ GCN layers and normalize the adjacency matrix to stabilize training, following  (Kipf and Welling, 2016).

A GCN layer then transforms the nodal features into a hidden representation:

where $W\in\mathbb{R}^{D\times d}$ is the feature transformation matrix. Applying $l$ GCN layers aggregates information from $l$-hop neighbors. We employ learnable pooling, among other methods such as average- and max-pooling, and investigate each in our ablation study. The graph embedding $\boldsymbol{z}_{G}$ contributes to the evolution of the latent state, as detailed in the next section, which introduces the latent ODE.

The transformation from discrete to continuous domain is predicated on the assumption that the step size, or learning rate, approaches zero, thereby approximating the time derivative of NN parameters $\Theta$ using $\frac{d\Theta}{dt}=-\frac{\partial\mathcal{L}}{\partial\Theta}$ (Su et al., 2016), where $\mathcal{L}$, $\Theta$ denotes the objective function and parameters from the source task. Assuming the parameters implicitly depends on time $t$, the time derivative of the training loss can be written as:

The ground truth ODE for loss is independent of time. Therefore we adopt autonomous differential equations to describe the continuous evolution of the learning curves. On the other hand, we do not directly use the exact formula, as this involves the computation of backpropagation of the source task, which is potentially computation intensive. Moreover, our primary goal is not to derive exact ODEs for each training configuration. Instead, we focus on efficiently inferring learning curves for new training configurations. To achieve this, we leverage the latent space, using the expressivity of hidden neurons to capture common patterns across learning curves from the same source task, despite variations in underlying architectures and hyperparameters.

Given that the training data for the source task remains fixed, the loss landscape varies depending on the architecture. Therefore, it’s sufficient to model a universal latent ODE as a function of the architecture, enabling it to describe the evolution of various learning curves, each corresponding to a different architecture. Our latent ODE is formalized by the equation:

Here, $\dot{\boldsymbol{z}}$ denotes the derivative of the latent state vector $\boldsymbol{z}$ with respect to time, driven by the function $f_{\theta_{2}}$, which takes as input both the latent state $\boldsymbol{z}$ and a graph-level embedding $\boldsymbol{z}_{G}$. This combination allows the model to simultaneously consider the dynamic properties of the learning curve and the static characteristics of the architecture, enhancing the predictive capability of the system. Eq (9) can be regarded as a single-agent representation of the dynamics induced by NN training, which involves multiple trajectories of neurons, edge weights, and loss. This reduced representation of coupled dynamical system has been explored in (Gao et al., 2016; Laurence et al., 2019) to study the tipping point of the original network dynamics. The difference from the prior dynamical system reduction approach is that the graph reduction mechanism is learnable so that the model can be adapted to unseen trajectories derived from different optimization trials.

Finally, numerical integration of Eq (9) yields a time series of latent states $\boldsymbol{z}_{i}$ ($n<i\leq m$). Each latent state is independently decoded by a function $\hat{y}_{i}=\mathrm{Decoder}_{\theta_{3}}(\boldsymbol{z}_{i})$.

To optimize our model, we maximize the evidence lower bound (ELBO), fomulated as follows

The first term in the ELBO equation represents the expected log-likelihood of observing future outputs given the latent states, as parameterized by $\theta=(\theta_{1},\theta_{2},\theta_{3})$, where $\theta_{1}$, $\theta_{2}$, and $\theta_{3}$ correspond to the architecture encoder, the ODE function, and the decoder, respectively. The second term penalizes the divergence (measured by the KL-divergence) between the posterior distribution of the latent states and their prior distribution, enforcing a regularization that anchors the posterior closer to the prior. This balance ensures that while the model remains flexible enough to capture complex patterns in data, it also maintains a level of generalization that prevents overfitting.

The computational cost of numerical integration is minimal because the ODE models only the evolution of the loss embedding with dimension $D$, rather than modeling each node in the architecture individually. The runtime of the forward pass is bounded by $O(D^{2}T)$, where $T=m-n$ is the number of integration time steps, assuming the ODE function is implemented as an MLP with a fixed number of layers. Consequently, the runtime of the ODE component remains independent of the overall size of the neural network.

The goal of this experiment is to evaluate the LC-GODE model against established learning curve prediction methods using real-world benchmarks. We train our model separately on the test loss curves of each source task. The condition length is set to 10 epochs for all methods in this experiment. The instantiation of LC-GODE that we report features: (i) an architecture encoder that utilizes 2 layers and employs a learnable pooling technique, (ii) an observed time-series encoder implemented using GRU, (iii) an ODE function with a 2-layer MLP and integrated using the Runge-Kutta 4 method (Butcher, 1996). Further details on the evaluation metrics and the training settings can be found in the supplemental materials.

We evaluate our model against six methods, including three Bayesian approaches and three general time-series prediction approaches. (i) LC-BNN: This method utilizes a Bayesian neural network to model the posterior distribution of future learning curves. The probability function is constructed from a combination of basis functions, following the approach described by Domhan et al. (2015). (ii) LC-PFN: A transformer-based model is trained on synthetic curves that are generated from a pre-defined prior. This method serves as an efficient alternative to traditional Markov Chain Monte-Carlo (MCMC) techniques for sampling from posterior distributions. (iii) VRNN: A probabilistic model utilizing random forests and Bayesian recurrent neural networks. (iv) LSTM: This Recurrent Neural Network captures sequential dependencies. Before the observation cutoff at $n$, input states are derived from actual observations. Post $n$, the model uses its own predictions from previous timestamps as inputs. (v) NODE: Latent Ordinary Differential Equations, focusing primarily on modeling the latent loss representation without incorporating architectural information. (vi) NSDE: Latent Stochastic Differential Equations, consisting of a drift term and a diffusion term. The drift term is the same as NODE, and the sequence encoder and decoder is the same as in our model. Both NODE and NSDE are trained using the variational framework.

We evaluate performance using Mean Absolute Percentage Error (MAPE) and Root Mean Squared Error (RMSE) with different prediction lengths. Table I shows the extrapolation error for test accuracy curves from four datasets. Due to space constraints, the comparison results for loss curve extrapolation are provided in the supplement. Our proposed LC-GODE outperforms all baselines on these datasets. Specifically, LC-GODE reduces the error on test accuracy curves by 36.13%, 30.72%, 34.97%, and 59.63%, and on test loss curves by 65.5%, 44.61%, 20.1%, and 23.45% compared to the NODE model without architecture information. This improvement is due to the incorporation of architecture information with graph embedding. The architectures with similar performance are mapped to nearby locations in the hidden space, as shown in Figure 2(a). To demonstrate the advantage of jointly training the time-series encoder with the graph encoder, we compare the optimal performance difference of two configurations with the distance between their initial latent states in Figure 2(b). When architecture information is included, the correlation between these distances increases from 0.77 to 0.83 for the CIFAR-10 test accuracy curves. For a detailed comparison over epochs, we plot MAPE at 10 prediction lengths for NODE, NSDE, and LC-GODE on both test accuracy and loss curves (Figure 3). Overall, models perform better and have larger improvement on test accuracy curves compared to loss curves. This could be attributed to the fact that classification accuracy is less sensitive to decision boundary changes than cross-entropy loss, making it less variable and easier to predict.

Regarding the baseline comparisons, LC-PFN emerges as the second most effective approach for extrapolating test accuracy curves of the car and segment source tasks, while the NSDE model ranks as the second in predicting test loss curves for CIFAR-10 and CIFAR-100. Both NSDE and NODE models demonstrate comparable performance, with NSDE marginally outperforming NODE. These results underscore the viability of employing time-series approaches for addressing learning curve extrapolation challenges. Notably, the slight advantage of NSDE over NODE suggests subtle benefits in capturing stochastic dynamics that may be present in complex learning scenarios. This comparison highlights the potential for refined time-series models to enhance predictive accuracy and adaptability in diverse training environments.

We evaluate our method’s efficacy in ranking training configurations by comparing the predicted and true optimal performances at a fixed snapshot, employing two metrics: regret (the performance discrepancy between the actual and predicted best configurations) and ranking (the position of the predicted best configuration according to the true learning curves). These metrics demonstrate whether the predicted performance can effectively guide the selection of a performant configuration.

As shown in Table II, our proposed approach, LC-GODE, enhances the ranking of the predicted best model by 3 positions on the segment dataset and by 8 positions on the CIFAR-10 dataset, and reduces regret by 96% on CIFAR-10 compared to the superior baseline among NODE and NSDE for test accuracy curves. Nevertheless, when employing test loss to identify the best model, the overall ranking for all methods declines, indicating that predicted test loss is less effective for model selection.

Our approach identified a performant model with a 20x speedup compared to exhaustive brute-force training using stochastic gradient descent (SGD) on the MLP datasets. The speedup is computed as the total runtime needed to fully train all configurations via SGD, divided by the sum of the actual training time for $n$ epochs and the inference time required by LC-GODE. The inference time for LC-GODE is approximately 0.8 seconds, and the majority of the model selection time is spent on the initial $n$ epochs of SGD training. Additional results on the training and inference runtime of all methods are provided in the supplement.

Figure 4 further shows the true vs predicted best test accuracy for the 2 tabular data and 2 image classification source tasks, further indicating a high correlation between predicted and true metrics.

We analyze the impact of three hyperparameters: maximal time $T_{\max}$, observation length, and hidden dimension (Figure 6). A larger $T_{\max}$ negatively impacts performance. As observation length increases, the model receives more information about the curve, requiring less extrapolation and thus reducing error. The performance remains relatively constant when the latent dimension is within the range of 20 to 50.