Joint Learning of Energy-based Models and their Partition Function

We thank the reviewer for the constructive feedback and relevant references.

Presenting at least one experiment comparing the proposed method and one of the SOTA methods for learning discrete EBMs is necessary

The original Table 3 contained generalized Fenchel-Young (GFY) losses comparisons, and Figures 1 and 3 contained comparison to logistic regression, which serves as a relevant baseline for structured prediction tasks learned without explicit partition function estimation. Following reviewers feedback, we have now added comparisons with MCMC and adversarial (min-max) training baselines.

Link to updated Table 1 and Table 3: https://anonymous.4open.science/r/ebms/tables.pdf

If the paper is accepted, we are going to include these tables in the main text, thanks to the extra page allowed by the camera-ready.

For convenience, we also show the new results below:

Table 1(label ranking)

Table 3 (multilabel)

Technical details: For the the min-max approach, we use optimistic ADAM as solver, an MLP as generator and we use REINFORCE (score function estimator) for gradient estimation. For MCMC, we use standard Metropolis–Hastings algorithm. Similarly to the min-min approach, we tune the learning rate and L2 regularization using grid search.

References [a-d]

Thank you for these references. We will add them.

Differences with AugSA [b]

Thank you for this very relevant reference. We will cite it and clarify the distinctions. First, this paper studies linear conditional random fields for whole-sentence LMs while we study more general energy-based models. In addition, our method supports the broader class of Fenchel-Young losses (including sparsemax), not just MLE. They indeed propose the idea to optimize the normalization constant but in their case it’s a constant, while in our case it’s a function, parameterized as a neural network that approximates the input-dependent log-partition function. Last, we empirically demonstrate (Figure 2) that our learned $\tau$ generalizes to unseen inputs $x$.

Not clear why the numerical metrics are not reported in the main text, which should be major results.

This was due to space constraints in the initial submission. If our paper is accepted, we plan to use the extra page allowed for the camera-ready version to move the expanded Table 3 (which now includes MCMC and a min-max approach) into the main text.

The right hand sides of Eq 2 and 5 are the same, which is confusing.

Please note that Eq. (2) is using an argmax while Eq. (5) is using a max.

Line 243: "We have therefore obtained convergence rates for learning EBMs in arbitrary combinatorial spaces." Not clear.

Our claim refers to the fact that our proposed objective can be optimized using standard stochastic gradient methods, with the convergence rate discussed after Proposition 2. This is achieved without requiring an oracle for the exact partition function (needed for traditional MLE gradient computation) or k-best solutions (needed for traditional sparsemax loss optimization), which are often intractable in arbitrary combinatorial spaces. We will clarify this context in the paper.

In section 3.3, where does the min-max formulation come from?

We derived the min-max formulation presented in Proposition 3 and Appendix C.4 ourselves. It serves as a comparison point to our proposed min-min formulation and draws parallels to adversarial learning approaches.

Figure 3: what are plots (a) and (b)?

Thank you for catching that typo. We were referring to the left and center columns of the Figure.

Thank you very much for your positive review and constructive feedback.

The approach is limited to a discrete target space

Our approach can in principle be applied to continuous output spaces $\mathcal{Y}$. However, we believe standard regression tasks are less common applications of conditional EBMs, compared to structured prediction over discrete spaces. This motivated our focus on discrete, combinatorially-large spaces like sets and permutations.

Figure 2: a single scalar number for each of the datasets would be more informative and a comparison with other techniques

Indeed, Figure 2 plots the learned vs. true log-partition for 100 individual test samples, rather than a time series. We chose this visualization to show the correlation across diverse samples. While scalar metrics like correlation coefficients could summarize this, the plot provides a better view of the approximation quality. We agree that summarizing with a scalar metric in the caption would be beneficial and will add this. A comparison with other techniques for learning the log-partition function is difficult as, to our knowledge, parameterizing and learning $\tau$ as a parametrized function is novel to our work.

Empirical comparison with other techniques

For multilabel classification, in Table 3, logistic (unary) corresponds exactly to using sigmoids as output layers and binary cross-entropy as loss. In addition, Table 3 includes the results when using the generalized Fenchel-Young loss, which is state-of-the-art for structured prediction tasks.

For label ranking, we use full rankings as supervision, which is naturally cast as a structured prediction task. We could potentially reduce full rankings to many pairwise rankings and use pairwise ranking losses instead, but we believe such a comparison is out of the scope of our paper.

However, we acknowledge that the paper did lack comparisons with MCMC and min-max baselines. We therefore added these comparisons to Table 1 and Table 3.

