summaries: This paper proposes a new architecture to strengthen message passing GNNs. Since the weaknesses of GNNs usually come from  the unidentifability of nodes, the authors suggest to first preprocess given graphs and add each node a new initial attribute, and then apply the outer GNN. The new node attributes come from another message passing GNN (called based GNN) applied to a subgraph of a certain depth (or height according to the paper's terminology) around each node. Then they prove that the new architecture can distinguish regular graphs with probability $1-o(1)$ (over the choice of a random regular graph), and so the new model is more expressive than MPNNs, since MPNNs can't. They finally evaluate their model in various settings.    For me the contributions of the paper are: (1) proposing nested GNNs (NGNNs); (2) showing that NGNNs can distinguish almost all regular graphs; (3) and evaluating the new model through various experiments.

summaries: This paper introduces Nested Graph Neural Networks (NGNNs), a novel kind of graph neural network that overcomes the limitations of typical message passing approaches by representing each node as a function of its nested subgraph, rather than its nested subtree.  The result is a model provably more expressive than 1-WL and 2-WL, at the cost of a constant increase in computational complexity (although this has practical limitations, as the authors point out).  NGNN works on two levels:   1. A "base GNN" is used to compute representations for each node, applying the GNN to the nested subgraph around the node and eventually pooling the subgraph to obtain the representation.   2. An "outer GNN" processes the resulting graph to obtain the final node representations, and a global pooling is optionally applied to get graph-level embeddings.   Therefore, NGNN can surpass 1-WL expressivity using only standard message-passing GNNs, making it very convenient and more computationally efficient than previous higher-order GNN variants.   The authors prove the expressive power of NGNN in distinguishing n-sized r-regular graphs.   The paper concludes with a thorough experimental analysis of NGNN on several benchmarks of graph classification and regression: QM9, ogb-molhiv, ogb-molpcba, and some TU datasets.  The method achieves state-of-the-art performance on ogb-molpcba and several QM9 subtasks.

summaries: This paper presents a framework that applies GNN to subgraphs around each node independently to enhance node representations, followed by a global pooling layer to get graph-level embedding. The framework can be applied to any GNNs. The paper also demonstrates it's relative improvement over several real-world datasets, although the improvement seems marginal. From theoretical side, the paper proved formally that this framework can distinguish k-regular graphs where most traditional message passing based neural network cannot.

summaries: This paper proposes Nested GNN, which encodes each node based on the subgraph rather than the subtree (as what GNN does) around it. The main difference between NGNN and GNN is that subgraph captures more information than the subtree, such as the interactions between a node's neighbors. The whole-graph representation can be pooled from each node representation, and can be applied in tasks such as graph classification. From theoretical perspective, this paper studies the limitations of expressive power of GNN and show that NGNN is more powerful. The proposed method shows good performance on multiple datasets.

summaries: Since GNNs based on rooted subtrees are known with limited expressive power, this paper proposes to represent nodes with their rooted subgraphs. Its expressive power with respect to graph isomorphism testing is proven beyond 2-WL. Such a method is compatible to plug into popular GNNs to enhance their representation power. Enhanced GNNs show strong performances on graph classification and regression tasks.

abstractText: Graph neural network (GNN)’s success in graph classification is closely related to the Weisfeiler-Lehman (1-WL) algorithm. By iteratively aggregating neighboring node features to a center node, both 1-WL and GNN obtain a node representation that encodes a rooted subtree around the center node. These rooted subtree representations are then pooled into a single representation to represent the whole graph. However, rooted subtrees are of limited expressiveness to represent a nontree graph. To address it, we propose Nested Graph Neural Networks (NGNNs). NGNN represents a graph with rooted subgraphs instead of rooted subtrees, so that two graphs sharing many identical subgraphs (rather than subtrees) tend to have similar representations. The key is to make each node representation encode a subgraph around it more than a subtree. To achieve this, NGNN extracts a local subgraph around each node and applies a base GNN to each subgraph to learn a subgraph representation. The whole-graph representation is then obtained by pooling these subgraph representations. We provide a rigorous theoretical analysis showing that NGNN is strictly more powerful than 1-WL. In particular, we proved that NGNN can discriminate almost all r-regular graphs, where 1-WL always fails. Moreover, unlike other more powerful GNNs, NGNN only introduces a constantfactor higher time complexity than standard GNNs. NGNN is a plug-and-play framework that can be combined with various base GNNs. We test NGNN with different base GNNs on several benchmark datasets. NGNN uniformly improves their performance and shows highly competitive performance on all datasets.

authors:
Bryan Perozzi
Rami Al-Rfou
Steven Skiena

title: Deepwalk: Online learning of social representations

venue: In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining,

title: Convolution kernels on discrete structures

venue: Technical report, Citeseer,

authors:
Nino Shervashidze
SVN Vishwanathan
Tobias Petri
Kurt Mehlhorn
Karsten M Borgwardt

title: Efficient graphlet kernels for large graph comparison

venue: In AISTATS,

authors:
Risi Kondor
Nino Shervashidze
Karsten M Borgwardt

title: The graphlet spectrum

venue: In Proceedings of the 26th Annual International Conference on Machine Learning,

authors:
Karsten M Borgwardt
Hans-Peter Kriegel

title: Shortest-path kernels on graphs

venue: IEEE International Conference on Data Mining, pages 8–pp. IEEE,

authors:
Marion Neumann
Roman Garnett
Christian Bauckhage
Kristian Kersting

title: Propagation kernels: efficient graph kernels from propagated information

venue: Machine Learning,

authors:
Nino Shervashidze
Pascal Schweitzer
Erik Jan van Leeuwen
Kurt Mehlhorn
Karsten M Borgwardt

title: Weisfeiler-lehman graph kernels

venue: Journal of Machine Learning Research,

authors:
Franco Scarselli
Marco Gori
Ah Chung Tsoi
Markus Hagenbuchner
Gabriele Monfardini

title: The graph neural network model

venue: IEEE Transactions on Neural Networks,

authors:
Joan Bruna
Wojciech Zaremba
Arthur Szlam
Yann LeCun

title: Spectral networks and locally connected networks on graphs

venue: arXiv preprint arXiv:1312.6203,

authors:
David K Duvenaud
Dougal Maclaurin
Jorge Iparraguirre
Rafael Bombarell
Timothy Hirzel
Alán Aspuru-Guzik
Ryan P Adams

title: Convolutional networks on graphs for learning molecular fingerprints

venue: In Advances in neural information processing systems,

authors:
Yujia Li
Daniel Tarlow
Marc Brockschmidt
Richard Zemel

title: Gated graph sequence neural networks

venue: arXiv preprint arXiv:1511.05493,

authors:
Thomas N Kipf
Max Welling

title: Semi-supervised classification with graph convolutional networks

venue: arXiv preprint arXiv:1609.02907,

authors:
Michaël Defferrard
Xavier Bresson
Pierre Vandergheynst

title: Convolutional neural networks on graphs with fast localized spectral filtering

venue: In Advances in Neural Information Processing Systems,

authors:
Hanjun Dai
Bo Dai
Le Song

title: Discriminative embeddings of latent variable models for structured data

venue: In Proceedings of The 33rd International Conference on Machine Learning,

authors:
Muhan Zhang
Zhicheng Cui
Marion Neumann
Yixin Chen

title: An end-to-end deep learning architecture for graph classification

venue: In AAAI,

authors:
Zhitao Ying
Jiaxuan You
Christopher Morris
Xiang Ren
Will Hamilton
Jure Leskovec

title: Hierarchical graph representation learning with differentiable pooling

authors:
Justin Gilmer
Samuel S Schoenholz
Patrick F Riley
Oriol Vinyals
George E Dahl

title: Neural message passing for quantum chemistry

venue: In Proceedings of the 34th International Conference on Machine Learning-Volume

authors:
Christopher Morris
Martin Ritzert
Matthias Fey
William L Hamilton
Jan Eric Lenssen
Gaurav Rattan
Martin Grohe

title: Weisfeiler and leman go neural: Higher-order graph neural networks

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Haggai Maron
Heli Ben-Hamu
Hadar Serviansky
Yaron Lipman

title: Provably powerful graph networks

authors:
Zhengdao Chen
Soledad Villar
Lei Chen
Joan Bruna

title: On the equivalence between graph isomorphism testing and function approximation with gnns

authors:
Boris Weisfeiler
AA Lehman

title: A reduction of a graph to a canonical form and an algebra arising during this reduction

venue: Nauchno-Technicheskaya Informatsia,

authors:
Pan Li
Yanbang Wang
Hongwei Wang
Jure Leskovec

title: Distance encoding–design provably more powerful gnns for structural representation learning

venue: arXiv preprint arXiv:2009.00142,

authors:
Christopher Morris
Gaurav Rattan
Petra Mutzel

title: Weisfeiler and leman go sparse: Towards scalable higher-order graph embeddings

authors:
Muhan Zhang
Yixin Chen

title: Weisfeiler-lehman neural machine for link prediction

venue: In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining,

authors:
Keyulu Xu
Weihua Hu
Jure Leskovec
Stefanie Jegelka

title: How powerful are graph neural networks

venue: arXiv preprint arXiv:1810.00826,

authors:
Andries E Brouwer
Willem H Haemers

title: Strongly regular graphs

venue: In Spectra of Graphs,

authors:
Ryan Murphy
Balasubramaniam Srinivasan
Vinayak Rao
Bruno Ribeiro

title: Relational pooling for graph representations

venue: In International Conference on Machine Learning,

authors:
Giorgos Bouritsas
Fabrizio Frasca
Stefanos Zafeiriou
Michael M Bronstein

title: Improving graph neural network expressivity via subgraph isomorphism counting

venue: arXiv preprint arXiv:2006.09252,

authors:
Jiaxuan You
Jonathan Gomes-Selman
Rex Ying
Jure Leskovec

title: Identity-aware graph neural networks

venue: arXiv preprint arXiv:2101.10320,

authors:
Dominique Beaini
Saro Passaro
Vincent Létourneau
William L Hamilton
Gabriele Corso
Pietro Liò

title: Directional graph networks

venue: arXiv preprint arXiv:2010.02863,

authors:
Haggai Maron
Heli Ben-Hamu
Nadav Shamir
Yaron Lipman

title: Invariant and equivariant graph networks

venue: arXiv preprint arXiv:1812.09902,

authors:
Haggai Maron
Ethan Fetaya
Nimrod Segol
Yaron Lipman

title: On the universality of invariant networks

venue: In International conference on machine learning,

authors:
Nicolas Keriven
Gabriel Peyré

title: Universal invariant and equivariant graph neural networks

venue: arXiv preprint arXiv:1905.04943,

authors:
Waïss Azizian
Marc Lelarge

title: Characterizing the expressive power of invariant and equivariant graph neural networks

venue: arXiv preprint arXiv:2006.15646,

authors:
George Dasoulas
Ludovic Dos Santos
Kevin Scaman
Aladin Virmaux

title: Coloring graph neural networks for node disambiguation

venue: arXiv preprint arXiv:1912.06058,

authors:
Zhengdao Chen
Lei Chen
Soledad Villar
Joan Bruna

title: Can graph neural networks count substructures? Advances in neural information processing

title: What graph neural networks cannot learn: depth vs width

venue: arXiv preprint arXiv:1907.03199,

authors:
Ryoma Sato
Makoto Yamada
Hisashi Kashima

title: Random features strengthen graph neural networks

venue: arXiv preprint arXiv:2002.03155,

authors:
Ralph Abboud
İsmail İlkan Ceylan
Martin Grohe
Thomas Lukasiewicz

title: The surprising power of graph neural networks with random node initialization

venue: arXiv preprint arXiv:2010.01179,

authors:
Clément Vignac
Andreas Loukas
Pascal Frossard

title: Building powerful and equivariant graph neural networks with structural message-passing

title: Link prediction based on graph neural networks

authors:
Muhan Zhang
Pan Li
Yinglong Xia
Kai Wang
Long Jin

title: Revisiting graph neural networks for link prediction

venue: arXiv preprint arXiv:2010.16103,

title: Inductive matrix completion based on graph neural networks

venue: In International Conference on Learning Representations,

authors:
Pim de Haan
Taco Cohen
Max Welling

title: Natural graph networks

venue: arXiv preprint arXiv:2007.08349,

authors:
Giannis Nikolentzos
George Dasoulas
Michalis Vazirgiannis

title: k-hop graph neural networks

venue: Neural Networks,

authors:
Sami Abu-El-Haija
Bryan Perozzi
Amol Kapoor
Nazanin Alipourfard
Kristina Lerman
Hrayr Harutyunyan
Greg Ver Steeg
Aram Galstyan

title: Mixhop: Higher-order graph convolutional architectures via sparsified neighborhood mixing

venue: In international conference on machine learning,

authors:
Kexin Huang
Marinka Zitnik

title: Graph meta learning via local subgraphs

venue: Advances in Neural Information Processing Systems,

authors:
Matthias Fey
Jan Eric Lenssen

title: Fast graph representation learning with pytorch geometric

venue: arXiv preprint arXiv:1903.02428,

authors:
Raghunathan Ramakrishnan
Pavlo O Dral
Matthias Rupp
O Anatole Von Lilienfeld

title: Quantum chemistry structures and properties of 134 kilo molecules

venue: Scientific data,

authors:
Zhenqin Wu
Bharath Ramsundar
Evan N Feinberg
Joseph Gomes
Caleb Geniesse
Aneesh S Pappu
Karl Leswing
Vijay Pande

title: Moleculenet: a benchmark for molecular machine learning

venue: Chemical science,

authors:
Kristian Kersting
Nils M. Kriege
Christopher Morris
Petra Mutzel
Marion Neumann

title: Benchmark data sets for graph kernels, 2016

authors:
Paul D Dobson
Andrew J Doig

title: Distinguishing enzyme structures from non-enzymes without alignments

venue: Journal of molecular biology,

authors:
Asim Kumar Debnath
de Compadre RL Lopez
Gargi Debnath
Alan J Shusterman
Corwin Hansch

title: Structure-activity relationship of mutagenic aromatic and heteroaromatic nitro compounds. correlation with molecular orbital energies and hydrophobicity

venue: Journal of medicinal chemistry,

authors:
Hannu Toivonen
Ashwin Srinivasan
Ross D King
Stefan Kramer
Christoph Helma

title: Statistical evaluation of the predictive toxicology challenge 2000–2001

authors:
Ida Schomburg
Antje Chang
Christian Ebeling
Marion Gremse
Christian Heldt
Gregor Huhn
Dietmar Schomburg

title: Brenda, the enzyme database: updates and major new developments

venue: Nucleic acids research,

authors:
Weihua Hu
Matthias Fey
Marinka Zitnik
Yuxiao Dong
Hongyu Ren
Bowen Liu
Michele Catasta
Jure Leskovec

title: Open graph benchmark: Datasets for machine learning on graphs

authors:
Brandon Anderson
Truong-Son Hy
Risi Kondor

title: Cormorant: Covariant molecular neural networks

venue: arXiv preprint arXiv:1906.04015,

authors:
Johannes Klicpera
Janek Groß
Stephan Günnemann

title: Directional message passing for molecular graphs

venue: arXiv preprint arXiv:2003.03123,

authors:
Zhuoran Qiao
Matthew Welborn
Animashree Anandkumar
Frederick R Manby
Thomas F Miller III

title: Orbnet: Deep learning for quantum chemistry using symmetry-adapted atomic-orbital features

venue: The Journal of Chemical Physics,

authors:
Hanqing Zeng
Muhan Zhang
Yinglong Xia
Ajitesh Srivastava
Andrey Malevich
Rajgopal Kannan
Viktor Prasanna
Long Jin
Ren Chen

title: Deep graph neural networks with shallow subgraph samplers

authors:
Will Hamilton
Zhitao Ying
Jure Leskovec

title: Inductive representation learning on large graphs

authors:
Federico Errica
Marco Podda
Davide Bacciu
Alessio Micheli

title: A fair comparison of graph neural networks for graph classification

authors:
Matthias Fey
Jan E. Lenssen

title: Fast graph representation learning with PyTorch Geometric

venue: In ICLR Workshop on Representation Learning on Graphs and Manifolds,

authors:
Guohao Li
Chenxin Xiong
Ali Thabet
Bernard Ghanem

title: Deepergcn: All you need to train deeper gcns

venue: arXiv preprint arXiv:2006.07739,

authors:
Gabriele Corso
Luca Cavalleri
Dominique Beaini
Pietro Liò
Petar

title: Veličković. Principal neighbourhood aggregation for graph nets

venue: arXiv preprint arXiv:2004.05718,

authors:
Rémy Brossard
Oriel Frigo
David Dehaene

title: Graph convolutions that can finally model local structure

venue: arXiv preprint arXiv:2011.15069,

authors:
Tuan Le
Marco Bertolini
Frank Noé
Djork-Arné Clevert

title: Parameterized hypercomplex graph neural networks for graph classification

venue: arXiv preprint arXiv:2103.16584,

authors:
Katsuhiko Ishiguro
Shin-ichi Maeda
Masanori Koyama

title: Graph warp module: an auxiliary module for boosting the power of graph neural networks

authors:
Douglas J Klein
Milan

title: Randić. Resistance distance

venue: Journal of Mathematical Chemistry,

heading: 1 Introduction

text:
Graph is an important tool to model relational data in the real world. Representation learning over graphs has become a popular topic of machine learning in recent years. While network embedding methods, such as DeepWalk [1], can learn node representations well, they fail to generalize to whole-graph representations, which are crucial for applications such as graph classification, molecule modeling, and drug discovery. On the contrary, although traditional graph kernels [2–7] can be used for graph classification, they define graph similarity often in a heuristic way, which is not parameterized and lacks some flexibility to deal with features.
In this context, graph neural networks (GNNs) have regained people’s attention and become the state-of-the-art graph representation learning tool [8–17]. GNNs use message passing to propagate features between connected nodes. By iteratively aggregating neighboring node features to the center node, GNNs learn node representations encoding their local structure and feature information. These node representations can be further pooled into a graph representation, enabling graph-level tasks such as graph classification. In this paper, we will use message passing GNNs to denote this class
∗Corresponding author: Muhan Zhang (muhan@pku.edu.cn). †Pan Li contributes Sec. 3.3 that proves the Theorem 1 and some implementation ideas.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
of GNNs based on repeated neighbor aggregation [18], in order to distinguish them from some high-order GNN variants [19–21] where the effective message passing happens between high-order node tuples instead of nodes.
GNNs’ message passing scheme mimics the 1-dimensional Weisfeiler-Lehman (1-WL) algorithm [22], which iteratively refines a node’s color according to its current color and the multiset of its neighbors’ colors. This procedure essentially encodes a rooted subtree around each node into its final color, where the rooted subtree is constructed by recursively expanding the neighbors of the root node. One critical reason for GNN’s success in graph classification is because two graphs sharing many identical or similar rooted subtrees are more likely classified into the same class, which actually aligns with the inductive bias that two graphs are similar if they have many common substructures [23].
Despite this, rooted subtrees are still limited in terms of expressing all possible substructures that can appear in a graph. It is likely that two graphs, despite sharing a lot of identical rooted subtrees, are not similar at all because their other substructure patterns are not similar. Take the two graphs G1 and G2 in Figure 1 as an example. If we apply 1-WL or a message passing GNN to them, the two graphs will always have the same representation no matter how many iterations/layers we use. This is because all nodes in the two graphs have identical rooted subtrees across all tree heights. However, the two graphs are quite different from a holistic perspective. G1 is composed of two triangles, while G2 is a hexagon. The intrinsic reason for such a failure is that rooted subtrees have limited expressiveness for representing general graphs, especially those with cycles.
To address this issue, we propose Nested Graph Neural Networks (NGNNs). The core idea is, instead of encoding a rooted subtree, we want the final representation of a node to encode a rooted subgraph (local h-hop subgraph) around it. The subgraph is not restricted to be of any particular graph type such as tree, but serves as a general description of the local neighborhood around a node. Rooted subgraphs offer much better representation power than rooted subtrees, e.g., we can easily discriminate the two graphs in Figure 1 by only comparing their height-1 rooted subgraphs.
To represent a graph with rooted subgraphs, NGNN uses two levels of GNNs: base (inner) GNNs and an outer GNN. By extracting a local rooted subgraph around each node, NGNN first applies a base GNN to each node’s subgraph independently. Then, a subgraph pooling layer is applied to each subgraph to aggregate the intermediate node representations into a subgraph representation. This subgraph representation is used as the final representation of the root node. Rather than encoding a rooted subtree, this final node representation encodes the local subgraph around it, which contains more information than a subtree. Finally, all the final node representations are further fed into an outer GNN to learn a representation for the entire graph. Figure 2 shows one NGNN implementation using message passing GNNs as the base GNNs and a simple graph pooling layer as the outer GNN.
One may wonder that the base GNN seems to still learn only rooted subtrees if it is message-passingbased. Then why is NGNN more powerful than GNN? One key reason lies in the subgraph pooling
layer. Take the height-1 rooted subgraphs (marked with red boxes) around v1 and v2 in Figure 1 as an example. Although v1 and v2’s height-1 rooted subtrees are still the same, their neighbors (labeled by 1 and 2) have different height-1 rooted subtrees. Thus, applying a one-layer message passing GNN plus a subgraph pooling as the base GNN is sufficient to discriminate G1 and G2.
The NGNN framework has multiple exclusive advantages. Firstly, it allows freely choosing the base GNN, and can enhance the base GNN’s representation power in a plug-and-play fashion. Theoretically, we proved that NGNN is more powerful than message passing GNNs and 1-WL by being able to discriminate almost all r-regular graphs (where 1-WL always fails). Secondly, by extracting rooted subgraphs, NGNN allows augmenting the initial features of a node with subgraphspecific structural features such as distance encoding [24] to improve the quality of the learned node representations. Thirdly, unlike other more powerful graph neural networks, especially those based on higher-order WL tests [19–21, 25], NGNN still has linear time and space complexity w.r.t. graph size like standard message passing GNNs, thus maintaining good scalability. We demonstrate the effectiveness of the NGNN framework in various synthetic/real-world graph classification/regression datasets. On synthetic datasets, NGNN demonstrates higher-than-1-WL expressive power, matching very well with our theorem. On real-world datasets, NGNN consistently enhances a wide range of base GNNs’ performance, achieving highly competitive results on all datasets.

heading: 2 Preliminaries

heading: 2.1 Notation and problem definition

text: We consider the graph classification/regression problem. Given a graph G = (V,E) where V = {1, 2, . . . n} is the node set and E ⊆ V × V is the edge set, we aim to learn a function mapping G to its class or target value y. The nodes and edges in G can have feature vectors associated with them, denoted by xi (for node i) and eij (for edge (i, j)), respectively.

heading: 2.2 Weisfeiler-Lehman test

text:
The Wesfeiler-Lehman (1-WL) test [22] is a popular algorithm for graph isomorphism checking. The classical 1-WL works as follows. At first, all nodes receive a color 1. Each node collects its neighbors’ colors into a multiset. Then, 1-WL will update each node’s color so that two nodes get the same new color if and only if their current colors are the same and they have identical multisets of neighbor colors. Repeat this process until the number of colors does not increase between two iterations. Then, 1-WL will return that two graphs are non-isomorphic if their node colors are different at some iteration, or fail to determine whether they are non-isomorphic. See [7, 26] for more details.
1-WL essentially encodes the rooted subtrees around each node at different heights into its color representations. Figure 1 middle shows the rooted subtrees around v1 and v2. Two nodes will have the same color at iteration h if and only if their height-h rooted subtrees are the same.

heading: 3 Nested Graph Neural Network

text: In this section, we introduce our Nested Graph Neural Network (NGNN) framework and theoretically demonstrate its higher representation power than message passing GNNs.

heading: 3.1 Limitations of the message passing GNNs

text:
Most existing GNNs follow the message passing framework [18]: given a graph G, each node’s hidden state ht+1v is updated based on its previous state h t v and the messages m t+1 v from its neighbors
ht+1v = Ut(h t v,m t+1 v ), where m t+1 v = ∑ u∈N(v|G) Mt(h t v,h t u, evu). (1)
Here Mt, Ut are the message and update functions at time stamp t, evu is the feature of edge (v, u), and N(v|G) is the set of v’s neighbors in graph G. The initial hidden states h0v are given by the raw node features xv . After T time stamps (iterations), the final node representations hTv are summarized into a whole-graph representation with a readout (pooling) function R (e.g., mean or sum):
hG = R({hTv |v ∈ G}). (2)
Such a message passing (or neighbor aggregation) scheme iteratively aggregates neighbor information into a center node’s hidden state, making it encode a local rooted subtree around the node. The final node representations will contain both the local structure and feature information around nodes, enabling node-level tasks such as node classification. After a pooling layer, these node representations can be further summarized into a graph representation, enabling graph-level tasks. When there is no edge feature and the node features are from a countable space, it is shown that message passing GNNs are at most as powerful as the 1-WL test for discriminating non-isomorphic graphs [27, 19].
For an h-layer message passing GNN, it will give two nodes the same final representation if they have identical height-h rooted subtrees (i.e., both the structures and the features on the corresponding nodes/edges are the same). If two graphs have a lot of identical (or similar) rooted subtrees, they will also have similar graph representations after pooling. This insight is crucial for the success of modern GNNs in graph classification, because it aligns with the inductive bias that two graphs are similar if they have many common substructures. Such insight has also been used in designing the WL subtree kernel [7], a state-of-the-art graph classification method before GNNs.
However, message passing GNNs have several limitations. Firstly, rooted subtree is only one specific substructure. It is not general enough to represent arbitrary subgraphs, especially those with cycles due to the natural restriction of tree structure. Secondly, using rooted subtree as the elementary substructure results in a discriminating power bounded by the 1-WL test. For example, all n-node r-regular graphs cannot be discriminated by message passing GNNs. Thirdly, standard message passing GNNs do not allow using root-node-specific structural features (such as the distance between a node and the root node) to improve the quality of the learned root node’s representation. We might need to break through such limitations in order to design more powerful GNNs.

heading: 3.2 The NGNN framework

text:
To address the above limitations, we propose the Nested Graph Neural Network (NGNN) framework. NGNN no longer aims to encode a rooted subtree around each node. Instead, in NGNN, each node’s final representation encodes the general local subgraph information around it more than a subtree, so that two graphs sharing a lot of identical or similar rooted subgraphs will have similar representations.
Definition 1. (Rooted subgraph) Given a graph G and a node v, the height-h rooted subgraph Ghv of v is the subgraph induced from G by the nodes within h hops of v (including h-hop nodes).
To make a node’s final representation encode a rooted subgraph, we need to compute a subgraph representation. To achieve this, we resort to an arbitrary GNN, which we call the base GNN of NGNN. For example, the base GNN can be simply a message passing GNN, which performs message passing within each rooted subgraph to learn an intermediate representation for every node of the subgraph, and then uses a pooling layer to summarize a subgraph representation from the intermediate node representations. This subgraph representation is used as the final representation of the root node in
the original graph. Take root node w as an example. We first perform T rounds of message passing within node w’s rooted subgraph Ghw. Let v be any node appearing in G h w. We have
ht+1 v,Ghw = Ut(h t v,Ghw ,mt+1 v,Ghw ), where mt+1 v,Ghw
= ∑
u∈N(v|Ghw)
Mt(h t v,Ghw ,htu,Ghw , evu). (3)
Here Mt, Ut are the message and update functions of the base GNN at time stamp t, N(v|Ghw) denotes the set of v’s neighbors within w’s rooted subgraph Ghw, and h
t+1 v,Ghw and mt+1 v,Ghw denote
node v’s hidden state and message specific to rooted subgraph Ghw at time stamp t + 1. Note that when node v attends different nodes’ rooted subgraphs, its hidden states and messages will also be different. This is in contrast to standard GNNs where a node’s hidden state and message at time t is the same regardless of which root node it contributes to. For example, ht+1v and m t+1 v in Eq. 1 do not depend on any particular rooted subgraph.
After T rounds of message passing, we apply a subgraph pooling layer to summarize a subgraph representation hGhw from the intermediate node representations {h T v,Ghw |v ∈ Ghw}.
hw : = hGhw = R0({h T v,Ghw |v ∈ Ghw}), (4)
where R0 is the subgraph pooling layer. This subgraph representation hGhw will be used as root node w’s final representation hw in the original graph. Note that the base GNNs are simultaneously applied to all nodes’ rooted subgraphs to return a final node representation for every node in the original graph, and all the base GNNs share the same parameters. With such node representations, NGNN uses an outer GNN to further process and aggregate them into a graph representation of the whole graph. For simplicity, we let the outer GNN be simply a graph pooling layer denoted by R1:
hG := R1({hw|w ∈ G}). (5)
The Nested GNN framework can be understood as a two-level GNN, or a GNN of GNNs—the inner subgraph-level GNNs (base GNNs) are used to learn node representations from their rooted subgraphs, while the outer graph-level GNN is used to return a whole-graph representation from the inner GNNs’ outputs. The inner GNNs all share the same parameters which are trained end-to-end with the outer GNN. Figure 2 depicts the implementation of the NGNN framework described above.
Compared to message passing GNNs, NGNN changes the “receptive field” of each node from a rooted subtree to a rooted subgraph, in order to capture better local substructure information. The rooted subgraph is read by a base GNN to learn a subgraph representation. Finally, the outer GNN reads the subgraph representations output by the base GNNs to return a graph representation.
Note that, when we apply the base GNN to a rooted subgraph, this rooted subgraph is extracted (copied) out of the original graph and treated as a completely independent graph from the other rooted subgraphs and the original graph. This allows the same node to have different representations within different rooted subgraphs. For example, in Figure 2, the same node B appears in four different rooted subgraphs. Sometimes it is the root node, while other times it is a 1-hop neighbor of the root node. NGNN enables learning different representations for the same node when it appears in different rooted subgraphs, in contrast to standard GNNs where a node only has one single representation at one time stamp (Eq. 1). Similarly, NGNN also enables using different initial features for the same node when it appears in different rooted subgraphs. This allows us to customize a node’s initial features based on its structural role within a rooted subgraph, as opposed to using the same initial features for a node across all rooted subgraphs. For example, we can optionally augment node B’s initial features with the distance between node B and the root—when node B is the root node, we give it an additional feature 0; and when B is a k-hop neighbor of the root, we give it an additional feature k. Such feature augmentation may help better capture a node’s structural role within a rooted subgraph. It is an exclusive advantage of NGNN and is not possible in standard GNNs.

heading: 3.3 The representation power of NGNN

text:
We theoretically characterize the additional expressive power of NGNN (using message passing GNNs as base GNNs) as opposed to standard message passing GNNs. We focus on the ability to discriminate regular graphs because they form an important category of graphs which standard GNNs cannot represent well. Using 1-WL or message passing GNNs, any two n-sized r-regular graphs will have the same representation, unless discriminative node features are available. In contrast, we prove that NGNN can distinguish almost all pairs of n-sized r-regular graphs regardless of node features.
Definition 2. If the message passing (Eq. 3) and the two-level graph pooling (Eqs. 4,5) are all injective given input from a countable space, then the NGNN is called proper.
A proper NGNN always exists due to the representation power of fully-connected neural networks used for message passing and Deep Set for graph pooling [28]. For all pairs of graphs that 1-WL can discriminate, there always exists a proper NGNN that can also discriminate them, because two graphs discriminated by 1-WL means they must have different multisets of rooted subtrees at some height h, while a rooted subtree is always included in a rooted subgraph with the same height.
Now we present our main theorem.
Theorem 1. Consider all pairs of n-sized r-regular graphs, where 3 ≤ r < (2 log n)1/2. For any small constant > 0, there exists a proper NGNN using at most d( 12 + ) logn log(r−1− )e-height rooted subgraphs and d lognlog(r−1− )e-layer message passing, which distinguishes almost all (1− o(1)) such pairs of graphs.
We include the proof in Appendix A. Theorem 1 has three implications. Firstly, since NGNN can discriminate almost all r-regular graphs where 1-WL always fails, it is strictly more powerful than 1-WL and message passing GNNs. Secondly, it implies that NGNN does not need to extract subgraphs with a too large height (about 12 logn log (r−1) ) to be more powerful. Moreover, NGNN is already powerful with very few layers, i.e., an arbitrarily small constant times lognlog (r−1) (as few as 1 layer). This benefit comes from the subgraph pooling (Eq. 4), freeing us from using deep base GNNs. We further conduct a simulation experiment in Appendix D to verify Theorem 1 by testing how well NGNN discriminates r-regular graphs in practice. The results match almost perfectly with our theory.
Although NGNN is strictly more powerful than 1-WL and 2-WL (1-WL and 2-WL have the same discriminating power [20]), it is unclear whether NGNN is more powerful than 3-WL. Our early-stage analysis shows both NGNN and 3-WL cannot discriminate strongly regular graphs with the same parameters [29]. We leave the exact comparison between NGNN and 3-WL to future work.

heading: 3.4 Discussion

text:
Base GNN. NGNN is a general plug-and-play framework to increase the power of a base GNN. For the base GNN, we are not restricted to message passing GNNs as described in Section 3.2. For example, we can also use GNNs approximating the power of higher-dimensional WL tests, such as 1-2-3-GNN [19] and PPGN/Ring-GNN [20, 21], as the base GNN. In fact, one limitation of these high-order GNNs is their O(n3) complexity. Using the NGNN framework we can greatly alleviate this by applying the higher-order GNN to multiple small rooted subgraphs instead of the whole graph. Suppose a rooted subgraph has at most c nodes, then by applying a high-order GNN to all n rooted subgraphs, we can reduce the time complexity from O(n3) to O(nc3). Complexity. We compare the time complexity of NGNN (using message passing GNNs as base GNNs) with a standard message passing GNN. Suppose the graph has n nodes with a maximum degree d, and the maximum number of nodes in a rooted subgraph is c. Each message passing iteration in a standard message passing GNN takes O(nd) operations. In NGNN, we need to perform message passing over all n nodes’ rooted subgraphs, which takes O(n · cd). We will keep c small (which can be achieved by using a small h) to improve NGNN’s scalability. Additionally, a small c enables the base GNN to focus on learning local subgraph patterns.
In Appendix B, we discuss some other design choices of NGNN.

heading: 4 Related work

text:
Understanding GNN’s representation power is a fundamental problem in GNN research. Xu et al. [27] and Morris et al. [19] first proved that the discriminating power of message passing GNNs is bounded by the 1-WL test, namely they cannot discriminate two non-isomorphic graphs that 1-WL fails to discriminate (such as r-regular graphs). Since then, there is increasing effort in enhancing GNN’s discriminating power beyond 1-WL [19, 21, 20, 30, 24, 31–33, 25]. Many GNNs have been proposed to mimic higher-dimensional WL tests, such as 1-2-3-GNN [19], Ring-GNN [21] and PPGN [20]. However, these models generally require learning the representations of all node tuples
of certain cardinality (e.g., node pairs, node triples and so on), thus cannot leverage the sparsity of graph structure and are difficult to scale to large graphs. Some works study the universality of GNNs for approximating any invariant or equivariant functions over graphs [34, 21, 35–37]. However, reaching universality would require polynomial(n)-order tensors, which hold more theoretical value than practical applicability. Dasoulas et al. [38] propose to augment nodes of identical attributes with different colors, which requires exhausting all the coloring choices to reach universality. Similarly, Relational Pooling (RP) [30] uses the ensemble of permutation-aware functions over graphs to reach universality, which requires exhausting all n! permutations to achieve its theoretical power. Its local version Local Relational Pooling (LRP) [39] applies RP over subgraphs around nodes, which is similar to our work yet still requires exhausting node permutations in local subgraphs and even more loses RP’s theoretical power. In contrast, NGNN maintains a controllable cost by only applying a message passing GNN to local subgraphs, and is guaranteed to be more powerful than 1-WL.
Because of the high cost of mimicking high-dimensional WL tests, several works have been proposed to increase GNN’s representation power within the message passing framework. Observing that different neighbors are indistinguishable during neighbor aggregation, some works propose to add one-hot node index features or random features to GNNs [40, 41]. These methods work well when nodes naturally have distinct identities irrespective of the graph structure. However, although making GNNs more discriminative, they also lose some of GNNs’ generalization ability by not being able to guarantee nodes with identical neighborhoods to have the same embedding; the resulting models are also no longer permutation invariant. Repeating random initialization helps with avoiding such an issue but gets much slower convergence [42]. An exception is structural message-passing (SMP) [43], which propagates one-hot node index features to learn a global n× d feature matrix for each node. The feature matrix is further pooled to learn a permutation-invariant node representation.
On the contrary, some works propose to use structural features to augment GNNs without hurting the generalization ability of GNNs. SEAL [44, 45], IGMC [46] and DE [24] use distance-based features, where a distance vector w.r.t. the target node set to predict is calculated for each node as its additional features. Our NGNN framework is naturally compatible with such distance-based features due to its independent rooted subgraph processing. GSN [31] uses the count of certain substructures to augment node/edge features, which also surpasses 1-WL theoretically. However, GSN needs a properly defined substructure set to incorporate domain-specific inductive biases, while NGNN aims to learn arbitrary substructures around nodes without the need to predefine a substructure set.
Concurrent to our work, You et al. [32] propose Identity-aware GNN (ID-GNN). ID-GNN uses different weight parameters between each root node and its context nodes during message passing. It also extracts a rooted subgraph around each node, and thus can be viewed as a special case of NGNN with: 1) the number of message passing layers equivalent to the subgraph height, 2) directly using the root node’s intermediate representation as its final representation without subgraph pooling, and 3) augmenting initial node features with 0/1 “identity”. However, the extra power of ID-GNN only comes from the “identity” feature, while the power of NGNN comes from the subgraph pooling— without using any node features, NGNN is still provably more discriminative than 1-WL. Another similar work to ours is natural graph network (NGN) [47]. NGN argues that graph convolution weights need not be shared among all nodes but only (locally) isomorphic nodes. If we view our distance-based node features as refining the graph convolution weights so that nodes within a center node’s neighborhood are no longer treated symmetrically, then our NGNN reduces to an NGN.
The idea of independently performing message passing within k-hop neighborhood is also explored in k-hop GNN [48] and MixHop [49]. However, MixHop directly concatenates the aggregation results of neighbors at different hops as the root representation, which ignores the connections between other nodes in the rooted subgraph. k-hop GNN sequentially performs message passing for k-hop, k − 1-hop, ..., and 0-hop node (the update of (i−1)-hop nodes depend on the updated states of i-hop nodes), while NGNN simultaneously performs message passing for all nodes in the subgraph thus is more parallelizable. Both MixHop and k-hop GNN directly use the root node’s representation as its final node representation. In contrast, NGNN uses a subgraph pooling to summarize all node representations within the subgraph as the final root representation, which distinguishes NGNN from other k-hop models. As Theorem 1 shows, the subgraph pooling enables using a much smaller number of message passing layers l (as small as 1) than the depth k of the subgraph, while MixHop and k-hop GNN always require l ≥ k. MixHop and k-hop GNN also do not have the strong theoretical power of NGNN to discriminate r-regular graphs. Like SEAL and k-hop GNN, G-Meta [50] is another work extracting subgraphs around nodes/links. It focuses specifically on a meta-learning setting.

