[Paper Review] Fishnets: Information-Optimal, Scalable Aggregation for Sets and Graphs

T. Lucas Makinen, Justin Alsing+1 more|arXiv (Cornell University)|Oct 5, 2023

Advanced Graph Neural Networks4 citations

TL;DR

Fishnets proposes an information-theoretic aggregation method for sets and graphs that learns optimal, scalable summaries using neural networks to parameterize score and inverse-Fisher information. It achieves state-of-the-art performance on ogbn-protein with fewer parameters and faster training by explicitly modeling uncertainty and information content, outperforming standard GNNs and deepsets in robustness and information saturation.

ABSTRACT

Set-based learning is an essential component of modern deep learning and network science. Graph Neural Networks (GNNs) and their edge-free counterparts Deepsets have proven remarkably useful on ragged and topologically challenging datasets. The key to learning informative embeddings for set members is a specified aggregation function, usually a sum, max, or mean. We propose Fishnets, an aggregation strategy for learning information-optimal embeddings for sets of data for both Bayesian inference and graph aggregation. We demonstrate that i) Fishnets neural summaries can be scaled optimally to an arbitrary number of data objects, ii) Fishnets aggregations are robust to changes in data distribution, unlike standard deepsets, iii) Fishnets saturate Bayesian information content and extend to regimes where MCMC techniques fail and iv) Fishnets can be used as a drop-in aggregation scheme within GNNs. We show that by adopting a Fishnets aggregation scheme for message passing, GNNs can achieve state-of-the-art performance versus architecture size on ogbn-protein data over existing benchmarks with a fraction of learnable parameters and faster training time.

Motivation & Objective

• To develop an optimal, scalable aggregation method for sets and graphs grounded in information theory.
• To address the limitations of fixed aggregators (e.g., mean, max, sum) in deep learning and Bayesian inference.
• To enable robust, information-saturating summaries under distributional shifts and noisy data.
• To integrate the formalism into GNNs as a drop-in replacement for improved performance and efficiency.

Proposed method

• Fishnets parameterize per-data-point score and Fisher information using twin neural networks to model local contributions to the log-likelihood.
• The method aggregates scores and inverse-Fisher weights additively, forming a pseudo-maximum likelihood estimate via Fisher scoring.
• It explicitly learns inverse-Fisher weights as a function of data, enabling uncertainty-aware aggregation.
• The formalism is applied to both Bayesian simulation-based inference (SBI) and graph neural networks (GNNs).
• In GNNs, Fishnets replaces standard aggregation layers, using edge uncertainty (e.g., N in coin toss simulations) as input for adaptive weighting.
• The approach is validated across PyTorch Geometric, JAX, and TensorFlow, demonstrating cross-framework compatibility.

Figure 1: Fishnets saturate information for datasets 20 times larger than the training set. Residual maximum likelihood estimates for slope ( left ) and intercept ( right ) scatter about the truth for linear regression test datasets of size $n_{\rm data}=10^{4}$ . The solid pink line is obtained fro

Experimental results

Research questions

• RQ1 Can an aggregation scheme be designed to saturate Bayesian information content in complex, ragged datasets?
• RQ2 How does Fishnets perform under distributional shift compared to standard deepsets and GNNs?
• RQ3 Can Fishnets achieve state-of-the-art GNN performance with significantly fewer learnable parameters?
• RQ4 Does explicit modeling of inverse-Fisher weights improve robustness to noisy or uncertain data?
• RQ5 Can Fishnets be used as a plug-in replacement for standard GNN aggregation layers without architectural changes?

Key findings

• Fishnets achieved a test ROC-AUC of 0.8110 ± 0.0021 on ogbn-protein with only 442,500 parameters, outperforming larger GCN-28 models.
• On a noisy edge variant of the proteins dataset, Fishnets achieved 71.98% ROC-AUC, a 10.5% improvement over GCN-28’s 64.71%.
• Fishnets saturated Bayesian information content and demonstrated robustness to distributional shifts, unlike standard deepsets.
• The method achieved information saturation in non-trivial inference problems, even in regimes where MCMC techniques fail.
• Fishnets enabled faster training and lower parameter counts than leading GNNs while maintaining or exceeding performance on benchmarks.
• The approach demonstrated explicit information-level inductive bias, improving generalization under uncertainty and edge noise.

Figure 2: Fishnets achieve the exact form of the score as a function of input data in the linear regression case, indicating information saturation. Slices of true (dark) and network predicted (pink) score vector components as a function of data inputs for the $n_{\rm data}=10^{4}$ test set.