Combating the Instability of Mutual Information-based Losses via Regularization

TL;DR

This paper addresses the instability issues in mutual information-based neural network training by introducing a novel regularization technique, demonstrating improved stability and performance across various benchmarks.

Contribution

The authors propose a new regularization method that stabilizes MI-based loss training and introduce a benchmark for evaluating MI estimation and downstream task performance.

Findings

Regularization stabilizes MI-based loss training.

Regularized losses outperform unregularized counterparts.

Benchmark evaluates MI estimation and downstream task performance.

Abstract

Notable progress has been made in numerous fields of machine learning based on neural network-driven mutual information (MI) bounds. However, utilizing the conventional MI-based losses is often challenging due to their practical and mathematical limitations. In this work, we first identify the symptoms behind their instability: (1) the neural network not converging even after the loss seemed to converge, and (2) saturating neural network outputs causing the loss to diverge. We mitigate both issues by adding a novel regularization term to the existing losses. We theoretically and experimentally demonstrate that added regularization stabilizes training. Finally, we present a novel benchmark that evaluates MI-based losses on both the MI estimation power and its capability on the downstream tasks, closely following the pre-existing supervised and contrastive learning settings. We evaluate six different MI-based losses and their regularized counterparts on multiple benchmarks to show that our approach is simple yet effective.