Paper: Estimating Unknown Population Sizes Using Hypergeometric Maximum Likelihood (ICLR-2024)

Reviewer comment (Theory):
I have some concerns about the motivation/applicability of the paper. In the first part, the authors contrast their approach to capture-recapture, and say that capture re-capture is not good because one must be able to sample (to tag) and then resample from the same population.  Yet, in their own setup (e.g., Equation (3)) the authors assume that one can repeatedly and independently sample from the same population.  I find this setup interesting from a statistics perspective, but unrealistic as an experimental design.  If I understand correctly, it would correspond to sampling without replacement from a population, returning all of those samples to the population, and then resampling without replacement and so on.  In some sense, this is essentially what the authors claimed was problematic about capture-recapture, but they are repeating it $T$ times instead of twice.  I would be happy to change my mind about this being problematic if the authors could provide some compelling, real life examples of datasets of practical interest with this experimental design.

Author response:
One example of a real-life application that directly matches this experimental design is a variation of the standard capture-recapture problem of counting fish in a lake. In the original problem statement, we wish to estimate the number of fish of a certain species in a lake. To do so, we catch a sample of fish, tag the fish of the species of interest, and return all caught fish to the lake. We then catch a second sample and count how many tagged fish have been recaptured. However, consider a variation where we are instead counting fish in a flowing river. Here it would be impossible to recapture any tagged fish since they are swept downriver. But if we can assume that the instantaneous distribution of fish at the sample point is relatively constant over the course of the experiment, then we can perform multiple resamplings, counting the fish of the species of interest each time, and obtain an estimate using our proposed method.

As a second hypothetical scenario, in the context of quality control we may wish to estimate the number of defects in a batched production line, where it is known that the same machines are used to produce each batch. Capture-recapture would require sampling a subset of items from a batch, marking those with defects, returning them, sampling another random subset from the same batch, and counting how many of the marked defective items reappear. This would require stopping the production line, or removing the entire batch until the experiment is complete. In contrast, here we propose sampling from each batch only once, without needing to mark and resample items twice from the same batch, allowing the production line to continue operating.

We believe that the key difference between the capture-recapture method and our repeated sampling approach is that the former requires resampling of the exact same population, whereas we only require resampling from the same distribution. Therefore in any experiment where we can assume the distribution remains unchanged across multiple instances, it would be possible to resample the same distribution without resampling the exact same population.

Additionally, capture-recapture cannot be used if carrying out the experiment changes the underlying distribution, for example if the act of sampling from a batch may cause additional defects. In our later experiment this is important as measuring gene transcripts requires destroying the cell they are in.

Having said that, we agree that the classical statistics regime, where there is only a single count distribution to be estimated, is not the most valuable setting to apply our approach. We mainly intend this initial set of simulations to demonstrate that it is theoretically possible to perform inference on the multivariate hypergeometric distribution with unknown population sizes, as we do not believe this has been demonstrated previously. We intend to set the scene for the later sections, and we believe that the primary method of practical value is in the case of count distributions conditional on a continuous latent variable. It would be impossible to apply capture-recapture to the single-cell genomics problem without substantial breakthroughs on the experimental side.

Reviewer comment (Theory):
At the bottom of Figure 3 it is claimed that Equation (4) is concave.  I plotted Equation (4) for a particular setup and it is _not_ concave.  One can also see this from Figure 1 where if one were to draw a well place diagonal line through the plot, one could obtain a linear slice where the negative log-likelihood first decreases, then increases, then decreases again, which is impossible for a convex function.  The error in reasoning comes from the fact that Equation (4) is the sum of both $\log \Gamma(x)$ and $-\log \Gamma(x)$ functions.  The former is strictly concave, but the latter is strictly convex.  The sum of convex and concave functions does not need to be either convex or concave.

Author response:
Thank you for pointing this out, we agree with you and so have removed this statement. Please note that the relevance of our results and conclusions remains intact. While the log-likelihood is not concave, it appears to be concave along the line of correct N_1/N_2 ratio, however have not yet been able to show this rigorously. This is important as the optimization path is mainly along this vector. We hope that deeper theoretical analysis of his likelihood will be undertaken in the future, and have added this to the Future Work section.

