The Hessian Penalty:
A Weak Prior for Unsupervised Disentanglement
William Peebles      John Peebles      Jun-Yan Zhu      Alexei Efros      Antonio Torralba
UC Berkeley
Yale
Adobe Research
MIT CSAIL
paper | code | 90 sec video | 10 min video
ECCV 2020 (Spotlight)

Download video

Abstract

Existing disentanglement methods for deep generative models rely on hand-picked priors and complex encoder-based architectures. In this paper, we propose the Hessian Penalty, a simple regularization term that encourages the Hessian of a generative model with respect to its input to be diagonal. We introduce a model-agnostic, unbiased stochastic approximation of this term based on Hutchinson's estimator to compute it efficiently during training. Our method can be applied to a wide range of deep generators with just a few lines of code. We show that training with the Hessian Penalty often causes axis-aligned disentanglement to emerge in latent space when applied to ProGAN on several datasets. Additionally, we use our regularization term to identify interpretable directions in BigGAN's latent space in an unsupervised fashion. Finally, we provide empirical evidence that the Hessian Penalty encourages substantial shrinkage when applied to over-parameterized latent spaces.

Video (10 min)

Download video

Complete Results

Below, we show all z components learned by (1) ProGAN fine-tuned with the Hessian Penalty, (2) ProGAN trained from scratch with the Hessian Penalty, (3) ProGAN trained with InfoGAN loss and (4) vanilla ProGAN. Each row corresponds to varying a single z component from -2 to +2, leaving the remaining ones fixed. Different columns in the same row re-sample the fixed z components. All models were trained with 12-dimensional latent z vectors---with the exception of CLEVR-U which used |z|=3---hence why there are 12 rows. These results are not cherry-picked.

Edges+Shoes

Hessian Penalty
(Fine-Tuned)

Hessian Penalty

InfoGAN

Vanilla

CLEVR-Simple

Hessian Penalty
(Fine-Tuned)

Hessian Penalty

InfoGAN

Vanilla

CLEVR-1FOV

Hessian Penalty
(Fine-Tuned)

Hessian Penalty

InfoGAN

Vanilla

CLEVR-U

Hessian Penalty
(Fine-Tuned)

Hessian Penalty

InfoGAN

Vanilla

CLEVR-Complex

Hessian Penalty
(Fine-Tuned)

Hessian Penalty

InfoGAN

Vanilla

BigGAN Direction Discovery Results

Below, we show several linear z-space directions learned via the Hessian Penalty. These directions are learned in BigGAN trained on ImageNet.

Try our code

We released portable implementations of the Hessian Penalty in and TensorFlow for your use in your code bases. They can be used with just one line of code.

[GitHub]

Paper

W. Peebles, J. Peebles, J.-Y. Zhu, A. Efros, A. Torralba
The Hessian Penalty: A Weak Prior for Unsupervised Disentanglement.
In ECCV, 2020 (Spotlight). arxiv

[Bibtex]

Acknowledgements

We thank Pieter Abbeel, Taesung Park, Richard Zhang, Mathieu Aubry, Ilija Radosavovic, Tim Brooks, Karttikeya Mangalam and all of BAIR for valuable discussions and encouragement. This work was supported, in part, by grants from SAP, Adobe and Berkeley DeepDrive. We further thank Taesung Park for his project page template.