Improving Gradient Flow with Unrolled Highway Expectation Maximization

TL;DR

This paper introduces HEMNet, a novel neural network architecture that enhances gradient flow in unrolled EM algorithms using scaled skip connections, leading to better performance in semantic segmentation tasks.

Contribution

The paper proposes HEMNet, which incorporates scaled skip connections into unrolled EM algorithms to improve gradient flow without extra computational costs, while maintaining EM convergence properties.

Findings

Significant performance improvements on semantic segmentation benchmarks.

Effective alleviation of gradient decay in unrolled EM networks.

Preservation of EM algorithm convergence properties.

Abstract

Integrating model-based machine learning methods into deep neural architectures allows one to leverage both the expressive power of deep neural nets and the ability of model-based methods to incorporate domain-specific knowledge. In particular, many works have employed the expectation maximization (EM) algorithm in the form of an unrolled layer-wise structure that is jointly trained with a backbone neural network. However, it is difficult to discriminatively train the backbone network by backpropagating through the EM iterations as they are prone to the vanishing gradient problem. To address this issue, we propose Highway Expectation Maximization Networks (HEMNet), which is comprised of unrolled iterations of the generalized EM (GEM) algorithm based on the Newton-Rahpson method. HEMNet features scaled skip connections, or highways, along the depths of the unrolled architecture, resulting in improved gradient flow during backpropagation while incurring negligible additional computation and memory costs compared to standard unrolled EM. Furthermore, HEMNet preserves the underlying EM procedure, thereby fully retaining the convergence properties of the original EM algorithm. We achieve significant improvement in performance on several semantic segmentation benchmarks and empirically show that HEMNet effectively alleviates gradient decay.