heading: 5 Experiments

text:
In this section, we study the effectiveness of the NGNN framework for graph classification and regression tasks. In particular, we want to answer the following questions:
Q1 Can NGNN reach its theoretical power to discriminate 1-WL-indistinguishable graphs? Q2 How often and how much does NGNN improve the performance of a base GNN? Q3 How does NGNN perform in comparison to state-of-the-art GNN methods in open benchmarks? Q4 How much extra computation time does NGNN incur?
We implement the NGNN framework based on the PyTorch Geometric library [51]. Our code is available at https://github.com/muhanzhang/NestedGNN.

heading: 5.1 Datasets

text: To answer Q1, we use a simulation dataset of r-regular graphs and the EXP dataset [42] containing 600 pairs of 1-WL-indistinguishable but non-isomorphic graphs. To answer Q2, we use the QM9 dataset [52, 53] and the TU datasets [54]. QM9 contains 130K small molecules. The task here is to perform regression on twelve targets representing energetic, electronic, geometric, and thermodynamic properties, based on the graph structure and node/edge features. TU contains five graph classification datasets including D&D [55], MUTAG [56], PROTEINS [55], PTC_MR [57], and ENZYMES [58]. We used the datasets provided by PyTorch Geometric [51], where for QM9 we performed unit conversions to match the units used by [19]. The evaluation metric is Mean Absolute Error (MAE) for QM9 and Accuracy (%) for TU. To answer Q3, we use two Open Graph Benchmark (OGB) datasets [59], ogbg-molhiv and ogbg-molpcba. The ogbg-molhiv dataset contains 41K small molecules, the task of which is to classify whether a molecule inhibits HIV virus or not. ROC-AUC is used for evaluation. The ogbg-molpcba dataset contains 438K molecules with 128 classification tasks. The evaluation metric is Average Precision (AP) averaged over all the tasks. We include the statistics for QM9 and OGB datasets in Table 1.