Reviewer comment (Evaluation):
Figures 2 and 3 should contain error bars and/or multiple replicates.  The fact that the error fluctuates across the number of trials (as opposed to monotonically decreasing) suggests that there is some substantial noise.  This noise makes it very difficult to assess the claim that increasing $K$ makes the problem easier.

Author response:
We are grateful for this suggestion. We have added confidence intervals to Figures 2 and 3 by repeating each trial with multiple random seeds. We’ve used 25/75% intervals for visual clarity, and the same figure with 5/95% intervals has been added to the appendix.

Reviewer comment (Evaluation):
My intuition is that the good performance in Figures 2 and 3 comes from the large number of trials.  E.g., even at thousands of trials, there is still substantial error.  In particular, if one only has a single trial, then I think that the MLE for the population counts should just be the observed counts, but the magnitude of this is driven entirely by the sample size and has nothing to do with the population size.  As a result, I don't think this is a sensible approach when there is only a single trial.  I suspect that this becomes problematic in the VAE setting -- since the likelihood is not particularly informative for any given point the prior will be incredibly important.

Author response:
We agree that this approach only works when there are many observations available, and not in the case of few trials. In our experiments we have observed that with few observations the MLE obtains the correct relative fraction of counts (as would be obtained using a multinomial likelihood), but that more observations are required to obtain the exact counts. The log-likelihood is quite flat, and therefore the minimum is sensitive when there are few observations.

We have not investigated if the estimate converges more quickly if a prior is used, and have focused on applications where there are enough trials to obtain a good estimate. It would be an interesting extension to see if fewer trials are required with a well-chosen prior.

Reviewer comment (Writing):
In Equation (1), I believe that $n_2$ should be $c_2$.

Author response:
Yes you are correct, we have fixed the typo.

Reviewer comment (Writing):
The notation in Equation (8) is a bit sloppy.  I understand the point the authors are making, but it seems sloppy to define $\hat{N}_i$ as a function of itself (i.e., $\hat{N}_i$ appears on both sides of the definition).

Author response:
That is a good point, we have changed the equation from an equality to an update.

Reviewer comment (Theory):
I believe that the continuous version of the likelihood (i.e., Equation (4) but using equation (6) for the binomial coefficients) does not define a proper probability distribution for non-integer $N_k$.  Is this problematic?

Author response:
You are correct, the continuous relaxation is no longer properly normalized for non-integer N_k. We do not believe this is a problem if we are using it as an optimization objective, as the parameters that produce the minimum remain unchanged.

Reviewer comment (Theory):
Is there a reason for including the violation penalty $C_\text{viol}$ as opposed to just performing bounded optimization?  The lower bound on $\hat{N}_i$ can be determined easily from the data, so it would not be difficult to enforce the bound during optimization.  Furthermore, the constraint set is quite simple (and convex), so something like projected gradient descent would be easy to perform.

Author response:
We did not pursue a bounded optimization approach as we were not able to obtain a lower bound on N_i when there are multiple distributions to estimate. We are in the unsupervised regime and do not know which observations belong to which distribution, and hence what lower bound is appropriate for each observation. However, we expect that there is room for more sophisticated ways of discouraging count estimate violations, and we have added bounded optimization as a possible next step to consider.

Reviewer comment (Writing):
The authors do not provide a clear motivation for the second part of the work which focuses on the variational autoencoder framework. Although the initial experiments focus on the recovery of the category sizes of a single multivariate hypergeometric distribution, they later explore a scenario where they have a mixture of distributions. The estimated counts are per observation, what is the ultimate goal in this case?

Author response:
Thank you for this feedback, we have added additional description of the motivation for the need to model a mixture of distributions to the introduction and Section 5.

In machine learning, count data is often binarized so that it can be modeled using a probability distribution over categories, such as by using a neural network with a softmax output layer. However, we believe that disregarding the actual counts leads to loss of information that may be valuable in certain applications. Therefore, departing from previous work, the ultimate goal of our approach is to directly estimate the observed counts without forced binarization, thus preserving this information.

Reviewer comment (Theory):
It seems to be a mistake in the likelihood of equations 3 and 4 compared to that at the end of page 3. What is the optimization problem that you solve?

Author response:
Thank you for pointing this out, we have updated the notation to make it clear which parameters are being optimized, and to distinguish between the parameters and the observations.

Reviewer comment (Writing):
Minor: N2 should be 30 in section 4.

Author response:
You are correct, we have fixed that typo - thank you.

