Improving Gradient Flow with Unrolled Highway Expectation Maximization
This work provides an incremental improvement for researchers integrating model-based EM algorithms into deep learning architectures, specifically by improving gradient flow during training.
This paper addresses the vanishing gradient problem when discriminatively training neural networks with unrolled Expectation Maximization (EM) layers. They propose Highway Expectation Maximization Networks (HEMNet), which uses scaled skip connections within the unrolled EM iterations to improve gradient flow, leading to significant performance improvements on semantic segmentation benchmarks.
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.