[Paper Review] Bag Graph: Multiple Instance Learning using Bayesian Graph Neural Networks

Soumyasundar Pal, Antonios Valkanas+2 more|arXiv (Cornell University)|Feb 22, 2022

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes Bag Graph, a Bayesian Graph Neural Network framework for Multiple Instance Learning that jointly learns bag-level representations and infers relational graph structures between bags when no prior graph is available. By integrating neural MIL models with a differentiable Bayesian GNN, it achieves state-of-the-art performance on classification benchmarks and superior results in text categorization, electoral prediction, and rental price regression tasks.

ABSTRACT

Multiple Instance Learning (MIL) is a weakly supervised learning problem where the aim is to assign labels to sets or bags of instances, as opposed to traditional supervised learning where each instance is assumed to be independent and identically distributed (IID) and is to be labeled individually. Recent work has shown promising results for neural network models in the MIL setting. Instead of focusing on each instance, these models are trained in an end-to-end fashion to learn effective bag-level representations by suitably combining permutation invariant pooling techniques with neural architectures. In this paper, we consider modelling the interactions between bags using a graph and employ Graph Neural Networks (GNNs) to facilitate end-to-end learning. Since a meaningful graph representing dependencies between bags is rarely available, we propose to use a Bayesian GNN framework that can generate a likely graph structure for scenarios where there is uncertainty in the graph or when no graph is available. Empirical results demonstrate the efficacy of the proposed technique for several MIL benchmark tasks and a distribution regression task.

Motivation & Objective

• Address the limitation of existing MIL methods that assume i.i.d. instances and fail to model dependencies between instances or bags.
• Overcome the challenge of lacking or noisy relational graph information between bags in real-world MIL applications.
• Develop an end-to-end framework that jointly learns bag representations and infers meaningful graph topologies between bags using a Bayesian GNN.
• Enable both inductive and transductive learning in MIL by leveraging uncertainty-aware graph structure inference.
• Demonstrate the method's effectiveness across diverse MIL tasks, including classification, text categorization, and distribution regression.

Proposed method

• Integrate instance-level neural networks (e.g., Deep Sets, Set Transformer) to learn permutation-invariant bag representations from sets of instances.
• Construct a Bayesian Graph Neural Network (Bayesian GNN) that jointly learns the graph structure between bags and performs label prediction via message passing.
• Use a variational inference framework to infer the posterior distribution over graph edges, enabling uncertainty-aware graph learning when no ground-truth graph exists.
• Train the entire architecture end-to-end using differentiable graph inference, allowing gradients to flow through both the bag representation and graph topology learning components.
• Employ a transductive setting during training and inference, where graph structure is jointly inferred for both training and test bags to improve generalization.
• Utilize observed spatial or relational side information (e.g., geographic proximity) as an initial graph prior, which is refined by the Bayesian GNN during training.

Experimental results

Research questions

• RQ1 Can a Bayesian GNN framework effectively infer meaningful relational structures between bags in the absence of prior graph information in MIL tasks?
• RQ2 How does end-to-end learning of bag representations and graph topology improve performance compared to models that treat bags as independent or use fixed graphs?
• RQ3 To what extent does the proposed method outperform state-of-the-art MIL models in both classification and regression settings?
• RQ4 Does the transductive graph inference setting provide a performance advantage over inductive or disconnected test-node settings?
• RQ5 Can the method generalize across diverse domains such as text categorization, electoral outcome prediction, and real estate price regression?

Key findings

• The proposed B-DS-GCN and B-ST-GCN models achieve the lowest RMSE (67.51 ± 16.39) and MAE (47.24 ± 10.21) in rental price regression, significantly outperforming baseline models like Deep Sets (RMSE: 86.37 ± 20.41).
• In the ablation study, transductive graph inference (including test nodes in graph learning) reduced RMSE by 10.6% compared to disconnected test-node settings, demonstrating its effectiveness.
• The Bayesian GNN variants (B-DS-GCN and B-ST-GCN) achieved the best performance in electoral result prediction and text categorization, outperforming non-graph baselines and GCN variants.
• The method achieved comparable or better performance than state-of-the-art MIL models on standard benchmark datasets, even when no ground-truth graph was provided.
• The use of a Bayesian framework enabled uncertainty-aware graph learning, allowing the model to adaptively learn complex relationships beyond simple geographic proximity.
• Empirical results show that the proposed method is robust across multiple tasks, with consistent improvements in RMSE, MAE, and MAPE across 100 random train-test splits in the regression experiment.