Reviewer comment (Experiments):
I believe the method is also applicable to the non-central Hypergeometric as in the cited Sutter et al. (2020) paper and it would have been interesting to apply the proposed method to scenarios that are similar to those described in that work.

Author response:
We were not able to compare against the experiments in Sutter et al. (2022) because the starting assumptions about which parameters in the model are known/unknown are quite different. These authors address the non-central hypergeometric distribution, which is an extension of the hypergeometric distribution where the items belonging to each category have different probabilities of being sampled. Using the urn example, this would occur if different colored balls had different sizes, leading to some colors being more or less likely to be drawn from the urn, all else being equal. The standard hypergeometric distribution can be seen as a special case of the non-central distribution with all category weights equal. Therefore they learn the category weights while assuming all category sizes are known, whereas we learn the category sizes, while fixing the category weights to be equal. The experiments they perform primarily evaluate these learned category weights, which are constant in our method.

Reviewer comment (Experiments):
It feels like there is a bit of a gap in the experiments, there is a jump from using 3 categories to more than 1000 categories. What about more intermediate cases, say with 100 categories? How does the proposed method work in this case with respect to n_max as well as the number of samples?

Author response:
During initial empirical exploration we found that the parameter inference worked well with an intermediate number of categories. We chose to use 1000 features for the experiment in Section 5 because in Section 6 we use 10,000 features, so across Sections 4-6 we span five orders of magnitude in terms of number of parameters to estimate. However, we agree that including an experiment comparing the effect of number of categories for Section 5 would strengthen the argument, and we are working to generate these results to add to the Appendix if time permits. We have added additional results for different f_max for the case of K=3 to the Appendix.

We have performed more targeted investigation into the effects of n_max and f_max. One factor hindering this is the effect of the level of correlation between features. Our experiments on empirically simulated data, where we assumed no correlation structure between features, suggest that estimation performance degrades below an f_max (max counts sampled vs total number of counts) of roughly 0.25, as shown in Figure 2. However, we found that for the single-cell genomics data described in Section 6, the method can still perform well with a lower f_max. We believe this is due to the presence of mutual correlation among count features, in this case because genes are often turned on in groups, and act in biological pathways. Because in our approach we do not factorize the hypergeometric distribution, its parameters are inferred jointly, and therefore the joint correlation between features is beneficial to the estimation process. We expect this correlation is common in most real dataset, but its effect remains to be confirmed.

Reviewer comment (Novelty):
Has a similar relaxation (replacing a Binomial coefficient or factorials with a Gamma function) been used elsewhere that you are aware of in estimation contexts?

Author response:
The same relaxation that we use in our approach was also used for the non-central hypergeometric distribution by Sutter et al. (2022). The Gamma relaxation has been used to extend the binomial coefficient to negative values [4]. We are also aware that the Gamma function is typically used to compute binomial coefficients efficiently, such as in scipy. All sources we have found directly state that this switch from discrete to continuous is valid since x! = Gamma(x+1), by definition of the Gamma function.

Reviewer comment (Writing):
There is a typo in equation (1), you should have c_2 and not n_2?

Author response:
Yes, you are correct, we have fixed that typo - thank you.

Reviewer comment (Theory):
In equation (7), shouldn't you define c_i to be the largest category count over the T samples, as that gives you the lower bound on the N_k?

Author response:
You raise a good point. In the case of a single underlying distribution your suggestion would make more sense, but in the later scenarios we address, where there are multiple distributions to be learned, taking the max would not work as the samples are unlabelled (we don’t know which distribution they are drawn from), so we cannot compute the max within each category. We have added a note to make this distinction.

Reviewer comment (Theory):
There are no theoretical justifications on how far or close the proposed relaxations could be from the true count-based discrete distributions.

Author response:
> Can there be theoretical arguments made on how bad/good the proposed relaxation would be?

At a minimum we can say that the continuous relaxation of the hypergeometric distribution will be exact for integer values of x, given that x! = Gamma(x+1).

In the only recent use of a continuous relaxation of the hypergeometric distribution that we could find, Sutter et al. (2022) use a Kolmogorov-Smirnov test to show that their proposed continuous relaxation of the non-central hypergeometric distribution is statistically indistinguishable in sample distribution from the discrete form. The authors similarly replaced the factorial with the Gamma function, although they factor the joint multivariate distribution into bivariate conditional distributions (see Section 4 of their paper). Because we do not need to factorize the joint distribution, our proposed relaxation should perform at least as well, and therefore we can conclude that our proposed relaxation is also identical in distribution to the discrete hypergeometric distribution.

