Smooth InfoMax -- Towards Easier Post-Hoc Interpretability

TL;DR

Smooth InfoMax (SIM) is a self-supervised learning method that enhances interpretability of neural network representations by creating smooth, disentangled latent spaces through probabilistic modules and the InfoNCE loss.

Contribution

The paper introduces SIM, a novel approach that integrates interpretability constraints into deep representations using probabilistic modules and regularization, improving post-hoc analysis.

Findings

Creates smooth, well-disentangled latent spaces

Improves effectiveness of post-hoc interpretability methods

Maintains large-scale training benefits of Greedy InfoMax

Abstract

We introduce Smooth InfoMax (SIM), a self-supervised representation learning method that incorporates interpretability constraints into the latent representations at different depths of the network. Based on $\beta$-VAEs, SIM's architecture consists of probabilistic modules optimized locally with the InfoNCE loss to produce Gaussian-distributed representations regularized toward the standard normal distribution. This creates smooth, well-defined, and better-disentangled latent spaces, enabling easier post-hoc analysis. Evaluated on speech data, SIM preserves the large-scale training benefits of Greedy InfoMax while improving the effectiveness of post-hoc interpretability methods across layers.