[论文解读] Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration

Hanjun Dai, Rishabh Singh+3 more|arXiv (Cornell University)|Jan 1, 2020

Machine Learning and Data Classification被引用 2

一句话总结

本文提出 ALOE，一种通过学习采样器模拟局部搜索来训练离散能量模型（EBM）的方法，实现通过变分幂迭代进行高效梯度估计。该方法在结构化预测任务中达到最先进性能，包括在软件测试中实现与 libfuzzer 相当的有效性。

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.

研究动机与目标

• 解决在离散结构化数据上训练能量模型的挑战，其中分区函数估计不可行。
• 在保持计算可处理性的同时，提升对自回归模型的灵活性，以建模复杂离散分布。
• 开发一种可微分的、受局部搜索启发的采样器，以实现对 EBM 的有效梯度估计。
• 通过变分幂迭代框架实现能量函数与采样器的端到端联合训练。
• 在真实应用中展示该方法的有效性，如程序语言建模与软件模糊测试。

提出的方法

• 提出一种学习采样器，通过在离散空间中探索邻近状态来模拟局部搜索，实现高效采样以支持梯度估计。
• 引入一种变分形式的幂迭代，联合优化能量函数与采样器，提升训练稳定性和收敛性。
• 通过使用学习采样器进行重要性采样，对分区函数进行可微分近似，以估计梯度。
• 利用采样器的探索能力，通过变分下界实现能量函数与采样器的端到端联合训练。
• 采用对比学习目标，促使有效结构的能量较低，无效结构的能量较高。
• 将该框架应用于无条件与条件建模，支持结构生成与模糊测试引导等应用。

实验结果

研究问题

• RQ1 学习的局部搜索采样器是否能在不显式计算分区函数的情况下，有效估计离散能量模型中的梯度？
• RQ2 所提出的变分幂迭代框架是否能实现能量函数与采样器的稳定且高效的联合训练？
• RQ3 ALOE 在建模复杂离散结构（如程序）方面与自回归模型相比表现如何？
• RQ4 训练后的 EBM 是否能像手工设计的工具（如 libfuzzer）一样，有效引导模糊测试发现漏洞？
• RQ5 采样器中的局部探索在多大程度上提升了离散 EBM 学习中的样本质量与训练效率？

主要发现

• ALOE 在程序语言建模中表现优异，生成的有效且多样化的代码序列优于自回归基线模型。
• 学习采样器有效捕捉了局部结构依赖关系，实现在无需显式归一化的情况下生成高质量样本。
• 变分幂迭代方法实现了能量函数与采样器的稳定训练，减少了训练发散现象。
• 在软件模糊测试中，EBM 引导的模糊测试器在发现漏洞方面与 libfuzzer 表现相当，展现出实际应用价值。
• 与标准离散 EBM 训练基线相比，该方法在样本质量与训练效率方面均有显著提升。
• 该方法在不同离散结构化领域（包括代码生成与漏洞发现）中均表现出良好的泛化能力。