heading: 5.2 Models

text:
QM9. We use 1-GNN, 1-2-GNN, 1-3-GNN, and 1-2-3-GNN from [19] as both the baselines and the base GNNs of NGNN. Among them, 1-GNN is a standard message passing GNN with 1-WL power. 1-2-GNN is a GNN mimicking 2-WL, where message passing happens among 2-tuples of nodes. 1-3-GNN and 1-2-3-GNN mimic 3-WL, where message passing happens among 3-tuples of nodes. 1-2-GNN and 1-3-GNN use features computed by 1-GNN as initial node features, and 1-2-3-GNN uses the concatenated features from 1-2-GNN and 1-3-GNN. We additionally include numbers provided by [53] and Deep LRP [39] as baselines. Note that we omit more recent methods [60–62] using advanced physical representations calculated from angles, atom coordinates, and quantum mechanics, which may obscure the comparison of models’ pure graph representation power. For NGNN, we uniformly use height-3 rooted subgraphs. For a fair comparison, the base GNNs in NGNN use exactly the same hyperparameters as when they are used alone, except for 1-GNN where we increase the number of message passing layers from 3 to 5 to make the number of layers larger than the subgraph height, similar to [63]. For subgraph pooling and graph pooling layers, we uniformly use mean pooling. All other settings follow [19].
TU. We use four widely adopted GNNs as the baselines and the base GNNs of NGNN: GCN [12], GraphSAGE [64], GIN [27], and GAT [15]. Since TU datasets suffer from inconsistent evaluation standards [65], we uniformly use the 10-fold cross validation framework provided by PyTorch Geomtric [66] for all the models to ensure a fair comparison. For GNNs, we search the number of message passing layers in {2, 3, 4, 5}. For NGNNs, we similarly search the subgraph height h in {2, 3, 4, 5}, so that both NGNNs and GNNs can have equal-depth local receptive fields. For NGNNs, we always use h+ 1 message passing layers instead of searching it together with h, because that will
make NGNNs have more hyperparameters to tune. All models have 32 hidden dimensions, and are trained for 100 epochs with a batch size of 128. For each fold, we record the test accuracy with the hyperparameters chosen based on the best validation performance of this fold. Finally, we report the average test accuracy across all the 10 folds.
OGB. We use GNNs achieving top places on the OGB graph classification leaderboard3 (at the time of submission) as the baselines, including GCN [12], GIN [27], DeeperGCN [67], Deep LRP [39], PNA [68], DGN [33], GINE [69], and PHC-GNN [70]. Note that those high-order GNNs [19–21, 25] are not included here, because despite being theoretically more discriminative, these GNNs are not among the GNNs with the best empirical performance on modern large-scale graph benchmarks, and their O(n3) complexity also raises a scalability issue. For NGNN, we use GIN as the base GNN (although GIN is not among the strongest baselines here). Some baselines additionally use the virtual node technique [18, 11, 71], which are marked by “*”. For NGNN, we search the subgraph height h in {3, 4, 5}, and the number of layers in {4, 5, 6}. We train the NGNN models for 100 and 150 epochs for ogbg-molhiv and ogbg-molpcba, respectively, and report the validation and test scores at the best validation epoch. We also find that our models are subject to high performance variance across epochs, likely due to the increased expressiveness. Thus, we save a model checkpoint every 10 epochs, and additionally report the ensemble performance by averaging the predictions from all checkpoints. The final hyperparameter choices and more details about the experimental settings are included in Appendix C. All results are averaged over 10 independent runs.
In the following, we uniformly use “Nested GNN” to denote an NGNN model using “GNN” as the base GNN. For example, Nested GIN denotes an NGNN model using GIN [27] as the base GNN. For the NGNN models in QM9, TU and OGB datasets, we augment the initial features of a node with Distance Encoding (DE) [24], which uses the (generalized) distance between a node and the root as its additional feature, due to DE’s successful applications in link-level tasks [44, 46]. Note that such feature augmentation is not applicable to the baseline models as discussed in Section 3.2. An ablation study on the effects of the DE features is included in Appendix E.

