Sourcerer: Sample-based Maximum Entropy Source Distribution Estimation

We thank the reviewer for the positive feedback of our work.

W/Q1: "Is there something that prevents you from comparing against NEB in Figure 4?" We thank the reviewer for their suggestion, and have now compared Neural Empirical Bayes (NEB) with Sourcerer on the high-dimensional simulators (SIR and Lotka-Volterra model). Results for this experiment can be found in the provided rebuttal supplement (Fig. R2). In summary, we find that Sourcerer performs better on both tasks: on SIR, Sourcerer achieves a C2ST accuracy (lower is better) of 55% between data and predictive simulations of learned source distribution, compared to 76% for NEB; similarly, on Lotka-Volterra, Sourcerer achieves a C2ST of 56%, compared to 62% for NEB. Thus, Sourcerer is demonstrably better than NEB on these high-dimensional tasks.

To perform this experiment, we trained surrogates for each simulator. While this was not necessary for Sourcerer, attempting to perform NEB on the differentiable likelihood of the differential equation simulators was computationally infeasible due to the large minibatches required for NEB. To enable surrogate training, we reduced the observation dimensionality of both models to 25 (instead of 50 and 100 respectively), and added Gaussian noise to the previously deterministic SIR simulator. For a fair comparison, we use the same surrogate model for both Sourcerer and NEB.

We argue that Sourcerer's improved performance over NEB on these tasks is due to the different ways in which the two approaches use the surrogate model - NEB is sensitive to the exact likelihood, and thus may suffer from a lack of robustness to misspecification as in the case of a surrogate model. Our method, on the other hand, requires only samples, and thus the exact likelihoods produced by the surrogate model do not have a detrimental effect.

W1.1 "The Wasserstein distance is hard to interpret. It is claimed that a low value is observed, but that is not clear to me why this value can be considered a low value as there are no baselines to compare with": We agree with the reviewer that the numerical value of the Sliced-Wasserstein Distance (SWD) is hard to interpret in isolation. We had therefore computed the expected SWD between two sets of independently generated observations (simulated from the true source distribution). This minimum achievable distance is indicated by the black dotted lines in Figure 4 of our original submission. We found that for both models, the distance between simulations (from the estimated source) and observations is very close to the expected distance between simulations, and will emphasize this in the revision. Finally, we also provide the more interpretable C2ST accuracies for the SIR and Lotka-Volterra experiments in the Appendix (Fig. 8). In both cases, Sourcerer achieves a C2ST accuracy close to 50%.

We thank the reviewer for their feedback. We address the reviewer's concerns below.

W1: "The computational constraints on the model are heavy": The reviewer raises concerns about computational constraints on the model, which we clarify here for the two terms in the cost function:

First, the entropy/KL-divergence term (term 2) imposes negligible computational constraints on the model because it is computed on samples in parameter space, and not in data space. In particular, given samples $\theta_i$ from our source model, we only need to compute their nearest-neighbor distances $d_i$ to estimate the entropy (Eq. 8 of our submission). Thus, its computation is completely independent of the simulator or surrogate. Hence, this concern is not a limitation of the method. In addition, while other distances (such as Sliced-Wasserstein or MMD, as suggested by the reviewer) can be used to regularize the source in the second term, we emphasize that the choice of entropy (KL-divergence wrt. uniform distribution) stems from the original motivation of the work (i.e., the theoretical justification in Prop 2.1.)

Second, we agree with the reviewer that requiring the simulator to be differentiable for term 1 is a computational constraint on our model. However, we emphasize that our approach is an improvement over prior work (in particular, Vandegar et al. (2020)) in this regard. As the reviewer identified, the requirement to train surrogates for non-differentiable models is a broader limitation of the field. However, typical works in the field require training a surrogate that has a tractable and differentiable likelihood. We only require differentiability. We do not require tractable likelihoods, which reduces the computational constraints. This is a relative strength of our method relative to previous approaches, not a limitation.

[Vandegar et al.] - Neural Empirical Bayes: Source distribution estimation and its applications to simulation-based inference. In International Conference on Artificial Intelligence and Statistics, 2020.

Finally, in terms of computational cost, the choices we make for the two terms in the loss function are highly efficient. For the mismatch term (term 1), the Sliced-Wasserstein distance is linearithmic in the number of observations, and thus scales well to millions of observations. Furthermore, the distance computation of individual 1-D slices is embarrassingly parallelizable. Our use of Kozachenko-Leonenko estimators to estimate entropy (term 2) is also scalable, since computing the estimate requires solving the all-nearest-neighbors problem, which can also be solved in linearithmic time. We now perform a runtime comparison between NEB and Sourcerer for the benchmark tasks, and find that Sourcerer is significantly faster (Table R1).

