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: The authors propose PACE, a Transformer-based architecture for directed acyclic graphs (DAGs). By providing a DAG-specific positional encoding (using canonicalisation and GINs), the claim is that a fully-connected self-attentional model can be executed over such an input, ameliorating the requirement of following the DAG's topological structure in an RNN-like fashion. The proposed model is evaluated over neural architecture search datasets, showing improvements over relevant baselines.