heading: 5.3 Results and discussion

text:
To answer Q1, we first run a simulation to test NGNN’s power for discriminating r-regular graphs. The results are presented in Appendix D. They match almost perfectly with Theorem 1, demonstrating that a practical NGNN can fulfil its theoretical power for discriminating r-regular graphs.
Table 2: Results (%) on EXP.
Method Test Accuracy GCN-RNI [42] 98.0±1.85 PPGN [20] 50.0±0.00 1-2-3-GNN [19] 50.0±0.00 3-GCN [42] 99.7±0.004 Nested GIN 99.9±0.26
We also test NGNN’s expressive power using the EXP dataset provided by [42], which contains 600 carefully constructed 1-WL indistinguishable but non-isomorphic graph pairs. Each pair of graphs have different labels, thus a standard message passing GNN cannot predict them both correctly, resulting in an expected classification accuracy of only 50%. We exactly follow the experimental settings and copy the baseline results in [42]. In Table 2, our Nested GIN model achieves a 99.9% classification accuracy, which outperforms
all the baselines and distinguishes almost all the 1-WL indistinguishable graph pairs. These results verified that NGNN’s expressive power is indeed beyond 1-WL and message passing GNNs.
To answer Q2, we adopt the QM9 and TU datasets. We show the QM9 results in Table 3. If the Nested version of a base GNN achieves a better result than the base GNN itself, we color that cell with light green. As we can see, NGNN brings performance gains to all base GNNs on most targets, sometimes by large margins. We also show the results on TU in Table 4. NGNNs also show improvement over their base GNNs in most cases. These results indicate that NGNN is a general framework for improving a GNN’s power. We further compute the maximum reduction of MAE for QM9 and maximum improvement of accuracy for TU before and after applying NGNN. NGNN reduces the MAE by up to 7.9 times for QM9, and increases the accuracy by up to 14.3% for TU. These results answer Q2, indicating that NGNN can bring steady and significant improvement to base GNNs.
To answer Q3, we compare Nested GIN with leading methods on the OGB leaderboard. The results are shown in Table 5. Nested GIN achieves highly competitive performance with these leading GNN models, albeit using a relatively weak base GNN (GIN). Compared to GIN alone, Nested GIN shows clear performance gains. It achieves test scores up to 79.86 and 30.07 on ogbg-molhiv and
3https://ogb.stanford.edu/docs/leader_graphprop/
Table 5: Results (%) on OGB datasets (* virtual node).
ogbg-molhiv (AUC) ogbg-molpcba (AP) Method Validation Test Validation Test CCN* 83.84±0.91 75.99±1.19 24.95±0.42 24.24±0.34 GIN* 84.79±0.68 77.07±1.49 27.98±0.25 27.03±0.23 Deep LRP 82.09±1.16 77.19±1.40 – – DeeperGCN* – – 29.20±0.25 27.81±0.38 HIMP – 78.80±0.82 – – PNA 85.19±0.99 79.05±1.32 – – DGN 84.70±0.47 79.70±0.97 – – GINE* – – 30.65±0.30 29.17±0.15 PHC-GNN 82.17±0.89 79.34±1.16 30.68±0.25 29.47±0.26 Nested GIN* 83.17±1.99 78.34±1.86 29.15±0.35 28.32±0.41 Nested GIN* (ens) 80.80±2.78 79.86±1.05 30.59±0.56 30.07±0.37
ogbg-molpcba, respectively, which outperform all the baselines. In particular, for the challenging ogbg-molpcba, our Nested GIN can achieve 30.07 and 28.32 test AP with and without ensemble, respectively, outperforming the plain GIN model (with 27.03 test AP) significantly. These results demonstrate the great empirical performance and potential of NGNN even compared to heavily tuned open leaderboard models, despite using only GIN as the base GNN.
To answer Q4, we report the training time per epoch for GIN and Nested GIN on OGB datasets. On ogbg-molhiv, GIN takes 54s per epoch, while Nested GIN takes 183s. On ogbg-molpcba, GIN takes 10min per epoch, while Nested GIN takes 20min. This verifies that NGNN has comparable time complexity with message passing GNNs. The extra complexity comes from independently learning better node representations from rooted subgraphs, which is a trade-off for the higher expressivity.
In summary, our experiments have firmly shown that NGNN is a theoretically sound method which brings consistent gains to its base GNNs in a plug-and-play way. Furthermore, NGNN still maintains a controllable time complexity compared to other more powerful GNNs.
Finally, we point out one memory limitation of the current NGNN implementation. Currently, NGNN does not scale to graph datasets with a large average node number (such as REDDIT-BINARY) or datasets with a large average node degree (such as ogbg-ppa) due to copying a rooted subgraph for each node to the GPU memory. Reducing batch size or subgraph height helps, but at the same time leads to performance degradation. One may wonder why materializing all the subgraphs into GPU memory is necessary. The reason is that we want to batch-process all the subgraphs simultaneously. Otherwise, we have to sequentially extract subgraphs on the fly, which results in a much higher latency. We leave the exploration of memory efficient NGNN to the future work.

