Fast and Differentiable Message Passing on Pairwise Markov Random Fields

TL;DR

This paper introduces two fast, GPU-accelerated, differentiable message passing algorithms for pairwise MRFs, significantly improving inference speed and effectiveness for applications like stereo matching, denoising, and semantic segmentation.

Contribution

The authors present ISGMR and TRWP, novel parallel algorithms that enhance message passing efficiency and differentiability, enabling faster end-to-end learning for MRF-based models.

Findings

ISGMR and TRWP outperform traditional methods in energy minimization.

TRWP is two orders of magnitude faster than TRWS.

CUDA implementations are 7x and 700x faster than PyTorch for forward/backward passes.

Abstract

Despite the availability of many Markov Random Field (MRF) optimization algorithms, their widespread usage is currently limited due to imperfect MRF modelling arising from hand-crafted model parameters and the selection of inferior inference algorithm. In addition to differentiability, the two main aspects that enable learning these model parameters are the forward and backward propagation time of the MRF optimization algorithm and its inference capabilities. In this work, we introduce two fast and differentiable message passing algorithms, namely, Iterative Semi-Global Matching Revised (ISGMR) and Parallel Tree-Reweighted Message Passing (TRWP) which are greatly sped up on a GPU by exploiting massive parallelism. Specifically, ISGMR is an iterative and revised version of the standard SGM for general pairwise MRFs with improved optimization effectiveness, and TRWP is a highly parallel version of Sequential TRW (TRWS) for faster optimization. Our experiments on the standard stereo and denoising benchmarks demonstrated that ISGMR and TRWP achieve much lower energies than SGM and Mean-Field (MF), and TRWP is two orders of magnitude faster than TRWS without losing effectiveness in optimization. We further demonstrated the effectiveness of our algorithms on end-to-end learning for semantic segmentation. Notably, our CUDA implementations are at least $7$ and $700$ times faster than PyTorch GPU implementations for forward and backward propagation respectively, enabling efficient end-to-end learning with message passing.