Link to updated Table 1 and Table 3: https://anonymous.4open.science/r/ebms/tables.pdf

If the paper is accepted, these tables are going to be included in the main text, thanks to the extra page allowed by the camera-ready.

Technical details: For the the min-max approach, we use optimistic ADAM as solver, an MLP as generator and we use REINFORCE (score function estimator) for gradient estimation. For MCMC, we use standard Metropolis–Hastings algorithm. Similarly to the min-min approach, we tune the learning rate and L2 regularization using grid search.

The variational perspective in equation (2) is generalised on page 6. NLL loss and KL divergence gives a standard "Bayes posterior" (which is not actually a posterior in this context). Arbitrary loss and KL divergence gives a "Gibbs posterior", which is equation (1). Arbitrary loss and arbitrary divergence (or, more specifically, f divergence) gives a generalised posterior. Maybe you would like to mention this (I guess "posterior" can be replaced by "measure")? --- see Knoblauch

This is a good remark. If $g(x, y) = \log p(x|y)$ and $q(y) = p(y)$, then Eq. (1) indeed gives the posterior $p(y|x)$ (in the sense of Bayes’ rule).

As far as I understand, the problem of sampling from the probabilistic EBM is not studied. Is that right? Would you expect any particular sampling algorithm to be especially well suited to this formulation of EBM, or are we no better off than generic samplers?

Indeed, we do not address the issue of sampling in this paper. We wonder whether the learned log-partition function can help design a better rejection sampling scheme. We leave this to future work.

The constraint is on the normalising constant equal to 1. Why not constrain the log normalising constant to be equal to zero? Wouldn't this then be of the same "scale" as the likelihood term in the loss?

We haven’t considered this. Our proof is based on the more general Fenchel-Young losses, where the Lagrange multiplier $\tau$ naturally appears as a consequence of Lemma 1.

We thank the reviewer for the constructive feedback.

Challenges of min-max implementation

To avoid any controversial claim, we will remove “notoriously hard to optimize”. Instead, we will explain the technical challenges that arise using this formulation in more detail.

The min-max objective is $\mathbb{E}_x \mathbb{E}{y \sim p(\cdot|x)} [g(x, y)] - \Omega(p(\cdot|x)) - \mathbb{E}{(x,y)} [g(x,y)]$, where $\Omega$ is the negative Shannon entropy. The variable $p$ appears in the sampling (expectation) and in the entropy. In practice, we parameterize $p$ as a neural network. If $p$ is an EBM, sampling from $p$ is difficult and the min-max formulation does not bring any benefit. Therefore, it is common to use a parameterization such that it is easy to sample from the mode instead.

Another challenge comes from the gradient computation. Since $p$ appears in the sampling, one typically uses the score function estimator (REINFORCE) to estimate the gradient w.r.t. the parameters of $p$. However, this estimator suffers from high variance. In our min-min formulation, the variable $\tau$ is a function, not a distribution and its gradients are easy to compute

High-variance of Monte-Carlo estimation of expectations