heading: 6 Conclusions

text:
We have proposed Nested Graph Neural Network (NGNN), a general framework for improving GNN’s representation power. NGNN learns node representations encoding rooted subgraphs instead of rooted subtrees. Theoretically, we prove NGNN can discriminate almost all r-regular graphs where 1-WL always fails. Empirically, NGNN consistently improves the performance of various base GNNs across different datasets without incurring the O(n3) complexity like other more powerful GNNs.
Acknowledge
The authors greatly thank the actionable suggestions from the reviewers to improve the manuscript. Li is partly supported by the 2021 JP Morgan Faculty Award and the National Science Foundation (NSF) award HDR-2117997.

title: Nested Graph Neural Networks


summaries: The paper describes a framework for learning and inference in probabilistic models with structured latent spaces. It operates in the tradition of IWAE and VSMC-like methods that use importance sampling to tighten a variational bound (or equivalently, uses a variational bound as an objective for learning importance sampler proposals). Unlike methods that optimize a single global bound, NVI is defined in terms of local pairwise divergences between a sequence of target distributions, which acts to decouple the stochastic optimization, reducing gradient variance and potentially offering a strategy for parallel inference.

summaries: The author proposes a nested variational inference (NVI) framework that combines nested importance sampling and variational inference. This framework allows the training of the important sampling proposal by minimizing KL divergences at each level of nesting. This also provides us a way to learn intermediate densities, which serves as a heuristics for guiding the importance sampler. Empirically, the author evaluates NVI by (1) sampling from a multimodal distribution with a learned annealing path; (2) learn heuristics for future observations in the state-space model; (3) perform amortized inference in the hierarchical deep generative model.

summaries: This paper proposes an inference procedure that combines ideas from variational and nested inference, resulting in a sequence of forward and reverse proposal distributions that are optimized with local objectives that encourage the forward and backward edge marginals to be similar. This allows the objective and therefore gradients to decouple. The addition of resampling, afforded by the connection to nested inference, additionally improves sample diversity.

summaries: The authors present a method to jointly learn a series of proposals and intermediate densities for e.g. SMC, or annealed importance sampling. They use a novel objective which sums divergences between distributions over each adjacent pair of variables in the chain, and argue that this leads to lower-variance gradient estimates than prior work which uses a single objective for the entire sequence of variables. They show experimentally the benefits of this method in three different settings.

summaries: The authors introduce a new inference and approximate weighted sampling framework that combines f-divergence variational inference and nested importance sampling. An original variation sampling density is transformed through a series of weight preserving operators that map the proposal to an approximation of an unnormalized target distribution. On the other hand, reverse transformations are use to transform the target back to the sampling distribution. The transformations are trained by minimizing an f-divergence at each nesting level.

abstractText: We develop nested variational inference (NVI), a family of methods that learn proposals for nested importance samplers by minimizing an forward or reverse KL divergence at each level of nesting. NVI is applicable to many commonly-used importance sampling strategies and provides a mechanism for learning intermediate densities, which can serve as heuristics to guide the sampler. Our experiments apply NVI to (a) sample from a multimodal distribution using a learned annealing path (b) learn heuristics that approximate the likelihood of future observations in a hidden Markov model and (c) to perform amortized inference in hierarchical deep generative models. We observe that optimizing nested objectives leads to improved sample quality in terms of log average weight and effective sample size.

authors: Ambrogioni Luca Lin Kate Fertig Emily Vikram Sharad Hinne Max Moore Dave Gerven Marcel

title: Automatic structured variational inference

venue: International Conference on Artificial Intelligence and Statistics. PMLR

authors:
Arbel
Michael
Matthews
Alexander GDG
Doucet
Arnaud

title: Annealed Flow Transport Monte Carlo

venue: arXiv preprint

authors: Bornschein Jörg Bengio Yoshua

title: Reweighted Wake-Sleep

venue: International Conference on Learning Representations

authors: Burda Yuri Grosse Roger Salakhutdinov Ruslan

title: Importance Weighted Autoencoders

venue: International Conference on Representations

authors:
Del Moral
Pierre
Doucet
Arnaud
Jasra
Ajay

title: Sequential Monte Carlo Samplers

venue: Journal of the Royal Statistical Society: Series B (Statistical Methodology)

authors: Doucet Arnaud Freitas Nando Gordon Neil eds

title: Sequential Monte Carlo Methods in Practice

authors:
Hoffman
Matthew D

title: Learning deep latent Gaussian models with Markov chain Monte Carlo

venue: Proceedings of the 34th International Conference on Machine Learning-Volume 70. JMLR. org

authors:
Huang
Chin-Wei
Tan
Shawn
Lacoste
Alexandre
Courville
Aaron C

title: Improving explorability in variational inference with annealed variational objectives

authors:
Kingma
Diederik P
Welling
Max

title: Auto-encoding variational bayes

venue: arXiv preprint arXiv:1312.6114

authors:
Lawson
Dieterich
Tucker
George
Naesseth
Christian A
Maddison
Chris
Adams
Ryan P
Teh
Yee Whye

title: Twisted variational sequential monte carlo

venue: Third workshop on Bayesian Deep Learning (NeurIPS)

authors:
Le
Tuan Anh
Igl
Maximilian
Rainforth
Tom
Jin
Wood
Frank

title: Auto-Encoding Sequential Monte Carlo

authors:
Le
Tuan Anh
A Kosiorek
N Siddharth
Teh
Yee Whye
Wood
Frank

title: Revisiting reweighted wake-sleep for models with stochastic control flow

authors:
Li
Yingzhen
Turner
Richard E
Liu
Qiang

title: Approximate inference with amortised mcmc

venue: arXiv preprint arXiv:1702.08343

authors:
Loaiza-Ganem
Gabriel
Cunningham
John P

title: The continuous Bernoulli: fixing a pervasive error in variational autoencoders

authors:
Maddison
Chris J
Lawson
Dieterich
Tucker
George
Heess
Nicolas
Norouzi
Mohammad
Mnih
Andriy
Doucet
Arnaud
Teh
Yee Whye

