ContextGNN: Beyond Two-Tower Recommendation Systems

Yuan, Yiwen, Zhang, Zecheng, He, Xinwei, Nitta, Akihiro, Hu, Weihua, Wang, Dong, Shah, Manan, Huang, Shenyang, Stojanovič, Blaž, Krumholz, Alan, Lenssen, Jan Eric, Leskovec, Jure, Fey, Matthias

arXiv.org Artificial Intelligence 

Recommendation systems predominantly utilize two-tower architectures, which evaluate user-item rankings through the inner product of their respective embeddings. However, one key limitation of two-tower models is that they learn a pair-agnostic representation of users and items. In contrast, pair-wise representations either scale poorly due to their quadratic complexity or are too restrictive on the candidate pairs to rank. To address these issues, we introduce Context-based Graph Neural Networks (ContextGNNs), a novel deep learning architecture for link prediction in recommendation systems. The method employs a pair-wise representation technique for familiar items situated within a user's local subgraph, while leveraging two-tower representations to facilitate the recommendation of exploratory items. A final network then predicts how to fuse both pair-wise and two-tower recommendations into a single ranking of items. We demonstrate that ContextGNN is able to adapt to different data characteristics and outperforms existing methods, both traditional and GNN-based, on a diverse set of practical recommendation tasks, improving performance by 20% on average.