Reviewer comment (Theory):
Can there be theoretical arguments made on how bad/good the proposed relaxation would be?

Author response:
> Can there be theoretical arguments made on how bad/good the proposed relaxation would be?

At a minimum we can say that the continuous relaxation of the hypergeometric distribution will be exact for integer values of x, given that x! = Gamma(x+1).

In the only recent use of a continuous relaxation of the hypergeometric distribution that we could find, Sutter et al. (2022) use a Kolmogorov-Smirnov test to show that their proposed continuous relaxation of the non-central hypergeometric distribution is statistically indistinguishable in sample distribution from the discrete form. The authors similarly replaced the factorial with the Gamma function, although they factor the joint multivariate distribution into bivariate conditional distributions (see Section 4 of their paper). Because we do not need to factorize the joint distribution, our proposed relaxation should perform at least as well, and therefore we can conclude that our proposed relaxation is also identical in distribution to the discrete hypergeometric distribution.

Reviewer comment (Experiments):
Is there a relationship of the proposed method to other hyper-parameter optimization methods? Why were no other baselines included that could possibly be used to solve the transcript problem? There is a huge literature on trying to estimate the hyper-parameters of count distributions including the buffet processes.

Author response:
> Is there a relationship of the proposed method to other hyper-parameter optimization methods ? Why were no other baselines included that could possibly be used to solve the transcript problem? There is a huge literature on trying to estimate the hyper-parameters of count distributions including the buffet processes.

We agree that there is a large literature on trying to estimate the hyperparameters of count distributions, and recently transcript count distributions in particular. However, the majority of these previous approaches attempt to learn the distribution of the observed (undersampled) counts, whereas here we aim to obtain the true distribution of counts before sampling. In other words, most methods attempt to estimate the parameters of distributions that account for the sparsity of the data (such as with a (zero inflated) binomial model), whereas we assume the sparsity is due to undersampling and try to estimate the parameters of the original distribution prior to sampling. Because we attempt to directly model the true underlying distribution, we are not able to provide a comparison.

As far as we are aware, no existing method has attempted to estimate this true underlying count distribution in the population under the assumption of sampling without replacement, which is why no baseline could be provided.

Reviewer comment (Experiments):
How exactly is the fact about with/without replacement helpful for the genome problem? The dataset has counts high enough that even if we use the methods that sample with replacement, they should possibly give reasonable results. Are these approaches known? If not, would it not make sense to compare against such baselines?

Author response:
> How exactly is the fact about with/without replacement helpful for the genome problem? The dataset has counts high enough that even if we use the methods that sample with replacement, they should possibly give reasonable results. Are these approaches known? If not, would it not make sense to compare against such baselines ?

We agree that the assumption that transcript counts are sampled with replacement (eg. a multinomial likelihood) may be reasonable for this problem if we are only interested in learning a useful latent space. However, there are two limitations in this scenario. The first and most important is that the assumption of sampling with replacement would only permit the model to learn the relative abundance of different feature counts (genes) through a probability distribution over the features. However, we believe that knowing the true count of RNA in a cell is valuable for downstream bioinformatic analyses. At a minimum, it significantly reduces sampling noise, which would improve the performance of further in-depth analyses. In the other applications we mention, such as recommender systems, we similarly argue that in some use cases the true number of counts is valuable. Second, gene counts have a high dynamic range, with some genes have many counts and some having very few, spanning 3-4 orders of magnitude. Therefore, while the approximation resulting from the assumption of sampling with replacement may work well for modeling genes with high counts, it introduces more significant errors for genes with low counts, where the probability of sampling a second transcript changes significantly after a first has been sampled. In fact, it is known that most signal in single-cell data is driven by the high abundance genes [3], which may be due to this bias.

Reviewer comment (Theory):
Why is N>>1 required for the continuous relaxation to be reasonable?

Author response:
> Why is N>>1 required for the continuous relaxation to be reasonable ?

We added this caveat specifically when thinking about the maximum percent error introduced by approximating an integer with a whole number, but in fact N>>1 is not required. We have removed that sentence for the sake of clarity.