PACE: A Parallelizable Computation Encoder for Directed Acyclic Graphs

1 INTRODUCTION . Directed acyclic graphs ( DAGs ) are ubiquitous in various real-world problems including neural architecture search ( Elsken et al. , 2019 ; Wen et al. , 2020 ) , source code modeling ( Allamanis et al. , 2018 ) , structure learning of Bayesian networks ( Koller & Friedman , 2009 ; Zhang et al. , 2019 ) , etc . One key challenge in DAG optimization problems is that it is not easy to use gradient strategies to quickly adjust the structure of a DAG towards the right direction due to the absence of gradient information . Some earlier works propose to directly optimize the discrete DAG structure through black-box optimization techniques such as reinforcement learning ( Zoph & Le , 2016 ) , evolutionary algorithms ( Liu et al. , 2017 ) , and Bayesian optimization ( Kandasamy et al. , 2018 ) , which are inherently less efficient . A more recent approach is to encode DAGs into some continuous space for searching , and various DAG encoders are developed . In general , these DAG encoding schemes fall into two categories : structure-aware encoding scheme ( Ying et al. , 2019 ; Wen et al. , 2020 ; Shi et al. , 2020 ) and computation-aware ( performance-aware ) encoding scheme ( Zhang et al. , 2019 ; Thost & Chen , 2021 ; Ning et al. , 2020 ) . Due to the superior graph representation learning ability , graph neural networks ( GNNs ) have broadly achieved state-of-art performance on various graph learning tasks , such as node classification ( Velickovic et al. , 2018 ; Hamilton et al. , 2017 ) , graph classification ( Xu et al. , 2019 ; Zhang et al. , 2018 ; Duvenaud et al. , 2015 ) , and link prediction ( Zhang & Chen , 2018 ; Zhang et al. , 2020 ) . Basically , GNNs follow the message passing scheme ( Gilmer et al. , 2017 ) where each node aggregates node features from its one-hop neighborhood repeatedly to update its own feature , and this aggregation happens at all nodes simultaneously . However , Thost & Chen ( 2021 ) suggest that such a framework can not exploit the inductive bias of the computation dependency defined by DAGs , thus failing to generate a smooth latent ( DAG encoding ) space beneficial to downstream optimization routines . Hence , in order to model the dependency between nodes in DAGs , various GNNs specifically designed for encoding DAGs , such as D-VAE ( Zhang et al. , 2019 ) and DAGNN ( Thost & Chen , 2021 ) , are developed to inject the computation dependency between nodes into the representation learning process . Instead of updating node features simultaneously , these DAG encoders are constructed upon a gated recurrent unit ( GRU ) and will not update the representation of a node until all of its predecessors are updated . One way to achieve this is to perform message passing sequentially following a topological ordering of the nodes . Such an asynchronous message passing scheme actually simulates how a real computation is performed along the DAG—the message passing order respects the computation dependency defined by the DAG , thus better exploiting the inductive bias . However , one key limitation of such DAG encoders is that the encoding process is inherently sequential , precluding processing all nodes in parallel . Although DAGNN proposes a topological batching trick to accelerate the training speed by partitioning nodes into disjoint batches where nodes within a batch can be processed in parallel , the time complexity is still lower-bounded by the longest path in the DAG and the fundamental constraint of the sequential computation nature still remains . Numerous efforts have been made to reduce the sequential computation cost in sequence modeling literature ( Cho et al. , 2014 ; Wu et al. , 2016 ) . For instance , ConvS2S ( Gehring et al. , 2017 ) utilizes convolutional layers as building blocks to compute output representations at different positions , while Transformer ( Vaswani et al. , 2017 ) proposes to inject the position ( order ) information into the model through positional encoding , and then dependency between positions can be captured through the attention mechanism ( Bahdanau et al. , 2015 ; Gehring et al. , 2017 ) in parallel instead of resorting to recurrent neural networks ( RNNs ) . However , the success of these techniques relies on the inherent linear order of symbols in the input/output sequences that automatically characterizes the dependency between symbols . That is , the dependency between symbols is fully captured by their positions in the sequence . Such a condition is not satisfied by nodes in a DAG G = ( V , E ) since each node can have multiple predecessors instead of only one like in plain sequences , and the dependency between nodes forms a ( strong ) partial order rather than a linear order . Thus , previous parallel sequence modeling methods would fail in the representation learning of DAGs . In this paper , we propose a novel Parallelizable Attention-based Computation structure Encoder , PACE , to improve the computation efficiency over existing GRU-based DAG encoders . In order to borrow the power of Transformer for sequence modeling to DAG modeling problems , we need to design a positional encoding scheme specifically for DAGs which can fully capture the dependency between nodes in a DAG before applying the pairwise self-attention mechanism . To achieve this , we propose a GNN-based dag2seq framework which is proved to injectively map DAGs to sequences of node embeddings . This means , we are able to fully recover the DAG structure from these produced node embeddings , the same as the positional encoding in the original Transformer . After that , a Transformer encoder ( with mask operation ) is applied to the node embedding sequence to simultaneously learn representations for all nodes in the DAG through the self-attention mechanism . This way , PACE incorporates the relational inductive bias ( Battaglia et al. , 2018 ; Xu et al. , 2020 ) carried by DAG structures into the encoding process , while eschewing the recurrence in previous works thus greatly improving the parallelization and encoding efficiency . To demonstrate the superiority of the proposed PACE model , we evaluate PACE against current state-of-art DAG encoders as well as other general-purpose graph encoders for undirected graphs . Experimental results show that PACE not only outperforms competitive baselines by generating smooth latent spaces that capture the similarity between DAGs thus facilitating the downstream search subroutines , but also significantly boosts the training and inference speed through parallelization . 2 BACKGROUNDS . 2.1 PARALLELIZABLE SEQUENCE MODELS . Encoding the complexities and nuances of sequences plays a central role in various machine learning tasks , including sentiment classification ( Medhat et al. , 2014 ) , speech recognition ( Abdel-Hamid et al. , 2014 ) , and other natural language processing ( NLP ) tasks ( Khan et al. , 2016 ) . For many years , sequential models , such as recurrent neural networks ( RNNs ) ( Medsker & Jain , 2001 ) , were the primal way to solve the sequence encoding problem . These models are computationally expansive due to the sequential encoding process . Hence , many parallelizable sequence models are proposed , including Transformer ( Vaswani et al. , 2017 ) , BERT ( Devlin et al. , 2019 ) , etc . Our proposed PACE model is built upon the Transformer ( encoder ) architecture . Transformer ( Vaswani et al. , 2017 ) is arguably the earliest translation model that solves the sequenceto-sequence task without using sequence-aligned RNNs or convolutional architectures . Transformer relates information from different positions in the sequence through the ( masked ) self-attention mech- anism to encode/decode the representation of items in the input/output sequence , and incorporates the order of the items into the encoding/decoding process through a positional encoding mechanism . Briefly , the positional encoding mechanism is an injective function fpe : N → Rd that represents the positions ( i.e . indices ) of items in the sequence as d-dimensional vectors . Hence , it consistently outputs a unique encoding for each position in the sequence . The Transformer architecture is inherently parallelizable and could capture the long-term dependency with ease , thus is broadly applied to sequence modeling tasks in following works ( Dai et al. , 2019 ; Al-Rfou et al. , 2019 ; Devlin et al. , 2019 ; Lewis et al. , 2020 ) . Recently , Transformer has been widely used in image processing , and has achieved the state-of-art performance on various image learning tasks , including object detection ( Carion et al. , 2020 ) and image recognition ( Dosovitskiy et al. , 2020 ) . 2.2 DAG ENCODING PROBLEM . A directed acyclic graph ( DAG ) G is represented as a pair ( V , E ) with V = { v1 , v2 , .. , vn } denoting the set of nodes andE ∈ V ×V denoting the set of directed edges . A DAG often carries a computation . We use O to denote an ( computational ) operation dictionary . For instance , the operation dictionary O for NAS-Bench-101 dataset contains five operations : “ Input ” , “ Output ” , “ 3× 3 convolution ” , “ 1× 1 convolution ” , and “ 3× 3 max-pool ” . Let G = ( V , E ) be a DAG whose nodes represent operations in O . Then the DAG G represents a computation structure in which dependencies between operations are determined by directed edges in E. Hence , isomorphic DAGs define the same computation structure . Then the objective of the DAG encoding problem is to develop encoders that can generate embeddings to distinguish the computation structures defined by DAGs in the encoding space . DAGs G = ( V , E ) show a close relation with partial order . For any two nodes vi , vj ∈ V , let be a binary relation such that vi vj if and only if there is a directed path from vi to vj , then the binary relation defines a partial order on the node set V . Based on the partial order , sequential DAG encoders , such as D-VAE ( Zhang et al. , 2019 ) and DAGNN ( Thost & Chen , 2021 ) , use GRU ( Cho et al. , 2014 ) to recursively encode nodes in the input DAG , where a node is not encoded until all of its predecessors ( those nodes with a relation with it ) are encoded . Because of the possibly very long dependency chains , these sequential DAG encoders inherently share the same limitations as RNNs , such as the slow training and inference speed and the difficulty to capture long-term dependencies . To address these limitations , it is intuitive to generalize Transformer to the DAG encoding problem , since Transformer brings undoubtedly a huge improvement over the RNN-based sequence models . However , Figure 2 illustrates that it can be ambiguous how to represent DAGs as sequences due to the complex topological structure of DAGs . Let o : V → O be the function mapping each node in G to an operation in O , and fpe be the original positional encoding function in Transformer . One intuitive way to linearize a DAG into a sequence is to sort its nodes with a topological order , which is also used in GRU-based DAG encoders . However , the topological orders of nodes in a DAG are often not unique . For instance , graph G1 in Figure 2 has two different topological orders : v1 , v2 , v3 , v4 , v5 and v2 , v1 , v3 , v4 , v5 , hence resulting in two different node sequences Seq1 and Seq2 . Note that this does not hurt GRU-based DAG encoders as they are invariant to which topological order is used . Figure 2 : Illustration of the ambiguity when applying sequence models ( such as Transformer ) to the DAG encoding problem . To avoid the ambiguity of topological order , another approach is to sort nodes according to a canonical order ( i.e. , node index in the canonical form of the DAG ) as suggested by Niepert et al . ( 2016 ) . For example , let v1 , v2 , v3 , v4 , v5 be the nodes sorted by a canonical order . Then , we can linearize a DAG into a sequence ( o ( v1 ) , fpe ( v1 ) ) , ( o ( v2 ) , fpe ( v2 ) ) , ... ( o ( vn ) , fpe ( vn ) ) , similar to how Transformer represents a sentence . Although canonical order guarantees the generated sequence is unique for the same DAG , different DAGs might have the same sequence . For instance , DAGs G1 and G2 in Figure 2 are not isomorphic , yet they will be represented as the same sequence Seq2 . This is because the positional encoding func- tion fpe only captures node positions in the sequence but fails to encode the node dependencies , causing a significant structural information loss . In summary , it is not straightforward to generalize Transformer to DAGs . We need to design a smart linearization method for DAGs which guarantees that the same DAG is transformed to the same sequence , while different DAGs are always transformed to different sequences , in order to achieve a lossless transformation .

Summary: This paper proposes an encoder architecture for directed acyclic graphs which is parallelizable in computation and thus more efficient than asynchronous message passing alternatives. The model is based on representing the graph in its canonical form, i.e. a unique representation of its isomorphism class, which is transformed into a positional encoding and combined with a Transformer-based encoder. In experiments on neural architecture searches and Bayesian networks, the proposed architecture outperforms related methods.