title: Filtering variational objectives

venue: Proceedings of the 31st International Conference on Neural Information Processing Systems

authors: Naesseth Christian Linderman Scott Ranganath Rajesh Blei David

title: Variational sequential monte carlo

authors: Naesseth Christian Lindsten Fredrik Schon Thomas

title: Nested Sequential Monte Carlo Methods

venue: International Conference on Machine Learning

authors:
Naesseth
Christian A
Linderman
Scott W
Ranganath
Rajesh
Blei
David M

title: Variational sequential monte carlo

venue: arXiv preprint arXiv:1705.11140

authors:
Naesseth
Christian A
Lindsten
Fredrik
Schön
Thomas B

title: Elements of Sequential Monte Carlo

authors:
Neal
Radford M

title: Annealed importance sampling

venue: Statistics and computing

authors:
Rainforth
Tom
Kosiorek
Adam
Le
Tuan Anh
Maddison
Chris
Igl
Maximilian
Wood
Frank
Teh
Yee Whye

title: Tighter Variational Bounds Are Not Necessarily Better

venue: en. International Conference on Machine Learning

authors: Ranganath Rajesh Gerrish Sean Blei David

title: Black Box Variational Inference

venue: en. Artificial Intelligence and Statistics. Apr

authors: Rezende Danilo Mohamed Shakir

title: Variational Inference with Normalizing Flows

authors:
Rezende
Danilo Jimenez
Mohamed
Shakir
Wierstra
Daan

title: Stochastic Backpropagation and Approximate Inference in Deep Generative Models

venue: Proceedings of the 31st International Conference on Machine Learning

authors:
Ritchie
Daniel
Horsfall
Paul
Goodman
Noah D

title: Deep Amortized Inference for Probabilistic Programs

