Learning Proposals for Practical Energy-Based Regression

TL;DR

This paper introduces a simple method to automatically learn proposal distributions for energy-based regression models, improving training and prediction efficiency, and demonstrates superior performance over traditional methods on real-world tasks.

Contribution

It proposes a unified training objective for jointly learning proposal distributions and EBMs, and applies this to enhance mixture density networks with an energy-based teacher.

Findings

Outperforms conventional MDNs on four regression tasks

Enables efficient importance sampling at test-time

Provides a unified training framework for proposals and EBMs

Abstract

Energy-based models (EBMs) have experienced a resurgence within machine learning in recent years, including as a promising alternative for probabilistic regression. However, energy-based regression requires a proposal distribution to be manually designed for training, and an initial estimate has to be provided at test-time. We address both of these issues by introducing a conceptually simple method to automatically learn an effective proposal distribution, which is parameterized by a separate network head. To this end, we derive a surprising result, leading to a unified training objective that jointly minimizes the KL divergence from the proposal to the EBM, and the negative log-likelihood of the EBM. At test-time, we can then employ importance sampling with the trained proposal to efficiently evaluate the learned EBM and produce stand-alone predictions. Furthermore, we utilize our derived training objective to learn mixture density networks (MDNs) with a jointly trained energy-based teacher, consistently outperforming conventional MDN training on four real-world regression tasks within computer vision. Code is available at https://github.com/fregu856/ebms_proposals.