W2: "The evaluation procedure or goal remains unclear to me": We apologize for not making our evaluation procedure sufficiently clear. For the benchmark tasks as well as the SIR and Lotka-Volterra experiment, we always evaluate the match between simulations and observations on a second, unseen ‘test’ set of observations that was generated independently using samples from the ground truth source distribution. Thus, we already do the evaluation in the way the reviewer asks, and just failed to point this out in the manuscript. We will clarify in the revised version.

More generally, our evaluation is always based on two criteria:

1. How well do the simulations (generated by the estimated source) match the set of observations? We measure this match using Classifier Two Sample Tests (C2ST) and the SWD.
2. How high is the entropy of the estimated source? We measure this by estimating the differential entropy with samples from the source using the Kozachenko-Leonenko estimator.

We benchmark our method against Neural Empirical Bayes (NEB), a state-of-the-art approach for source distribution estimation. For the benchmark tasks, we showed that Sourcerer obtains sources that are comparable to NEB in terms of simulation fidelity (point 1), while having a higher entropy (point 2). We now also perform an additional comparison on the high-dimensional data tasks of SIR and LV (Fig. 2), and find that our method also significantly outperforms NEB in terms of simulation fidelity (point 1). On SIR, Sourcerer achieves a C2ST accuracy (lower is better) of 55% between data and predictive simulations of learned source distribution, compared to 76% for NEB; similarly, on Lotka-Volterra, Sourcer achieves a C2ST of 56%, compared to 62% for NEB.

We thank the reviewer for their feedback and careful reading. We address the main concerns and questions raised by the reviewer below. Thank you for identifying typos, we will correct them.

W1: "My main concern is that the work presented looks like it is not a large or significant contribution beyond what is already present in the literature": We summarize the two main contributions of our work:

First, we provide an approach to resolve the ill-posedness of the source estimation problem by using the maximum entropy principle, and show that it leads to a unique solution. Empirically, we demonstrate that our approach can estimate high-entropy sources without sacrificing the quality of the match between simulations (from the estimated source) and observations. Second, our approach is fully sample-based. This provides significant computational advantages over previous methods, since our method can use differentiable simulators directly, without training surrogate models.

W1.1: "Detailed comparison to Neural Empirical Bayes": Our work makes conceptual, theoretical, and empirical contributions that go beyond Neural Empirical Bayes (NEB, Vandegar et al. (2020)):

First, a limitation of NEB - as directly acknowledged in their publication (Section 2, third paragraph) - is that the source distribution estimation problem is ill-posed. This ill-posedness is not addressed by the NEB method. In our work, we propose to use the maximum entropy principle as a way to resolve this ill-posedness, and make the theoretical contribution of showing that the maximum entropy source distribution is unique (Prop. 2.1). Furthermore, we demonstrate empirically (Table 1, Figures 3,4,5) that our method finds higher entropy source distributions than NEB, without any decline in the predictive performance of the estimated source distributions.

Second, our approach is entirely sample-based, i.e. it does not require access to likelihoods (or surrogate models with likelihoods). We show the advantages of our sample-based approach over NEB in the context of high-dimensional, differentiable simulators. In this setting, our method (unlike NEB) does not require the training of surrogate models. We now additionally show (Fig. R2) that our sample-based objective is more robust than NEB even in the case where a surrogate is available for higher-dimensional tasks. For the SIR task, our method achieves a C2ST accuracy (lower is better) of 55% between data and predictive simulations of the learned source distribution, compared to 76% for NEB; similarly, on Lotka-Volterra, our method achieves a C2ST of 56%, compared to 62% for NEB. Thus, our method is demonstrably superior to NEB on these high-dimensional tasks.

W1.2: "Comparison to additional works": The reviewer also asked us to contrast our approach with previous work.

Regarding Knoblauch et al. (2022) [1]: The idea of using other distance measures between distributions is the key approach of Generalized Bayesian Inference (GBI) methods, as discussed in our related work and in [1]. Our method shares similarities with GBI in that it uses non-likelihood-based objective functions. However, we emphasize that our work is concerned with source distribution estimation. This is a different problem than inference, in that we seek to identify a distribution of parameters that is consistent with a population of independent observations (many to many), rather than to explain a single or multiple observations from a single parameter (i.e., posterior inference, one to one or many to one). These are not technical differences, but rather that these methods simply address different goals.

