Graph Neural Network-Informed Predictive Flows for Faster Ford-Fulkerson and PAC-Learnability

TL;DR

This paper introduces a GNN-guided augmentation method for max-flow problems that improves efficiency in image segmentation by learning edge importance, guiding augmenting path selection, and reducing the number of augmentations needed.

Contribution

It presents a novel GNN-based framework that predicts edge importance to accelerate max-flow algorithms, maintaining optimality while reducing computational steps.

Findings

Reduces the number of augmentations in max-flow computation.

Preserves max-flow/min-cut optimality with fewer iterations.

Provides a theoretical link between prediction quality and efficiency.

Abstract

We propose a learning-augmented framework for accelerating max-flow computation and image segmentation by integrating Graph Neural Networks (GNNs) with the Ford-Fulkerson algorithm. Rather than predicting initial flows, our method learns edge importance probabilities to guide augmenting path selection. We introduce a Message Passing GNN (MPGNN) that jointly learns node and edge embeddings through coupled updates, capturing both global structure and local flow dynamics such as residual capacity and bottlenecks. Given an input image, we propose a method to construct a grid-based flow network with source and sink nodes, extract features, and perform a single GNN inference to assign edge probabilities reflecting their likelihood of belonging to high-capacity cuts. These probabilities are stored in a priority queue and used to guide a modified Ford-Fulkerson procedure, prioritizing augmenting paths via an Edmonds-Karp-style search with bottleneck-aware tie-breaking. This avoids repeated inference over residual graphs while leveraging learned structure throughout optimization. We further introduce a bidirectional path construction strategy centered on high-probability edges and provide a theoretical framework relating prediction quality to efficiency via a weighted permutation distance metric. Our method preserves max-flow/min-cut optimality while reducing the number of augmentations in practice. We also outline a hybrid extension combining flow warm-starting with edge-priority prediction, establishing a foundation for learning-guided combinatorial optimization in image segmentation.