Source-linked AI summary
D-VAE: A Variational Autoencoder for Directed Acyclic Graphs
Muhan Zhang, Shali Jiang, Zhicheng Cui, Roman Garnett, Yixin Chen
TL;DR
DAG optimization is challenging because candidate evaluation can be expensive and standard black-box methods are designed for continuous spaces. D-VAE uses dependency-respecting asynchronous message passing in a GNN-based variational autoencoder to encode DAG computations, and experiments show valid novel generation and effective latent-space optimization across neural architecture search and Bayesian network structure learning.
Problem
DAG optimization combines expensive graph evaluation with discrete structures that hinder direct use of standard continuous-space black-box optimization.
Method
D-VAE is a GNN-based variational autoencoder whose asynchronous message passing respects DAG dependencies and encodes computations in a continuous latent space.
Results
D-VAE generates novel and valid DAGs and finds better architectures and Bayesian networks through latent-space Bayesian optimization.
Takeaways & Limitations
D-VAE provides a latent-space framework for optimizing DAGs in neural architecture search and Bayesian network structure learning.
Takeaways & Limitations
D-VAE injectively encodes computations rather than graph structures and currently does not model semantic similarities among node types.
Abstract
from arXiv · showhide
Graph structured data are abundant in the real world. Among different graph types, directed acyclic graphs (DAGs) are of particular interest to machine learning researchers, as many machine learning models are realized as computations on DAGs, including neural networks and Bayesian networks. In this paper, we study deep generative models for DAGs, and propose a novel DAG variational autoencoder (D-VAE). To encode DAGs into the latent space, we leverage graph neural networks. We propose an asynchronous message passing scheme that allows encoding the computations on DAGs, rather than using existing simultaneous message passing schemes to encode local graph structures. We demonstrate the effectiveness of our proposed DVAE through two tasks: neural architecture search and Bayesian network structure learning. Experiments show that our model not only generates novel and valid DAGs, but also produces a smooth latent space that facilitates searching for DAGs with better performance through Bayesian optimization.
1 Introduction
DAG optimization appears across neural architecture search, Bayesian network structure learning, and circuit design, but discrete structures make evaluation and black-box optimization difficult. D-VAE addresses this by encoding DAG computations into a smoother continuous latent space for Bayesian optimization.
- DAG optimization arises in neural architecture search, Bayesian network structure learning, and electronic circuit design.
- DAG optimization is difficult because evaluating candidate graphs can be time-consuming and common black-box optimizers primarily operate in continuous spaces.
- Embedding DAGs into a relatively smooth continuous space could enable Bayesian optimization, or gradient methods when gradients are available.
- D-VAE combines a variational autoencoder for DAGs with Bayesian optimization in its latent space.
- The proposed asynchronous message passing scheme respects DAG dependencies and is designed to encode whole computations rather than only local graph structures.
- Experiments apply D-VAE to neural architecture search and Bayesian network structure learning, reporting novel valid DAG generation and smooth latent spaces effective for optimization.
2 Related work
Prior graph generative models use token, adjacency-matrix, or graph representations, while related NAS and Bayesian-network methods address discrete structure optimization through specialized search or continuous mappings. VAEs provide a general latent-variable framework for generation, including discrete outputs modeled through sequential decisions.
- VAEs learn an approximate posterior encoder and a generative decoder, then generate data by decoding latent vectors sampled from a prior.
- Token-based graph models represent graphs as token sequences but require task-specific grammars, limiting their generality.
- Adjacency-matrix-based and graph-based models provide alternative representations for deep graph generation.
- Neural architecture search methods include reinforcement learning, Bayesian optimization with architecture-similarity kernels, and continuous-relaxation approaches for gradient optimization.
- Bayesian network structure learning commonly uses score-based search over a discrete design space, with BIC and BDeu among its scoring functions.
3 DAG variational autoencoder (D-VAE)
D-VAE is a GNN-based variational autoencoder that encodes and decodes DAG computations with dependency-respecting asynchronous message passing. Its encoder is permutation-invariant under suitable aggregation and injective for computations, supporting Bayesian optimization in a continuous latent space, while vertex semantics remain an open extension.
- 3 DAG variational autoencoder (D-VAE): D-VAE encodes and decodes DAGs using asynchronous message passing, targeting computations rather than only graph structures.
- 3 DAG variational autoencoder (D-VAE): A computation is represented by a DAG whose directed edges specify signal flow between operations, and distinct DAG structures can represent the same computation.
- 3.1 Encoding: The encoder assumes a single starting node, adding a virtual starting node when multiple predecessor-free nodes exist.
- 3.1 Encoding: For each node, D-VAE aggregates predecessor hidden states and updates the node state only after those predecessor states have been computed.
- 3.1 Encoding: Following a topological order, the encoder uses the ending node’s hidden state to parameterize the posterior approximation qφ(z|G).
- 3.1 Encoding: The encoder is invariant to node permutations when its aggregation function is order-invariant, so isomorphic DAGs receive the same encoding.
- 3.1 Encoding: Under injective aggregation and update functions, D-VAE injectively encodes computations, assigning each computation a unique latent embedding suitable for continuous optimization.
- 3.1 Encoding: The injectivity guarantee concerns computations rather than graph structures, allowing different structures that represent the same computation to share an encoding.
4 Experiments
D-VAE is evaluated on neural architecture search and Bayesian network structure learning, testing reconstruction, generation, latent-performance prediction, Bayesian optimization, and latent-space smoothness. Across these experiments, it generates valid and novel DAGs and supports effective optimization through a continuous latent space.
- Experimental setup: The experiments cover reconstruction, prior validity, uniqueness, novelty, latent-performance prediction, Bayesian optimization, and latent-space visualization on two DAG datasets.The tasks are neural architecture search and Bayesian network structure learning.
- Basic abilities of VAE models: D-VAE and S-VAE generally perform best on basic VAE requirements, while D-VAE, S-VAE, and GraphRNN achieve near-perfect reconstruction, validity, and novelty.D-VAE and S-VAE also show higher uniqueness, indicating more diverse valid generations.
- Predictive performance of latent representation: D-VAE’s encoded means significantly outperform other models on both RMSE and Pearson’s r for predicting DAG performance.The authors attribute this to encoding computations rather than only local graph structures.
- Bayesian optimization: 94.80% was D-VAE’s highest neural-architecture accuracy, versus 92.79% for S-VAE; all five D-VAE architectures exceeded 94%.The architectures were fully trained on CIFAR-10 to measure true test accuracy.
- Bayesian optimization: D-VAE found a Bayesian network with BIC -11125.75, surpassing the best training-network score of -11141.89 and all 180,000 training networks.The reported BIC comparison uses the convention that higher scores are better.
- Latent space visualization: D-VAE’s latent spaces show smoother performance variation than S-VAE’s, including smoother BIC changes and better separation of high- and low-score regions.For neural architectures, both models show relatively smooth great-circle interpolations; the Bayesian-network visualization uses a principal 2-D subspace.
5 Conclusion
D-VAE is a GNN-based generative model for DAGs that respects partial-order dependencies when encoding computations. Bayesian optimization in its latent spaces offers directions for neural architecture search and Bayesian network structure learning.
- D-VAE uses asynchronous message passing to encode DAGs according to their partial order and explicitly model computations.
- Bayesian optimization in D-VAE latent spaces is applied to neural architecture search and Bayesian network structure learning.
Appendices
The appendices establish how D-VAE differs from related DAG optimization methods and prove permutation invariance and injective encoding of computations under suitable conditions.
- Related Work: Unlike NAO, D-VAE directly represents neural architectures as graphs and uses unsupervised rather than supervised learning.NAO uses strings and requires evaluating many randomly sampled graphs on a large dataset.
- Related Work: D-VAE differs from prior Bayesian optimization methods by searching in a learned latent space rather than heuristically operating locally in discrete graph space.Earlier BNSL approaches used Gaussian processes as proxies while retaining discrete local search.
- DAGs, Computations, and Functions: The paper distinguishes DAG structures, computations, and functions, and models computations because functions are harder to model and depend on training.Different DAGs can represent the same computation, while identical architectures can represent different functions after different training.
- Assumptions: Multiple initial input signals are reduced to the single-input setting by adding an initial assignment operation.The paper assumes one input signal throughout for presentation.
- Permutation Invariance: The encoder is invariant to node permutations, proved by induction from the permutation-invariant aggregation function and the single starting node.The proof propagates invariant predecessor states through each node and reaches the ending node.
- Injective Encoding: D-VAE can injectively encode computations when its update, operation encoding, and aggregation components are injective.The proof establishes an injective mapping from each composite computation to its node hidden state by induction.
E Modifications for Encoding Neural Architectures
For neural architectures, D-VAE adapts aggregation to preserve input order when operations such as concatenation make order meaningful, while the plain aggregator remains competitive.
- Aggregation Design: The aggregation function must reflect whether an operation treats predecessor inputs as ordered or unordered.Summation and averaging permit order-invariant aggregation, whereas concatenation requires order awareness.
- Order-Aware Aggregation: D-VAE adds each layer’s global ID to the aggregation function so it can respect concatenation order.The modification uses one-hot encodings of layer IDs.
- Empirical Effect: The order-aware aggregation improved D-VAE’s neural-architecture performance empirically, but the plain aggregation still outperformed all baselines.The passage reports both the within-model improvement and the baseline comparison without specifying a numeric value.
F Modifications for Encoding Bayesian Networks
For Bayesian networks, D-VAE encodes parent-variable dependencies and aggregates all node states to represent decomposable network scores, supporting injective encoding of the overall score function.
- Encoding Modifications: Bayesian-network encoding replaces predecessor hidden states with node-type features because Bayesian networks represent conditional dependencies rather than computational flow.The modification reflects the structure of Bayesian-network scoring.
- Score Decomposition: The Bayesian-network score decomposes into local terms for variables given their parents.The example decomposes the overall score into terms such as s(X3 | X1, X2) and s(X5 | X3, X4).
- Local Dependencies: Parent features are used instead of parent hidden states to d-separate a node from its grandparents, while the update function remains unchanged.This modification targets the local conditional-dependence structure used by Bayesian-network scoring.
- Graph-Level Encoding: The encoder sums all node states as the final output rather than using only the ending node state, matching score decomposability.The same decomposition motivates corresponding decoding modifications.
- Injective Encoding: Combining the modified aggregation and update functions can injectively encode complete Bayesian-network conditional dependencies and the overall score function.The argument relies on injective modeling of conditional dependence and injective set-function summation.
- Latent-Space Implications: Injective computation encoding is necessary for performance-smooth latent representations, but the paper provides no theoretical guarantee that smoothness must hold.The authors report empirical predictive and Bayesian-optimization performance instead.
H More Details about Neural Architecture Search
The neural architecture search experiment trains D-VAE on ENAS-generated architectures, optimizes weight-sharing accuracy in latent space, and then fully evaluates selected architectures on CIFAR-10. D-VAE is positioned as an unsupervised alternative that can be reused across NAS tasks, although the BO target is only a proxy for fully trained accuracy.
- Experimental setup: ENAS generates seed architectures, whose weight-sharing accuracies are used to train D-VAE and search for new architectures in its latent space.The search uses architectures proposed by ENAS and a VAE trained on the resulting data.
- Experimental setup: 1000 ENAS epochs generate 20 architectures per epoch, followed by a 90%/10% training–test split for D-VAE.Architectures after the burn-in period are evaluated with shared weights before splitting the data.
- Optimization target: Weight-sharing accuracy is optimized because it takes around 0.5 second to evaluate, whereas fully training a network takes over 12 hours.The paper explicitly treats weight-sharing accuracy as the BO target rather than true validation or test accuracy.
- Final evaluation: After BO, selected architectures are fully trained for 310 epochs on CIFAR-10 to measure their true test accuracy.The final reported test evaluation follows the original ENAS training setting.
- Limitation: BO is not performed directly on fully trained validation accuracy because limited computational resources make that optimization impractical.The paper describes true-validation BO as a more principled procedure reserved for future exploration.
- Motivation: D-VAE is unsupervised and can be applied to NAS tasks targeting different datasets after training, unlike task-specific ENAS, which must be retrained for each task.The paper also identifies learned architecture embeddings as useful for visualization and classification.
- Results: 94.80% test accuracy is achieved by D-VAE’s best architecture on CIFAR-10, using 3 million parameters.The comparison discusses different search spaces and substantially greater resources used by NAONET, while emphasizing DAG generation and optimization as the paper’s main goal.
J Baselines
The baselines represent DAGs as strings or encode local neighborhoods through simultaneous message passing, whereas D-VAE is designed to preserve DAG computation dependencies. The reported comparisons emphasize brittleness of string representations and weak convergence of simultaneous-message-passing models.
- String-based baselines: S-VAE represents each DAG as a sequence of node strings containing node types and indicators for edges from previous nodes.Each node string combines a one-hot type encoding with an adjacency-column indicator vector.
- String-based baselines: GraphRNN decomposes adjacency columns into entries generated sequentially and is augmented with S-VAE’s encoder for latent-space comparison.The original GraphRNN is a pure generative model without an encoder.
- String-based baselines: Non-unique topological orderings can give the same DAG multiple string representations and different encoded representations, violating permutation invariance.This is identified as a drawback shared by S-VAE and GraphRNN.
- String-based baselines: Two bits of change in a string representation can completely change the computational purpose of the represented graph.The paper uses Figure 9 to illustrate the brittleness of string representations.
- D-VAE comparison: D-VAE differentiates subtle edge changes by propagating different predecessor hidden states through the affected node and its successors.Changing edge (2,3) to (1,3) changes the aggregated message received by node 3.
- Message-passing baselines: GCN uses simultaneous incoming-neighbor message passing and a summed graph state, while D-VAE follows directed computational paths.GCN is equipped with D-VAE’s decoder to form a VAE baseline.
- Message-passing baselines: DeepGMG is adapted from undirected graph generation using simultaneous message passing and a VAE encoder.Its decoder uses a sequential node-and-edge generation scheme similar to D-VAE’s.
- Results: DeepGMG’s training loss never approaches zero, whereas D-VAE can be trained to near-zero loss.The paper presents this as evidence of a limitation of simultaneous message passing for encoding DAGs.
K VAE Training Details
The experiments use matched model settings, a shared latent dimension and prior, and a reduced KL-divergence weight to improve VAE reconstruction. D-VAE’s bidirectional encoding is enabled for neural architectures but not where it gives no improvement.
- Shared settings: All four models use the same applicable settings, including single-layer GRUs with hidden state size 501 and latent dimension 56.The shared setup is intended to make comparisons as fair as possible.
- Shared settings: All VAE models use the N(0, I) prior and a diagonal-covariance normal posterior qφ(z|G).The posterior is defined for the input DAG G.
- D-VAE configuration: D-VAE’s bidirectional encoding is enabled for neural architectures and disabled for Bayesian networks and other models when it provides no better results.This configuration is task-dependent.
- Optimization: The VAE objective is ReconstructLoss + αKLDivergence with α = 0.005 instead of the original value 1.The smaller KL weight follows prior implementation practice after α = 1 produced poor reconstruction accuracy.
- Optimization: Mini-batch SGD with Adam trains the models, using batch size 32 for neural architectures and 300 epochs for most models.DeepGMG is the exception to the 300-epoch training schedule.
L SGP Training Details
The SGP models standardized performance targets before training, with task-specific data sampling for Bayesian networks. BO is evaluated against random latent-space sampling across neural-architecture and Bayesian-network experiments, using task-specific validity checks.
- SGP model: Sparse Gaussian process regression serves as the predictive model, using an open-source implementation.The SGP is used in the latent-space optimization procedure.
- SGP preprocessing: Training and testing performances are standardized using the training mean and standard deviation before SGP input.RMSE and Pearson’s r are also computed on standardized performances.
- SGP optimization: The SGP is trained for 100 epochs with Adam, batch size 1,000, and learning rate 5E-4.These settings are applied constantly in the reported SGP training procedure.
- Task-specific sampling: Bayesian-network SGP training randomly samples 5,000 examples each time instead of using all 180,000 examples.The paper cites realism for expensive evaluations and more stable BO performance as reasons.
- Validity evaluation: Prior validity is evaluated after rescaling latent samples to reduce effects from different posterior–prior KL alignment across models.This procedure is intended to focus on intrinsic DAG-generation validity rather than convergence differences.
- Validity checks: Neural architectures must satisfy input/output, path-connectivity, and adjacency validity checks before ENAS can read them.The checks include one starting node, one ending type, no isolated or blocked paths, and required sequential connectivity.
- Validity checks: Bayesian networks must contain exactly 8 nodes, each type in ASTLBEXD exactly once, and remain acyclic.These conditions are required for bnlearn evaluation on the Asia dataset.
- Validity checks: The training graphs generated by the original software already satisfy the stated validity constraints.This applies to the constraints used for the generated neural architectures and Bayesian networks.
M.3 More D-VAE experiments
Ablations compare asynchronous and simultaneous message passing, while FAST trades decoder complexity for speed and supports deeper NAS experiments. Visualizations show decoded architectures and Bayesian networks in smooth latent spaces.
- Ablation studies: D-VAE generally achieves the best performance among D-VAE, D-VAE (SMP), D-VAE (FAST), and S-VAE on 6-layer neural architectures.The comparison covers generative ability and latent space predictive ability.
- Ablation studies: Simultaneous message passing reduces reconstruction accuracy, whereas FAST retains similar generative ability to D-VAE and exceeds S-VAE in latent space predictive ability.FAST remains inferior to D-VAE in latent space predictive ability.
- Ablation studies: Using asynchronous message passing in both the encoder and decoder benefits D-VAE compared with applying it only to DAG encoding.The ablation supports asynchronous message passing throughout the model.
- 12-layer neural architectures: 3.88% error rate was obtained for the final 12-layer neural architecture found by D-VAE (FAST).FAST enables deeper architectures with much less training time because its decoding complexity is linear.
- Latent space visualization: The paper visualizes decoded neural architectures and Bayesian networks in two dimensions; both D-VAE and S-VAE show smooth latent spaces for Bayesian networks.Figure 13 compares D-VAE with S-VAE for neural architectures, and Figure 14 makes the same comparison for Bayesian networks.