Our experiments (e.g., Table 1, Table 3) were designed to show scalability with respect to the output space size, governed by the number of classes $k$ ($|Y|=2^k$ or $k!$). Scalability with the number of training samples $n$ is addressed through standard SGD optimization. Regarding the MC estimation in Eq. (13), while MC estimates can introduce variance, our doubly stochastic approach (Alg. 1) proved stable and effective in practice across our experiments. Fig. 1 illustrates that even with a small number of prior samples ($y'$), the method converges effectively, although more samples accelerate convergence towards the exact MLE objective.

Lack of baseline comparison with other approaches, such as the variational-based min-max method.

Our paper included logistic regression (“Exact MLE” in Fig. 1 & 3) and generalized Fenchel-Young losses (GFY) (Table 3) as baselines. However, our paper did lack comparisons with MCMC and min-max approaches.

Links to updated Table 1 and Table 3 https://anonymous.4open.science/r/ebms/tables.pdf.

If the paper is accepted, we are going to include these tables in the main text, thanks to the extra page allowed by the camera-ready.

For convenience, we also show the new results below:

Table 1

Table 3

Technical details: For the the min-max approach, we use optimistic ADAM as solver, an MLP as generator and we use REINFORCE. For MCMC, we use Metropolis–Hastings algorithm. We tune the learning rate and L2 regularization using grid search.

Whether the proposed min-min formulation could also perform well in a probabilistic EBM setting

We indeed focus primarily on the prediction setting (i.e., finding the mode, $\mathrm{argmax}_y ~ p(y|x)$), because our approach, while learning the model $p(y|x)$ and approximating its normalization constant, doesn’t directly address the challenge of sampling from this distribution, which would be central to evaluating performance in a generative/probabilistic setting. Sampling in structured discrete spaces can be challenging. Thank you for pointing out the work of Ou et al, that we will cite.

Missing references

Thank you, we will add these relevant citations.

Thank you for your constructive feedback.

Lack of baselines, comparison with MCMC and adversarial approaches

While our paper includes GFY and exact MLE (logistic loss) baselines, it did lack comparisons with MCMC and min-max approaches. We have run additional experiments with these baselines.

Link to updated Table 1 and Table 3: https://anonymous.4open.science/r/ebms/tables.pdf

If the paper is accepted, these tables are going to be included in the main text, thanks to the extra page allowed by the camera-ready.

Technical details: For the the min-max approach, we use optimistic ADAM as solver, an MLP as generator and we use REINFORCE (score function estimator) for gradient estimation. For MCMC, we use standard Metropolis–Hastings algorithm. Similarly to the min-min approach, we tune the learning rate and L2 regularization using grid search.

Amortized optimization in the dual problem, neural OT [1]

We agree that parameterizing dual variables as neural networks is a commonly-used technique in the OT literature. We already cited Seguy et al. (2018) and will gladly add the suggested citation [1]. Technically, the dual variable in [1] is a convex function, hence the use of ICNNs. In our work, the Lagrange multiplier $\tau$ is a continuous function.

That said, we believe our application of that idea to EBMs remains novel and significant. To our knowledge, we are the first to parameterize the log-partition function of conditional EBMs as a neural network. We believe that parameterizing dual variables with neural networks is a powerful concept with potential applications beyond OT and EBMs.

Regarding amortized optimization, its goal is typically to accelerate the process of repeatedly solving similar optimization problems. In our paper, the primary goal of parameterizing the log-partition function $\tau$ as a neural network is different: it is to enable the approximation of the log-partition function to generalize to unseen data points $x$.

Why EBMs are particularly suited to this approximation over Monte Carlo-based methods?

As discussed in Section 2.4, sampling from EBMs, especially in discrete spaces, is often challenging and typically requires MCMC methods like Langevin-MCMC (for continuous spaces). For discrete output spaces, constructing an MCMC sampler is less explored and can be more challenging. Typically each particular output space (such as sets or permutations in our work) will require a specifically-designed sampler. In contrast, our proposed approach offers a simpler, more generic alternative that avoids MCMC sampling during training by jointly learning the energy and the log-partition function.

Relocating section on Fenchel-Young losses

We chose to include the Fenchel-Young losses because they encompass the sparsemax loss. Our work presents, to our knowledge, the first tractable method for optimizing the sparsemax loss for EBMs in general combinatorial spaces without needing a k-best oracle. Thanks to the extra page in the camera-ready, we believe we do not need to move this.

Experiments [in Figures 1 and 2] appear to be limited in scale and provide minimal insight

Our primary goal in these experiments was to demonstrate scalability concerning the number of classes $k$, which determines the size of the output space ($|Y| = 2^k$ for multilabel classification, $|Y| = k!$ for label-ranking). Scalability with respect to the number of training samples $n$ is handled via standard SGD.

Table 1 not discussed enough

The key takeaway from Table 1 is that the Birkhoff polytope representation tends to perform better with linear models for label ranking, likely due to its ability to capture more interactions. However, with more expressive MLP models, both the Birkhoff and Permutahedron representations achieve strong performance.

Eq 12 and 13 should be better introduced

Thank you for the suggestion. By minimizing the overall objective with respect to both $g$ and $\tau$, the optimization process encourages $\tau$ to approximate the true log-partition function while simultaneously learning the energy function $g$ that fits the data under the EBM framework. We will add a brief, self-contained explanation to the paper.

When using f-divergences what are the challenges in terms of optimization?

When $f$ is strictly convex, which is the case for the sparsemax loss, the loss is smooth, leading to a well-behaved optimization landscape. The primary challenge associated with the sparsemax loss in structured prediction is the need for a k-best oracle, which our approach alleviates (similarly, we do not need an LSE oracle in the logistic case).

Which other f-divergences could be interesting ?

We may consider alpha divergences, which interpolate between the logistic loss ($\alpha=1$) and the sparsemax loss ($\alpha=2$). Our approach can readily be applied with these divergences by selecting the appropriate $f$ function.