Regarding Dyer et al. (2024) [2]: This problem setting is different from ours. In this paper, the authors seek to find a distribution that produces simulations “matching closely a target, user-specified scenario”. This is different from the source distribution estimation problem, where we try to find a distribution that reproduces observed data. In addition, the entropy regularization has a very different interpretation. In Dyer et al. (2024), the entropy regularization “exhibits the trade-off between (a) the diversity of synthesized populations and scenarios and (b) the ability to identify the most extreme manifestations of the target scenario”. In our work, we show theoretically (Prop. 2.1) and empirically (Table 1, Figures 3,4,5) that no such trade-off is necessary for source distribution estimation.

In summary, we believe that our work is a significant step upon previous work and hope that our response addresses the reviewer’s concern that our contribution is limited to the use of a sample-based loss.

Q1: "Is there anything that we lose by using something like the Sliced Wasserstein Distance instead of a likelihood-based discrepancy?": We thank the reviewer for their question. We believe that the main drawback of using sample-based distances is in the low sample limit, where the number of observations $N$ in the true data is very small. Likelihood-based objectives may perform better in this setting, as the value of the likelihood itself contains more information than the sample alone.

We also evaluated Sourcerer’s performance in a low sample count regime ($N=100$ instead of $N=10000$) in a new experiment for the Two Moons task (Fig. R1a). We observe that our method, using the sample-based Sliced-Wasserstein distance, is still able to recover good sources in this case.

Q2: "Perhaps something additional to discuss would be limitations associated with simulators that involve discrete randomness": We thank the reviewer for pointing out this limitation. While we believe that an investigation of discrete simulators is outside the scope of this work, we agree that it merits acknowledgement in the limitations section.

We thank the reviewer for their positive assessment of our work. We address the reviewer’s questions below.

W1: "It would be interesting to see a more thorough study for models with even higher dimensionality": We thank the reviewer for their suggestion to investigate the performance of our approach in high-dimensional parameter spaces. To do so, we perform an additional experiment on the Gaussian Mixture benchmark task. The model is identical to the previous Gaussian Mixture model described in our benchmark task (Appendix 1.4), but now the dimension has been increased to 25. We again find sources that reproduce the data well, and that when we apply entropy regularization, we obtain higher entropy source distributions without any degradation to the pushforward accuracy (Fig. R1b). We also observe that the threshold value of $\lambda$, at which the pushforward C2ST degrades is now lower than in the 2-dimensional case. We believe that this is due to the entropy of the higher dimensional sources varying on a larger scale, thus restricting our method to smaller values of $\lambda$ to avoid the entropy term dominating the mismatch term.

Q1: "Maybe some statements that make the maximum entropy principle stand out as a sort of gold standard could be toned down a bit?": We agree with the reviewer that statements about maximum entropy being the "gold standard" need to be toned down, as this was not our intention. To clarify, we propose maximum entropy as one possible way to solve the ill-posedness of the source distribution estimation problem, which has a number of advantages such as uniqueness, and coverage of all feasible parameters for downstream inference tasks. These are not "intrinsically" optimal (and there might be specific cases where other regularizers would be more appropriate), and we will clarify the relevant statements in the revised version.

Q2: "Would it be possible for the authors to show what could happen with the choice of other or wrong distance types?": We thank the reviewer for their suggestion to use distance metrics other than the Sliced-Wasserstein distance (SWD). Therefore, we repeat our experiments on the benchmark tasks, replacing the SWD with MMD using an RBF kernel and the median distance heuristic for selecting the kernel length scale (Fig. R3). We observe that the results are comparable in terms to the SWD results presented in our work.

The main considerations in choosing the SWD in our work were computational. We expect that any sample-based, differentiable, and computationally efficient distance metric to measure the mismatch in our objective function (Eqs. 3, 4) will provide reasonable performance. Our initial choice of using SWD was motivated by its simplicity and scalability - we did not need to choose kernels or other hyperparameters (except for the number of slices).

Q3: "An interesting case the authors could discuss is when multiple losses may be required for different subsets of parameters?": We strongly agree with the reviewer that there are applications where a general mismatch function such as the SWD might be insufficient to measure the distance between distributions, and some combination of distance functions would be required for different subsets of the data. For example, for simulators that produce multimodal outputs, the SWD may not be an appropriate metric. Instead, we could consider designing a distance custom distance metric that combines the distances over the different modalities of the data. We believe that this is outside the scope of the work, but believe that this is an important point to add to the discussion in our submission.

We thank the reviewer for their thorough feedback. The reviewer raises some questions and concerns about our approach, which we address below.

W1: "The proposed method only applies to differential simulators": In the case where the simulator is neither differentiable nor provides explicit likelihood evaluations, our method can still be used, but then requires training a surrogate model. On our benchmark tasks (Table 1), we have shown that using a surrogate instead of the simulator led to comparable results in terms of simulation fidelity to data, and entropy. The results for the realistic Hodgkin-Huxley task (Fig. 5) were also obtained using a surrogate model. We emphasize that our method has less requirements than previous approaches: We do not require that the simulator have an explicit and differentiable likelihood, but only that the simulator be differentiable. Thus, the set of simulators to which we can apply our method without training a surrogate is larger.

