Timezone: »

Wasserstein Barycenter Matching for Graph Size Generalization of Message Passing Neural Networks
Xu Chu · Yujie Jin · Xin Wang · Shanghang Zhang · Yasha Wang · Wenwu Zhu · Hong Mei

Tue Jul 25 05:00 PM -- 06:30 PM (PDT) @ Exhibit Hall 1 #225

Graph size generalization is hard for Message passing neural networks (MPNNs). The graph-level classification performance of MPNNs degrades across various graph sizes. Recently, theoretical studies reveal that a slow uncontrollable convergence rate w.r.t. graph size could adversely affect the size generalization. To address the uncontrollable convergence rate caused by correlations across nodes in the underlying dimensional signal-generating space, we propose to use Wasserstein barycenters as graph-level consensus to combat node-level correlations. Methodologically, we propose a Wasserstein barycenter matching (WBM) layer that represents an input graph by Wasserstein distances between its MPNN-filtered node embeddings versus some learned class-wise barycenters. Theoretically, we show that the convergence rate of an MPNN with a WBM layer is controllable and independent to the dimensionality of the signal-generating space. Thus MPNNs with WBM layers are less susceptible to slow uncontrollable convergence rate and size variations. Empirically, the WBM layer improves the size generalization over vanilla MPNNs with different backbones (e.g., GCN, GIN, and PNA) significantly on real-world graph datasets.

Author Information

Xu Chu (Tsinghua University, Tsinghua University)
Yujie Jin (Peking University)
Xin Wang (Tsinghua University)
Shanghang Zhang (UC Berkeley)
Yasha Wang (Peking University)
Wenwu Zhu (Tsinghua University)

Wenwu Zhu is currently a Professor of Computer Science Department of Tsinghua University and Vice Dean of National Research Center on Information Science and Technology. Prior to his current post, he was a Senior Researcher and Research Manager at Microsoft Research Asia. He was the Chief Scientist and Director at Intel Research China from 2004 to 2008. He worked at Bell Labs New Jersey as a Member of Technical Staff during 1996-1999. He has been serving as the chair of the steering committee for IEEE T-MM since January 1, 2020. He served as the Editor-in-Chief for the IEEE Transactions on Multimedia (T-MM) from 2017 to 2019. And Vice EiC for IEEE Transactions on Circuits and Systems for Video Technology (TCSVT) from 2020-2021 He served as co-Chair for ACM MM 2018 and co-Chair for ACM CIKM 2019. His current research interests are in the areas of multimodal big data and intelligence, and multimedia networking. He received 10 Best Paper Awards. He is a member of Academia Europaea, an IEEE Fellow, AAAS Fellow, and SPIE Fellow.

Hong Mei (Peking University)

More from the Same Authors