venue: en. arXiv:1610.05735 [cs, stat] (Oct

authors:
Ritchie
Daniel
Thomas
Anna
Hanrahan
Pat
Goodman
Noah D

title: Neurally-guided procedural models: learning to guide procedural models with deep neural networks

venue: arXiv preprint arXiv:

authors:
Roeder
Geoffrey
Wu
Yuhuai
Duvenaud
David K

title: Sticking the Landing: Simple, Lower-Variance Gradient Estimators for Variational Inference

authors: Salimans Tim Kingma Diederik Welling Max

title: Markov chain monte carlo and variational inference: Bridging the gap

authors: Schulman John Heess Nicolas Weber Theophane Abbeel Pieter

title: Gradient Estimation Using Stochastic Computation Graphs

venue: Advances in Neural Information Processing Systems 28

authors:
Tucker
George
Lawson
Dieterich
Gu
Shixiang
Maddison
Chris J

title: Doubly reparameterized gradient estimators for monte carlo objectives

venue: arXiv preprint arXiv:1810.04152

authors:
Wang
Tongzhou
Wu
Yi
Moore
Dave
Russell
Stuart J

title: Meta-learning MCMC proposals

authors: Wingate David Weber Theo

title: Automated Variational Inference in Probabilistic Programming

venue: arXiv preprint arXiv:1301.1299

authors:
Wu
Hao
Zimmermann
Heiko
Sennesh
Eli
Le
Tuan Anh
Meent
Jan-Willem van de

title: Amortized Population Gibbs Samplers with Neural Sufficient Statistics

heading: 1 Introduction

text:
Deep generative models provide a mechanism for incorporating priors into methods for unsupervised representation learning. This is particularly useful in settings where the prior defines an inductive bias that reflects the structure of the underlying domain. Training models with structured priors, however, poses some challenges. A standard strategy for training deep generative models is to maximize a reparameterized evidence lower bound with respect to both the generative and an inference model [Kingma, Welling, 2013; Rezende et al., 2014]. This approach works well in variational autoencoders with isotropic Gaussian priors, but often fails for models with more structured priors or likelihoods.
In recent years, a range of strategies for improving upon standard reparameterized variational inference have been put forward. These include wake-sleep style variational methods that minimize the forward KL-divergence [Bornschein, Bengio, 2015; Le et al., 2019] as well as sampling schemes that incorporate annealing [Huang et al., 2018], Sequential Monte Carlo [Le et al., 2018; Naesseth et al., 2017; Maddison et al., 2017], Gibbs sampling [Wu et al., 2019; Wang et al., 2018], and MCMC updates [Salimans et al., 2015; Hoffman, 2017; Li et al., 2017]. While these methods offer flexible inference, typically resulting in better approximations to the posterior compared to traditional variational inference methods, they are either model-specific, requiring specialized sampling schemes and gradient estimators, or can not be easily composed with other techniques.
In this paper, we propose nested variational inference, a framework for combining nested importance sampling and variational inference. Nested importance sampling formalizes the construction of proposals by way of calls to other importance samplers [Naesseth et al., 2015; Naesseth et al., 2019], Many existing importance samplers are instances of nested samplers, including methods based on annealed importance sampling [Neal, 2001] and sequential Monte Carlo [Del Moral et al., 2006].
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
NVI learns proposals by optimizing a divergence at each level of nesting. Additionally, we combine nested variational objectives with importance resampling to further improve sampling quality, without the need to undergo extra steps to maintain differentiability due to the local nature of the objective. Resampling also allows to compute gradient estimates based on incremental weights, which depend only on variables that are sampled locally, rather than on all variables in the model. Doing so yields lower variance weights, whilst maintaining a high sample diversity relative to existing methods.

heading: 2 Background

heading: 2.1 Stochastic Variational Inference.

text:
Stochastic variational methods approximate a target density π(z; θ) = γ(z; θ)/Z with parameters θ using a variational density q(z;φ) with parameters φ. Two common variational objectives are the forward and reverse Kullback-Leibler (KL) divergence, which are both instances of f -divergences
Df ( π ∥∥ q) = E
q
[ f ( π(z; θ)
q(z;φ)
)] , (1)
with f(w) = w logw and f(w) = − logw respectively. We are typically interested in the setting where π(z; θ) is the posterior pθ(z |x) of a model with latent variables z and observations x. In this case, γ(z; θ) = pθ(x, z) is the joint density of the model, and Z = pθ(x) is the marginal likelihood.
Reverse KL-divergence. When optimizing KL (q ‖π), known as the reverse or exclusive KL divergence, it is common practice to maximize a lower boundL = Eq[log(γ/q)] = logZ−KL (q ‖π) with respect to θ and φ. The gradient of L can be approximated using reparameterized samples [Kingma, Welling, 2013; Rezende et al., 2014], likelihood-ratio estimators [Wingate, Weber, 2013; Ranganath et al., 2014], or a combination of the two [Schulman et al., 2015; Ritchie et al., 2016a].
Forward KL-divergence. In the case of the forward divergence KL (π ‖ q), also known as the inclusive KL divergence, stochastic variational methods typically optimize separate objectives for the inference and generative model. A common strategy is to train the generative model by either maximizing the ELBO or likelihood on the data. Optimizing the inference model requires samples from π, which itself requires approximate inference. A common strategy, which was popularized in the context of reweighted wake-sleep (RWS) methods [Bornschein, Bengio, 2015; Le et al., 2019], is to use q as a proposal in an importance sampler.

heading: 2.2 Importance Sampling

text:
Self-Normalized Importance Samplers. An expectation Eπ[g(z)] with respect to π can be rewritten with respect to a proposal density q by introducing an unnormalized importance weight w,
E π
[ g(z) ] = 1
Z E q
[ w g(z) ] , w = γ(z; θ)
q(z;φ) . (2)
Self-normalized estimators use weighted samples (w(s), z(s))Ss=1 to both approximate the expectation with respect to q, and to compute an estimate Ẑ of the normalizing constant,
E π [g(z)] ' ĝ = 1 Ẑ 1 S S∑ s=1 w(s) g ( z(s) ) , Ẑ = 1 S S∑ s=1 w(s), z(s) ∼ q(·;φ). (3)
The resulting estimator is consistent, i.e. ĝ a.s.−→ Eπ[g(z)] as the number of samples S increases , but it is not unbiased, since it follows from Jensen’s inequality that 1/Z = 1/Eq[Ẑ] ≤ Eq[1/Ẑ] . If the discrepancy between q and π is large, the importance weights will have large variance, which also has an impact on the bias. When q = π, the importance weight w = Z has zero variance, and the Jensen’s inequality is tight. In the context of stochastic variational inference, this means that gradient estimates might initially be strongly biased, since there will typically be a large discrepancy between q and π. However, the variance will typically decrease as the quality of the approximation improves.
Nested importance sampling and proper weighting Nested Importance sampling formalizes the construction of proposals by way of calls to other importance samplers. Extending the example above, we now assume q is some unnormalized density itself that we cannot directly generates samples from. We can employ another importance sampler with proposal η to simulate samples from q,
z(l) ∼ η(·, ψ), w(l) = q(z (l);φ)
η(z(l);ψ) , Ẑq =
1
L L∑ l=1 w(l)
The resulting weighted samples (w(l), z(l))Ll=1 can subsequently be used compute the normalizing constant Zq of q and to generate samples that are approximately distributed w.r.t. q via resampling. This allows us to use q as a proposal in a subsequent importance sampling step, whose weighted samples can again be used to compute consistent estimates w.r.t. the target density of interest π. This can be generalized, by introducing the notion of proper weighting. Definition 2.1 (Proper weighting). Let π be a probability density. For some constant c > 0, a random pair (w, z) ∼ Π is properly weighted (p.w.) for an unnormalized probability density γ ≡ Zπ if w ≥ 0 and for all measurable functions g it holds that
E w,z∼Π [w g(z)] = c
∫ dz γ(z) g(z) = cZ E
z∼π [g(z)] .
As long as we can generate properly weighted samples for the target density π from some sampler Π, we can compute consistent estimates
1 S ∑S s=1 w
(s) g(z(s)) 1 S ∑S s=1 w (s) = cZ Ez∼π [g(z)] cZ = E z∼π [g(z)] . (4)

heading: 3 Nested Variational Inference

text:
A widely used strategy in importance sampling is to decompose a difficult sampling problem into a series of easier problems. A common approach is to define a sequence of unnormalized densities {γk}Kk=1 that interpolate between an initial density π1 = γ1/Z1, for which sampling is easy, and the final target density πK = γK/ZK . At each step, samples from the preceding density serve to construct proposals for the next density, which is typically combined with importance resampling or application of a Markov chain Monte Carlo (MCMC) operator to improve the average sample quality.
Nested variational objectives. NVI defines objectives for optimizing importance samplers that target a sequence of densities. At every step, it minimizes the discrepancy between a forward density π̂k = γ̂k/Zk−1, which acts as the proposal, and a reverse density π̌k = γ̌k/Zk, which defines an intermediate target. We define the forward density by combining the preceding target γk−1 with a forward kernel qk, and the reverse density by combining the next target γk with a reverse kernel rk−1,
γ̌k(zk, zk−1) = γk(zk) rk−1(zk−1|zk), γ̂k(zk, zk−1) = qk(zk|zk−1) γk−1(zk−1). (5)
Our goal is to learn pairs of densities π̌k and π̂k that are as similar as possible. To do so, we minimize a variational objective D that comprises an f -divergence for each step in the sequence, along with a divergence between the first intermediate target π1 and an initial proposal q1. Concretely, the NVI objective can be formulated as,
DNVI = Df ( π1 ∣∣∣∣ q1)+ K∑
k=2
Df ( π̌k ∣∣∣∣ π̂k). (6)
Since each intermediate density πk occurs in both π̌k and in π̂k+1, this defines a trade-off between maximizing the similarity to πk+1 and the similarity to πk−1. To optimize this objective, we need to be able to simulate samples from the intermediate densities.
Sampling from intermediate densities. Given a pair (wk−1, zk−1) that is properly weighted for the previous target density γk−1, we can use a sequential Monte Carlo sampling [Del Moral et al., 2006] construction to define a pair (wk, zk) that is properly weighted for γk,
zk ∼ qk(· | zk−1), wk = vk wk−1, vk = γ̌k(zk, zk−1)
γ̂k(zk, zk−1) . (7)
We refer to the ratio vk as the incremental weight. In this construction, (wk, zk−1:k) is properly weighted for γ̌k which implies that (wk, zk) is also properly weighted for γk, since∫
dzk−1 γ̌k(zk, zk−1) = ∫ dzk−1 γk(zk)rk−1(zk−1 | zk) = γk(zk). (8)
Sequential importance sampling can be combined with other operations that preserve proper weighting, including rejection sampling, application of an MCMC transition operator, and importance resampling. This defines a class of samplers that admits many popular methods as special cases, including sequential Monte Carlo (SMC) [Doucet et al., 2001], annealed importance sampling (AIS) [Neal, 2001], and SMC samplers [Del Moral et al., 2006]. In this work we only consider variational transition kernels and importance resampling. These samplers vary in the sequences of densities they define. For example, AIS and SMC samplers both define intermediate density over a fixed common domain. Here, a common strategy is to define an annealing path γk(zk) = γ1(zk)1−βkγK(zk)βk for 0 = β1 < β2 < . . . < βK = 1. In contrast, when using SMC for state-space models we can define the intermediate target πk to be the the filtering distribution on the first k states. In this setting, the dimensionality of the support increases at each step and we can omit the reverse kernel rk−1. We are exploring both of these cases in Section 4.

heading: 3.1 Computing Gradient Estimates

text:
The NVI objective can be optimized with respect to three sets of densities. We will use θk, φ̂k, and φ̌k to denote the parameters of the densities πk, qk, and rk respectively. For notational convenience, we use ρ̌k = {θk, φ̌k−1} to refer to the parameters of the reverse density π̌k, and ρ̂k = {φ̂k, θk−1} to refer to the parameters of the forward density π̂k,
Gradients of the Forward KL divergence. When we employ the forward KL as the objective, the derivative with respect to ρ̂k can be expressed as (see Appendix E.3),
− ∂ ∂ρ̂k
KL ( π̌k ∥∥ π̂k) = Ě
πk
[ ∂
∂ρ̂k log γ̂k
( zk, zk−1; ρ̂k )] − E πk−1 [ ∂ ∂ρ̂k log γk−1 ( zk−1; θk−1 )] . (9)
This case is the nested analogue of RWS-style variational inference. We can move the derivative into the expectation, since π̌k does not depend on ρ̂k. We then decompose log π̂k = log γ̂k − logZk−1 and use the identity from Equation 21 to express the gradient logZk−1 as an expectation with respect to πk−1. The resulting expectations can be approximated using self-normalized estimators based on the outgoing weights wk and incoming weights wk−1 respectively.
The gradient of the forward KL with respect to ρ̌k is more difficult to approximate, since the expectation is computed with respect to π̌k, which depends on the parameters ρ̌k. The gradient of this expectation has the form (see Appendix E.3)
− ∂ ∂ρ̌k
KL ( π̌k ∥∥ π̂k) = − Ě
πk
[ log vk ∂
∂ρ̌k log π̌k
( zk, zk−1; ρ̌k )] (10)
= − Ě πk
[ log vk ∂
∂ρ̌k log γ̌k
( zk, zk−1; ρ̌k )] + Ě πk [ log vk ] E πk [ ∂ ∂ρ̌k log γk ( zk; θk )] .
In principle, we can approximate this gradient using self-normalized estimators based on the outgoing weight wk. We experimented extensively with this estimator, but unfortunately we found it to be unstable, particularly for the gradient with respect to the parameters of the reverse kernel rk−1. For this reason, our experiments employ the reverse KL when learning reverse kernels.
Our hypothesis is that the instability in this estimator arises because the gradient decreases the probability of high-weight samples and increases the probability of low-weight samples, rather than the other way around. This could lead to problems during early stages of training, when the estimator will underrepresent low-weight samples, for which the probability should increase.
Gradients of the Reverse KL divergence. When computing the gradient of the reverse KL with respect to ρ̂k, we obtain the nested analogue of methods that maximize a lower bound. Here we can either use reparameterized samples [Kingma, Welling, 2013; Rezende et al., 2014] or likelihood-ratio estimators [Wingate, Weber, 2013; Ranganath et al., 2014]. We will follow Ritchie et al. (2016b) and define a unified estimator in which proposals are generated using a construction
wk = vkwk−1, zk = g(z̃k, φ̂k), z̃k ∼ q̃k(z̃k | zk−1, φ̂k), wk−1, zk−1 ∼ Πk−1. (11) This construction recovers reparameterized samplers in the special case when q̃k does not depend on parameters, and recovers non-reparameterized samplers when zk = z̃k. This means it is applicable to models with continuous variables, discrete variables, or a combination of the two. The gradient of the reverse KL for proposals that are constructed in this manner becomes (see Appendix E.2)
− ∂ ∂ρ̂k
KL ( π̂k ∥∥ π̌k) = Ê
πk
[ ∂
∂zk log γ̌k
( zk, zk−1; ρ̂k )∂zk ∂ρ̂k ] + Ê πk [ log vk ∂ ∂ρ̂k log γ̂k ( zk, zk−1; ρ̂k )] (12)
− Ê πk
[ log vk ] E
πk−1
[ ∂
∂ρ̂k log γk−1
( zk−1; θk−1 )] .
In this gradient, the first term represents the pathwise derivative with respect to reparameterized samples. The second term defines a likelihood-ratio estimator in terms of the unnormalized density γ̂k, and the third term computes the contribution of the gradient of the log normalizer logZk−1.
Computing the gradient of the reverse KL with respect to ρ̌k is once again straightforward, since we are computing an expectation with respect to π̂k, which does not depend on ρ̌k. This means we can move the derivative into the expectation, which yields a gradient analogous to that in Equation 9,
− ∂ ∂ρ̌k
KL ( π̂k ∥∥ π̌k) = Ê
πk
[ ∂
∂ρ̌k log γ̌k
( zk, zk−1; ρ̌k )] − E πk [ ∂ ∂ρ̌k log γk ( zk; θk )] . (13)
Variance Reduction. To reduce the variance of the gradient estimates we use the expected logincremental weight as a baseline for score function terms and employ the sticking-the-landing trick [Roeder et al., 2017] when reparameterizing the forward kernel as described in Appendix E.

heading: 3.2 Relationship to Importance-Weighted and Self-Normalized Estimators

text:
There exists a large body of work on methods that combine variational inference with MCMC and importance sampling. We refer to Appendix A for a comprehensive discussion of related and indirectly related approaches. To position NVI in the context of the most directly related work, we here focus on commonly used importance-weighted and self-normalized estimators.
NVI differs from existing methods in that it defines an objective for pairs of variables (zk, zk−1) at each level of nesting, rather than a single objective for the entire sequence of variables (z1, . . . , zK). A standard approach for combining importance sampling and variational inference is to define an “importance-weighted” lower bound L̂K = log ẐK [Burda et al., 2016]. By Jensen’s inequality, E[L̂K ] ≤ logE[ẐK ] = logZK , which implies that we can define a stochastic lower bound using any properly-weighted sampler for γK , including samplers based on SMC [Le et al., 2018; Naesseth et al., 2018; Maddison et al., 2017]. For purposes of learning the target density γK , this approach is equivalent to computing an RWS-style estimator of the gradient,
∂
∂θK L̂K =
1
ẐK
1
S S∑ s=1 wsK ∂ ∂θK log γK(z s K ; θK). (14)
However, these two approaches are not equivalent for purposes of learning the proposals. We can maximize a stochastic lower bound to learn qk, but this requires doubly-reparameterized estimators
[Tucker et al., 2018] in order to avoid problems with the signal-to-noise ratio in this estimator, which can paradoxically deteriorate with the number of samples [Rainforth et al., 2018]. The estimators in NVI do not suffer from this problem, since we do not compute the logarithm of an average weight.
NVI is also not equivalent to learning proposals with RWS-style estimators. If we use sequential importance sampling (SIS) to generate samples, a self-normalized gradient for the parameters of qk that is analogous to the one in Equation 24 has the form
E πK ,rK−1,...,r1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k) ] ' 1 ẐK 1 S S∑ s=1 wsK ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (15)
Note that this expression depends on the final weight wK . By contrast, a NVI objective based on the forward KL yields a self-normalized estimator that is defined in terms of the intermediate weight wk
E πk,rk−1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k) ] ' 1 Ẑk 1 S S∑ s=1 wsk ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (16)
If instead of SIS we employ sequential importance resampling (i.e. SMC), then the incoming weight wk−1 is identical for all samples. This means that we can express this estimator in terms of the incremental weight vk rather than the intermediate weight wk
E πk,rk−1
[ ∂
∂φ̂k log qk(zk | zk−1 ; φ̂k)
] ' S∑ s=1 vsk∑S s′=1 v s′ k ∂ ∂φ̂k log qk(z s k | zsk−1 ; φ̂k). (17)
We see that NVI allows us to compute gradient estimates that are localized to a specific level of the sampler. In practice, this can lead to lower-variance gradient estimates.
Having localized gradient computations also offers potential memory advantages. Existing methods typically perform reverse-mode automatic differentiation on an objective that is computed from the final weights (e.g. the stochastic lower bound). This means that memory requirements scale as O(SK) since the system needs to keep the entire computation graph in memory. In NVI, gradient estimates at level k do not require differentiation of the incoming weights wk−1, This means that it is possible to perform automatic differentiation on a locally-defined objective before proceeding to the next level of nesting, which means that memory requirements would scale as O(S). It should therefore in principle be possible to employ a large number of levels of nesting K in NVI, although we do not evaluate the stability of NVI at large K in our experiments.

heading: 4 Experiments

text: We evaluate NVI on three tasks, (1) learning to sample form an unnormalized target density where intermediate densities are generated using annealing, (2) learning heuristic factors to approximate the marginal likelihood of future observations in state-space models, and finally (3) inferring distributions over classes from small numbers of examples in deep generative Bayesian mixtures.

heading: 4.1 Sampling from Multimodal Densities via Annealing