W2: "I also wish there is more theoretical results in terms of the statistical consistency and error bounds of using the sliced Wasserstein distance": The Sliced Wasserstein distance (SWD) is well established and its theoretical properties are actively studied. In particular, the sample complexity of the convergence rate has been studied in several existing works. For example, Nadjahi et al. (2020) showed that the finite sample estimate of the SWD between two distributions converges with rate $\sqrt{N}$ in the number of samples $N$. Furthermore, Nadjahi et al. (2019) showed that generative models trained to minimize the SWD using finite sample size $N$ also converge to the true optimizer (in distribution) at rate $\sqrt{N}$. We will add a summary of these properties in our revisions.

We agree with the reviewer that our benchmark tasks would benefit from exploring how our method performs in the low data limit. Therefore, we provide an additional empirical result by repeating the Two Moons benchmark task with a smaller observed dataset of only 100 samples (Fig. R1a). We observe similar behavior to the baseline case presented in our work for 10000 observations.

[Nadjahi et al. (2020)] - Statistical and Topological Properties of Sliced Probability Divergences. NeurIPS, 2020.
[Nadjahi et al. (2019)] - Asymptotic Guarantees for Learning Generative Models with the Sliced-Wasserstein Distance. NeurIPS, 2019.

W3: "I'm not sure if the maximum entropy approach is aligned with the intuition": We apologize for not making the intuition for using the maximum entropy principle sufficiently clear. We provide further clarification here and will reflect this in our revisions.

The maximum entropy principle does not conflict with the case where “the source distribution is highly concentrated in an area with high probability mass”. If parameters outside this area are not consistent with the data observed, our method will result only in distributions that are concentrated in this area. We observe this empirically in the deterministic SIR experiment (Fig. 4), where the source distribution is concentrated and regularization does not lead to higher entropy sources.

Instead, our use of the maximum entropy principle can be viewed as resolving the non-uniqueness of the source distribution estimation problem. According to the maximum entropy principle, given a choice between two source distributions that are consistent with the data, one should conservatively choose the source distribution with the higher entropy (which is more dispersed). This is often desirable when source estimation is used to learn prior distributions, as it ensures that a wide range of possible parameters is considered.

Q1: "What is the motivation of using the sliced Wasserstein distance?": We thank the reviewer for their question. We agree that other choices of sample-based, differentiable distances are possible. We now perform a new experiment on the benchmark tasks replacing SWD with MMD with an RBF kernel and the median distance heuristic for selecting the kernel length scale (Fig. R3). We observe that the results are comparable to the SWD results presented in our work in terms of simulation fidelity to data and entropy.

The constraints on the distance function D are computational; D should be differentiable, fast to compute, and not require probability densities. Our new empirical results show that indeed other distances that satisfy these constraints lead to good performance of our method. Our initial choice of using SWD was motivated by its simplicity and scalability - we did not need to choose kernels or other hyperparameters (except for the number of slices).

Q2: "Is there any principle of how to choose the hyperparameter?": We empirically find that a small value of $\lambda$ (e.g. $\lambda=0.015$), together with our linear decay schedule, is sufficient to obtain substantially higher entropy source distributions in all our experiments. Thus, starting with a small $\lambda$ is likely to be sufficient. In addition, a run without any regularization can be performed to verify that the regularization does not negatively affect the quality of the estimated source. We agree with the reviewer that more robust methods for choosing $\lambda$ are an interesting question for future research, and will reflect this in our revisions.

Q3: "Is there any empirical result of the convergence speeds for the proposed method vs the benchmark NEB?": We thank the reviewer for their question, and agree that our work can benefit from explicitly comparing the empirical convergence speed of Sourcerer to the baseline. We provide an empirical measurement of the convergence times of our method (with and without regularization) as compared to NEB, measured on the benchmark tasks (Table R1). The source model architectures $q_\phi$ are the same for all methods, as reported in Table 1 of our original submission. Sourcerer is faster than NEB on all benchmark tasks.

Dear all,

Thank you for your efforts so far in reviewing the paper and in answering reviewers' questions.

The evaluation is currently borderline, with an average score of 5.8 and a spread from 4 to 7.

The authors-reviewers discussion will close today on August 13. Please take this last opportunity to discuss the paper and the reviews and to reach a consensus, or at least to gather all the necessary information to make a fair and informed decision (in favor of acceptance or rejection).

Thank you for your attention and cooperation. The AC