# GNN-based Unified Deep Learning

arXiv: https://arxiv.org/abs/2508.10583 · HTML: https://mlsanity.com/paper/2508.10583/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 15 of its 24 references that are in the map.
- First posted on arXiv 14 Aug 2025.

## Builds on (most cited first, 15 of 15)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classification](https://mlsanity.com/paper/2110.14795/) (2021)
- [Deep Learning for Medical Image Analysis](https://mlsanity.com/paper/1708.08987/) (2017)
- [Learning to Generate Novel Domains for Domain Generalization](https://mlsanity.com/paper/2007.03304/) (2020)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Unsupervised domain adaptation in brain lesion segmentation with adversarial networks](https://mlsanity.com/paper/1612.08894/) (2016)
- [Graph Neural Networks for Learning Equivariant Representations of Neural Networks](https://mlsanity.com/paper/2403.12143/) (2024)
- [Morpho-MNIST: Quantitative Assessment and Diagnostics for Representation Learning](https://mlsanity.com/paper/1809.10780/) (2018)
- [Domain Generalization for Medical Image Analysis: A Review](https://mlsanity.com/paper/2310.08598/) (2023)
- [Domain Generalisation via Risk Distribution Matching](https://mlsanity.com/paper/2310.18598/) (2023)

## Built on by (most cited first, 1 of 1)

- [UnifiedFL: A Dynamic Unified Learning Framework for Equitable Federation](https://mlsanity.com/paper/2510.26350/) (2025)

## Abstract (arXiv)

Deep learning models often struggle to maintain generalizability in medical imaging, particularly under domain-fracture scenarios where distribution shifts arise from varying imaging techniques, acquisition protocols, patient populations, demographics, and equipment. In practice, each hospital may need to train distinct models - differing in learning task, width, and depth - to match local data. For example, one hospital may use Euclidean architectures such as MLPs and CNNs for tabular or grid-like image data, while another may require non-Euclidean architectures such as graph neural networks (GNNs) for irregular data like brain connectomes. How to train such heterogeneous models coherently across datasets, while enhancing each model's generalizability, remains an open problem. We propose unified learning, a new paradigm that encodes each model into a graph representation, enabling unification in a shared graph learning space. A GNN then guides optimization of these unified models. By decoupling parameters of individual models and controlling them through a unified GNN (uGNN), our method supports parameter sharing and knowledge transfer across varying architectures (MLPs, CNNs, GNNs) and distributions, improving generalizability. Evaluations on MorphoMNIST and two MedMNIST benchmarks - PneumoniaMNIST and BreastMNIST - show that unified learning boosts performance when models are trained on unique distributions and tested on mixed ones, demonstrating strong robustness to unseen data with large distribution shifts. Code and benchmarks: https://github.com/basiralab/uGNN