text:
A common strategy when sampling from densities with multiple isolated modes is to anneal from an initial density γ1, which is typically a unimodal distribution that we can sample from easily, to the
target density γK , which is multimodal [Neal, 2001]. Recent work on annealed variational objectives (AVOs) learns forward kernels qk and reverse kernels rk−1 for an annealing sequence by optimizing a variational lower bound at each level of nesting [Huang et al., 2018].
max qk,rk LAVOk , LAVOk = E q1,...,qk
[ log vk ] , γk(z) = π1(z) 1−βkγK(z) βk , k = 1, . . . ,K. (18)
NVI allows us to improve upon AVO in two ways. First, we can perform importance resampling at every step to optimize an SMC sampler rather than an annealed importance sampler. Second, we can learn the annealing schedule (β1, . . . , βK) as part of the intermediate densities γk such that intermediate densities are scheduled more equidistantly in terms of KL-divergence.
We illustrate the effect of these two modifications in Figure 3, in which we compare AVO to NVI with resampling and a learned path, which we refer to as NVIR∗. Both methods minimize the reverse KL at each step 1. The learned annealing path in NVIR∗ results in a smoother interpolation between the initial and final density. We also see that AVO does not assign equal mass to all 8 modes, whereas NVIR∗ yields a more even distribution.
In Figure 4a we compute the reverse KL between targets at each step in the sequence. For the standard linear annealing path, the KL decreases with k, suggesting that later intermediate densities are increasingly redundant. By constrast, in NVIR∗ we see that the KL is approximately constant across the path, which is what we would expect when minimizing D with respect to βk. This is also the case in an ablation without resampling, which we refer to as NVI∗.
1As noted in Section 3.1, we found optimization of the forward KL to be unstable when learning rk.
Figure 5: Ground truth densities (GT) and samples from final target density for NVIR∗ with 2 intermediate densities (K=4).
Figure 4b shows a rolling average of the ESS and its variance during training. We compare NVI-based methods to SVI and a variational SMC sampler [Le et al., 2018; Maddison et al., 2017; Naesseth et al., 2017]. NVIR* has consistently higher ESS and significantly lower variance compared to baselines. These plots also provides insight into the role of resampling in training dynamics. In NVI∗, we observe a cascading convergence, which is absent in NVIR∗. We hypothesize that resampling reduces the reliance on high-quality proposals from step k − 1 when estimating gradients at step k. Annealed NVI has similar use cases as normalizing flows [Rezende, Mohamed, 2015]. Inspired by concurrent work of Arbel et al. (2021), which explores a similar combination of SMC samplers and normalizing flows, we compare flow-based versions of NVI to planar normalizing flows, which maximize a standard lower bound (SVI-flow). We find that a normalizing flows can be effectively trained with NVI, in that samplers produce better estimates of the normalizing constant and higher ESS compared to SVI-flow. We also find that flow based models are able to produce high-quality samples with fewer intermediate densities (Figure 8). Moreover, we see that combining a flow-based proposal with learned βk values (NVI∗-flow) results in a more accurate approximation of the target than in an ablation with a linear interpolation path (AVO-flow)2.
In Table 1 we report sample quality in terms of the stochastic lower bound L̂K = log ẐK and the effective sample size ESS = ( ∑ s w s K) 2/ ∑ s(w s K)
2. The first metric can be interpreted as a measure of the average sample quality, whereas the second metric quantifies sample diversity. We compare NVI with and without resampling (NVIR∗ and its NVI∗) to ablations with a linear annealing path (NVIR and NVI), AVO, and a standard SVI baseline in which there are no intermediate densities. We additionally compare against AVO-flow and NVI∗-flow, which employ flows. We observe that NVIR∗ and NVI∗-flow outperform ablations and baselines in terms of log Ẑ, and are competitive in terms of ESS. We show qualitative results for two additional target densities in Figure 5. For more details we refer to Appendix G.1.

heading: 4.2 Learning Heuristic Factors for State-space Models

text:
Sequential Monte Carlo methods are commonly used in state-space models to generate samples by proposing one variable at a time. To do so, they define a sequence of densities πk = γk/Zk on the first k time points in a model, which are also known as filtering distributions,
γk(z1:k, η) = p(x1:k, z1:k, η) = p(η) p(x1, z1 | η) ∏k
l=2 p(xl, zl | zl−1, η), Zk = p(x1:k). (19)
Here z1:k and x1:k are sequences of hidden states and observations, and η is a set of global variables of the model. The densities γk differ from those in the annealing task in Section 4.1 in that the dimensionality of the support increases at each step. We define a forward density γ̂k that combines the preceding target γk−1 with a proposal qk for the time point zk, and define a reverse density γ̌k = γk that is equal to the next intermediate density (which means that we omit rk−1),
γ̌k(z1:k, η) = γk(z1:k, η), γ̂k(z1:k, η) = qk(zk |xk, z1:k−1, η) γk−1(z1:k−1, η). (20)
A limitation of this construction is that the filtering distribution πk−1 is not always a good proposal, since it does not incorporate knowledge of future observations. Ideally, we would like to define
2AVO-flow is itself a novel combination of AVO and flows, albeit an incremental one.
intermediate densities γk(z1:k, η) = p(x1:K , z1:k, η) that correspond to the smoothing distribution, but this requires computation of the marginal likelihood of future observations p(xk+1:K | zk, η), which is intractable. This is particularly problematic when sampling η as part of the SMC construction. The first density π1(z1, η) = p(z1, η |x1) will be similar to the prior, which will result in poor sample efficiency, since the smoothing distribution p(z1, η |x1:K) will typically be much more concentrated. To overcome this problem, we will use NVI to learn heuristic factors ψθ that approximate the marginal likelihood of future observations. We define a sequence of densities (γ0, . . . , γK),
γ0(η) = p(η) ψθ(x1:K |η), γk(z1:k, η) = p(x1:k, z1:k, η) ψθ(xk+1:K | η), k = 1, 2, ...,K. Our goal is to learn parameters θ of the heuristic factor such that that intermediate densities approximate the smoothing distribution. This approach is similar to recently proposed work on twisted variational SMC [Lawson et al., 2018], which maximized a stochastic lower bound.
To evaluate the this approach, we will learn heuristic factors for a hidden Markov model (HMM). While HMMs are a well-understood model class, they are a good test case, in that they give rise to significant sample degeneracy in SMC and allow to compute tractable heuristics. We optimize an NVI objective based on the forward KL with respect to the heuristic factor ψθ, an initial proposal q0(η |x1:K ;φ) and a forward kernel qk(zk |xk, zk−1, η;φk). Figure 6 shows qualitative and quantitative results. We compare NVIR∗ with a neural GMM-heuristic to a baseline without a heuristic and a baseline that uses a hand-coded GMM-heuristic. In addition, we train models which employ the convex update operator proposed in the Automatic Structured Variational Inference (ASVI) [Ambrogioni et al., 2021]. Training ASVI by optimizing the ELBO failed due to the high variance of the score function estimator, which is needed for the discrete latent variables. However, we were able to train ASVI-based models with NVIR∗. We refer to this novel combination as NASVI. Because NVI let us treat each step as a separate optimization problem, we also compare partial optimization with respect to γ̂k only to full optimization with respect to both γ̂k and γ̌k. While full optimization yields poor results, partial optimization learns a neural heuristic whose performance is similar to the GMM heuristic, which is a strong baseline in this context. For more details see Appendix G.2.

heading: 4.3 Meta Learning with Deep Generative models

text:
In this experiment, we evaluate NVI in the context of deep generative models with hierarchicallystructured priors. Concretely, we consider the task of inferring class weights from a mini-batch of images in a fully unsupervised manner. For this purpose, we employ a variational autoencoder with a prior in the form of a Bayesian Gaussian mixture model. Unlike the previous experiments, here we both train a generative and an inference model. The forward kernels in this setting act as the generative model and the reverse kernels are the inference model. Moreover, in this model sampling simplifies because we design the intermediate densities to be tractable, hence no importance sampling construction is needed. Regardless, we still can construct an NVI objective to train hierarchical deep generative models more efficiently and with a more accurate posterior inference.
Model Description. We define a hierarchical deep generative model for batches of N images of the form (see Appendix G.3 for a graphical model and architecture description)
λ ∼ Dir(· ; α) cn ∼ Cat(· |λ) zn ∼ N (· |µcn , 1/τcn) xn ∼ p(· | zn; θ) for n = 1 . . . N. Here λ, cn, zn, and xn refer to the cluster probabilities, cluster assignments, latent codes, and observed images respectively. In the likelihood p(x|z; θx), we use a convolutional network to parameterize a continuous Bernoulli distribution [Loaiza-Ganem, Cunningham, 2019]. We define
proposals q(z|x;φz), q(c|z;φc), and q(λ|c;φλ), which are also parameterized by neural networks. We refer to this model as Bayesian Gaussian mixture model VAE (BGMM-VAE).
Objective. To construct an NVI objective, we define intermediate densities for c and z. Unlike in previous experiments, we employ tractable densities in the form of a categorical πc(c; θc) for cluster assignments and a 8-layer planar flow πz(z; θz) for the latent codes. Subsequently, we can define the forward and reverse densities at each step as:
π̂2(λ, c) = π(λ)p(c|λ), π̂3(c, z) = π(c)p(z|c), π̂4(z, x) = π(z)p(x|z), π̌2(λ, c) = π(c)q(λ|c), π̌3(c, z) = π(z)q(c|z), π̌4(z, x) = π̃(x)q(z|x),
where π̃(x) is an empirical distribution over mini-batches of training data. We minimize the forward KL for the first two steps and the reverse KL at the final step.
Since the intermediate densities are tractable in this model, no nested importance sampling is required to optimize this nested objective; we can compute gradient estimates based on a (single) sample from p(λ)p(c|λ) in the first term, πc(c)q(λ|c) in the second, and p̂(x)q(z|x) in the final term. To learn the parameters {µ, τ, θx} of the generative model, we maximize a single-sample approximation of a lower bound L = Eq [ log ( p(x, z, c, λ) / q(z, c, λ|x) )] .
Results. We evaluate NVI for the BGMM-VAE using the following procedure. We generate minibatches with a sampled λ∗ (for which we make use of class labels that are not provided to the model). We then compute the expectation of λ under q(λ, c, z|x) by sampling from the inference model, and compare this value against λ∗. Additionally, we generate a new mini-batch given the inferred λ by running the generative model forward. We compare NVI against RWS, where we use 10 samples to estimate the gradients at each iteration. The results are shown in Figure 7. The cluster indices are rearranged based on the association of clusters to true classes. We observe that for RWS, even though the sample quality is reasonable, the posterior approximation is quite poor. When we train with NVI however, both the inferred λ and the generated samples match the test mini-batch reasonably well.

heading: 5 Conclusion

text:
We developed NVI, a framework that combines nested importance sampling and variational inference by optimizing a variational objective at every level of nesting. This formulation allows us to learn proposals and intermediate densities for a general class of samplers, which admit most commonly used importance sampling strategies as special cases. Our experiments demonstrate that samplers trained with NVI are able to outperform baselines when sampling from multimodal densities, Bayesian state-space models, and hierarchical deep generative models. Moreover, our experiments show that learning intermediate distributions results in better samplers.
NVI is particularly useful in the context of deep probabilistic programming systems. Because NVI can be applied to learn proposals for a wide variety of importance samplers, it can be combined with methods for inference programming that allow users to tailor sampling strategies to a particular probabilistic program. Concretely, NVI can be applied to nested importance samplers that are defined using a grammar of composable inference combinators [Stites et al., 2021], functions that implement primitive operations which preserve proper weighting, such as using samples from one program as a proposal to another program, sequential composition of programs, and importance resampling.

heading: Acknowledgements

text: This work was supported by the Intel Corporation, the 3M Corporation, NSF awards 1835309 and 2047253, startup funds from Northeastern University, the Air Force Research Laboratory (AFRL), and DARPA.

title: Nested Variational Inference