[Paper Review] Neural Relational Inference with Fast Modular Meta-learning

Ferran Alet, Erica Weng+2 more|arXiv (Cornell University)|Oct 10, 2023

Domain Adaptation and Few-Shot Learning19 references25 citations

TL;DR

The paper reframes neural relational inference as modular meta-learning over GNNs, learning a library of node and edge modules and using a fast proposal-assisted search (BounceGrad) to infer relational structure, achieving data-efficient, model-based inference and scalable planning for unseen or latent nodes.

ABSTRACT

extit{Graph neural networks} (GNNs) are effective models for many dynamical systems consisting of entities and relations. Although most GNN applications assume a single type of entity and relation, many situations involve multiple types of interactions. extit{Relational inference} is the problem of inferring these interactions and learning the dynamics from observational data. We frame relational inference as a extit{modular meta-learning} problem, where neural modules are trained to be composed in different ways to solve many tasks. This meta-learning framework allows us to implicitly encode time invariance and infer relations in context of one another rather than independently, which increases inference capacity. Framing inference as the inner-loop optimization of meta-learning leads to a model-based approach that is more data-efficient and capable of estimating the state of entities that we do not observe directly, but whose existence can be inferred from their effect on observed entities. To address the large search space of graph neural network compositions, we meta-learn a extit{proposal function} that speeds up the inner-loop simulated annealing search within the modular meta-learning algorithm, providing two orders of magnitude increase in the size of problems that can be addressed.

Motivation & Objective

• Infer relational structure and dynamics of multi-relational systems from observations using modular meta-learning.
• Learn a library of neural modules for nodes and edges that can be composed to model various interaction types.
• Increase data efficiency and enable inference of latent/unobserved entities by treating relational inference as meta-learning with a model-based approach.
• Speed up the combinatorial search over GNN structures through a learned proposal function to enable larger-scale problems.

Proposed method

• Frame relational inference as modular meta-learning with node modules m_i and edge modules h_j.
• Use BounceGrad: alternate simulated annealing over modular structure and gradient descent to optimize module weights.
• Meta-train by learning module weights across many tasks, meta-test by inferring the best structure S for a new task.
• Meta-learn a proposal function to accelerate inner-loop structure search, trained using outcomes of SA on training batches.
• Batch training by creating a super-graph of many tasks to exploit parallelism and maintain constant parameter cost.
• Leverage model-based inference to predict unseen entities by optimizing their latent states within the inner loop.

Figure 1: Modular meta-learning with graph networks; adapted from Alet et al. ( 2018 ) . The system meta-learns a library of node and edge modules, represented as small neural networks; at performance (meta-test) time, it is only necessary to infer the combination of modules that best predict the ob

Experimental results

Research questions

• RQ1 Can modular meta-learning with a library of neural modules match or surpass variational NRI in inferring relational structures from trajectory data?
• RQ2 Does joint, structure-level optimization of edge/types improve relational inference compared to independent edge predictions?
• RQ3 Can a learned proposal function dramatically speed up the inner-loop search to scale to larger task distributions and modular compositions?
• RQ4 Is it possible to infer latent/unobserved entities and their relations by integrating them into the inner-loop optimization?

Key findings

• Modular meta-learning with jointly inferred edge modules achieves comparable or superior trajectory-prediction performance to baselines on Springs and Charged-particles datasets.
• Edge-type prediction accuracy reaches near-gold performance, with modular meta-learning matching or surpassing prior methods in both domains.
• The model-based approach exhibits higher data efficiency than variational methods, particularly in low-data regimes.
• A learned proposal function speeds up the inner-loop search by orders of magnitude, enabling larger modular spaces and meta-datasets (up to 50,000 tasks).
• The framework can accommodate inferring latent/unobserved nodes by optimizing their trajectories within the inner loop, improving predictive accuracy.
• Batching across tasks and mapping computations onto a super-graph yields substantial speed-ups in training and evaluation.

Figure 2: Task setup, taken from Kipf et al. ( 2018 ) : we want to predict the evolution of a dynamical system by inferring the set of relations between the entities, such as attraction and repulsion between charged particles.