Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration

TL;DR

This paper introduces ALOE, a novel algorithm for learning energy-based models on discrete structures using a learned local search sampler, improving flexibility and efficiency in applications like software testing.

Contribution

ALOE employs a new variational power iteration method to train EBMs with learned local search, enhancing modeling of discrete structured data.

Findings

Significant improvements in discrete structure modeling.

Energy model guided fuzzer achieves performance comparable to libfuzzer.

Efficient training of EBMs with local search in complex domains.

Abstract

Discrete structures play an important role in applications like program language modeling and software engineering. Current approaches to predicting complex structures typically consider autoregressive models for their tractability, with some sacrifice in flexibility. Energy-based models (EBMs) on the other hand offer a more flexible and thus more powerful approach to modeling such distributions, but require partition function estimation. In this paper we propose ALOE, a new algorithm for learning conditional and unconditional EBMs for discrete structured data, where parameter gradients are estimated using a learned sampler that mimics local search. We show that the energy function and sampler can be trained efficiently via a new variational form of power iteration, achieving a better trade-off between flexibility and tractability. Experimentally, we show that learning local search leads to significant improvements in challenging application domains. Most notably, we present an energy model guided fuzzer for software testing that achieves comparable performance to well engineered fuzzing engines like libfuzzer.