Fast and Differentiable Message Passing on Pairwise Markov Random Fields

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. 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 parallelization 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 GPU by exploiting massive parallelism. Specifically, ISGMR is an iterative and revised version of the standard SGM for general second-order MRFs with improved optimization effectiveness, whereas TRWP is a highly parallelizable version of Sequential TRW (TRWS) for faster optimization. Our experiments on standard stereo benchmarks demonstrate that ISGMR achieves much lower energies than SGM and TRWP is two orders of magnitude faster than TRWS without losing effectiveness in optimization. Furthermore, our CUDA implementations are at least 7 and 650 times faster than PyTorch GPU implementations in the forward and backward propagation, respectively, enabling efficient end-to-end learning with message passing.