S CALING D ENSITIES FOR I MPROVED D ENSITY R ATIO E STIMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Estimating the discrepancy between two densities ( p and q ) is central to machine learning.
Most frequently used methods for the quantiﬁcation of this discrepancy capture it as a function of the ratio of the densities p/q .
In practice, closed-form expressions for these densities or their ratio are rarely available.
As such, estimating density ratios accurately using only samples from p and q is of high signiﬁcance and has led to a ﬂurry of recent work in this direction.
Among these, binary classiﬁcation based density ratio estimators have shown great promise and have been extremely successful in specialized domains.
However, estimating the density ratio using a binary classiﬁer, when the samples from the densities are well separated , remains challenging.
In this work, we ﬁrst show that the state-of-the-art solutions for such well separated cases have limited applicability, may suffer from theoretical inconsistencies or lack formal guarantees and therefore perform poorly in the general case.
We then present an alternative framework for density ratio estimation that is motivated by the scaled-Bregman divergence.
Our proposal is to scale the densities p and q by another density m and estimate log p/q as log p/m − log q/m .
We show that if the scaling measures are constructed such that they overlap with p and q , then a single multi-class logistic regression can be trained to accurately recover p/m and q/m on samples from p, q and m .
We formally justify our method with the scaled-Bregman theorem and show that it does not suffer from the issues that plague the existing solutions.
We provide a large battery of empirical evaluations of our method with both synthetic and real datasets on the tasks of density ratio estimation, mutual information estimation, and representation learning.
Finally, we demonstrate that our method can be applied to improve non-classiﬁcation based model-free density ratio estimators as well.
I NTRODUCTION
Quantiﬁcation of discrepancy between two distributions underpins the foundation of a large number of machine learning techniques.
Of especial prominence are distribution discrepancy measures known as fdivergences (Csisz´ar, 1964), which are deﬁned as expectations of the functions of the ratio of the two densities.
As such, density ratio estimation is often a central task in deep generative modeling, mutual information and divergence estimation, representation learning, etc.
(Sugiyama et al., 2012; Gutmann & Hyv¨arinen, 2010; Goodfellow et al., 2014;
Nowozin et al., 2016; Srivastava et al., 2017; Belghazi et al., 2018;
Oord et al., 2018; Srivastava et al., 2020).
However, estimating density ratio by modeling each of the densities is challenging for most problems of interest as modeling high dimensional densities is a harder problem than estimating their ratios Sugiyama et al.
Therefore, in practice, estimators are employed that only require samples from the pair of densities in order to estimate their ratio without explicitly modeling the individual densities.
One of the most commonly used density ratio estimators (DRE) is a binary classiﬁer based DRE (BC-DRE) which can be formally shown to estimate the ground truth density ratio when trained to correctly discriminate be-
BC-DREs have been tremendously successful in problems involving minimization of the discrepancy between the data and the model distributions even in high-dimensional settings (Nowozin et al., 2016;
Radford et al., 2015).
However, they do not fare as well when applied to the task of estimating the discrepancy between two distributions that are far apart or easily separable otherwise.
This issue has been characterized recently as the density-chasm problem by Rhodes et al.
The discrepancy between two distributions can manifest itself in many different ways.
Here, we roughly group these ways into two classes, First Order Discrepancy (FOD) and Higher Order Discrepancy (HOD) .
Two distributions can easily be separated when the ﬁrst order statistics, such as their means, are far apart from each other.
We refer to this type of discrepancy between p and q as FOD.
Similarly, if the discrepancy between p and q (with no FOD) originates due to the difference in higher order statistics we refer to this as HOD.
For example, consider Figure 1 where in all the three cases, the
KL [ p (cid:107) q ] = 50 , but in (a) it originates from FOD since N ( µ 1 = − 5 , σ 1 = 1) , N ( µ 1 = 5 , σ 1 = 1) , in (b) it originates from HOD since N ( µ 1 = 0 , σ 1 = 1) , N ( µ 2 = 0 , σ 1 = 0 . 097) , and in (c) it originates from FOD and HOD since N ( µ 1 = − 1 . 5 , σ 1 = 0 . 5) , N ( µ 2 = 1 , σ 1 = 0 . 25) .
In a recent work, Choi et al.
(2021) attempt to address the density-chasm problem.
They considered distributions p and q that are clearly separated and proposed that in such cases density ratios can be better estimated if the mixture of the two densities is mapped onto a unit-sphere using a ﬂow-based map (Rezende & Mohamed, 2015).
Using a bijective map ensures that the density ratio between the original distributions is maintained post mapping.
While this method (F-DRE) may resolve the density-chasm problem due to FOD, we found that it struggles when the densities are separated due to HOD.
On the other hand, Rhodes et al.
(2020) proposed to estimate the density ratio between p and q as pp 1 × p 1 p 2 × · · · × p K q .
This telescoping is achieved by training K BC-DREs where the distributions p 1 , . .
. , p k are designed to be closely-packed such that all of the K ratios can be easily estimated using BC-DREs without suffering from the density-chasm problem.
They call this estimator, telescoping density ratio estimator or TRE.
Empirically, we found that while TRE performs relatively better than BC-DRE and F-DRE in cases with density-chasm problem due to HOD, it cannot deal with the density-chasm problem originating due to FOD.
Intuitively, this happens because TRE trains K DREs, each of which is trained on K different pairs of distributions.
This leads to the distribution shift issue at test time, especially when p and q have high FOD.
Further, care needs to be taken when designing intermediate distributions p 1 , . .
, p K to ensure that each of the ratios are well deﬁned.
In fact, the construction schemes introduced in the original work may render TRE theoretically ill-deﬁned if both p and q do not have full support (Appendix A).
We will show that neither of these state-of-the-art methods are applicable when the density ratio estimation problem involves both FOD and HOD.
Unfortunately, most problems in practice tend to have both types of discrepancy.
This leaves a clear technical gap in the task of density ratio estimation.
Hence, in this work, we present S CALED D ENSITY R ATIO E STIMATOR ( S DRE), a novel method for estimating density ratio that:
Theoretically well-deﬁned and does not suffer from distribution shift.
Resolves both FOD and HOD types of density-chasm issues.
Does not require training additional FLOW-based models or BC-DREs and therefore is computationally more efﬁcient.
N OTATION AND B ACKGROUND
Bregman distance between two d -dimensional vectors a and b is deﬁned as B φ ( a, b ) =
φ ( a ) − φ ( b ) − ∇ φ ( b )( a − b ) .
Here φ : R d (cid:55)→ R is a convex function.
Following the notation in Stummer & Vajda (2012) the scaled-Bregman Divergence (sBD) between probability measures P and Q scaled by an arbitrary measure M can be deﬁned as,
[Equation]
Here, f : X → R is a convex function, f (cid:48) is its right derivative.
Further, all the measures P, Q and M are deﬁned on the measurable space ( X , A ) and are absolutely continuous with respect to a σ -ﬁnite measure λ , therefore admitting the following densities,
[Equation]
It is easy to show that for f ( t ) =
t log( t ) , sBD reduces to KL divergence and is independent of the scaling measure M and λ ,
[Equation]
We deﬁne
[Equation]
as the density ratios between p & m , q & m and p & q respectively.
Given samples { x ip } Ni =1 and { x jq } Mj =1 from distributions p and q respectively and assuming p << q , the goal of DRE ˆ r p/q is to estimate r p/q .
2. B INARY C LASSIFICATION TO DRE USING BD
Proposition 3 (Menon & Ong, 2016) formally establishes the link between binary classiﬁcation and density ratio estimation using sBD.
We re-state it here.
Proposition 1 (Proposition 3 (Menon & Ong, 2016) rephrased)
Given a class-probability estimator ˆ η : X (cid:55)→ [0 , 1] such that the density ratio estimator, ˆ r p/q ( x ) =
ˆ η 1 − ˆ η .
Then for any convex differentiable f : [0 , 1] (cid:55)→ R ,
[Equation]
[Equation] where D f is Bregman Divergence (BD), f † : [0 , ∞ ) (cid:55)→ R | f † ( x ) =
(1 + x ) f ( x 1+ x ) , D f † is sBD and p marginal is the marginal distribution of X .
2. D ENSITY C HASM
Kato & Teshima (2021) note that most empirical estimators of density ratio, especially those implemented using a deep neural network tend to over-ﬁt the loss function in some way or the other.
This phenomenon was recently characterized by Rhodes et al.
(2020) for the BC-DRE as the density chasm problem. They found that the BC-DRE tends to substantially underestimate the ground-truth KL-divergence (as a function of the estimated ratio) between p and q when these distributions are signiﬁcantly separated from each other.
We conjecture that this underestimation occurs because in the case of ﬁnite training data, when the two sets of samples are easily separable, the estimated ratios do not need to be as large as they theoretically had to be to achieve perfect classiﬁcation performance.
In such a case, many different decision boundaries achieve equally good performance.
R ELATED W ORK
Our work is most closely related to and improves upon the shortcomings of the recently proposed
TRE method.
Therefore, in order to better understand these shortcomings, we begin by analyzing
TRE under Proposition 3 of Menon & Ong (2016) as described in equation 5 above.
TRE uses a two step, divide-and-conquer strategy to resolve the density chasm problem.
In the ﬁrst step, they construct K waymark distributions p 1 , . . .
, p K by gradually transporting samples from p towards samples from q .
Then, they train K BC-DRE models, one for each consecutive pairs of distributions.
This allows for estimating the ratio r p/q as the product of the K BC-DREs, [Figure] pq = pp 1 × · · · × p K q .
They introduced two schemes for creating waymark distributions that make sure that consecutive pairs of distributions are packed close-enough to ensure that none of the K
BC-DREs suffers from the density-chasm issue.
Hence, TRE attempts to resolve the density-chasm issue by replacing the ratio between p and q with a bridge of K intermediate density ratios that by design (of the waymark distribution) do not suffer from density-chasm problem.
Application of Proposition 3 to TRE shows that each of the K individual binary classiﬁers, upon training lead to K DREs that minimize the sBD to their corresponding true density ratios on the samples from their corresponding denominator densities.
This however, does not guarantee that a product of such individual DREs minimizes the sBD to r p/q on samples from q (or all K intermediate distributions for K > 1 ).
This issue is acknowledged by the authors who note that TRE’s performance suffers signiﬁcantly from the distribution shift issue especially when the distributions p and q have high FOD.
An example of this can be seen in Figure 2a where we train TRE to estimate the density ratio between p = N ( − 1 , 0 . 1) and q = N (+1 , 0 . 2) using 3 intermediate distributions p 1 , p 2 , p 3 constructed using the linear-mixing construction described in Section 4.2.
In the ﬁrst row we evaluate the ratios p/p 1 , p 1 /p 2 , p 2 /p 3 and p 3 /q on samples from their respective denominator densities and in the second row we evaluate log pq on samples from all the K + 2 distributions.
As per equation 5, individual DREs are in fact able to estimate their corresponding density ratios accurately on the samples from their corresponding denominator densities.
But, their combination does not lead to an accurate estimation of p/q for any of the K + 2 distributions except p 2 , making TRE not useful for estimating discrepancy measures such as KL divergence as the model is not accurate on samples from p .
Another interesting related work comes from Choi et al.
(F-DRE), who propose to train BCDRE in a learned feature space where the two distributions p and q are mapped onto a unit sphere.
To this end, they employ a FLOW-based model on the input and train it to map the mixture of the samples from the two distributions to a standard Gaussian.
It is easy to show that any bijective map will preserve the original density ratio r p/q in the feature space as the Jacobian correction term simply cancels out.
However, such a method cannot bring the distributions closer along both FOD and HOD since doing that will change the original density ratio.
Due to this, in practice, we found that F-DRE tends to work only in fairly limited settings.
D ENSITY R ATIO E STIMATION BY S CALING D ENSITIES
We tackle the density chasm problem using the notion of scaling densities from the scaled-Bregman divergence (equation 1, sBD).
sBD introduces a scaling measure m in the estimation of the discrepancy between two distributions p and q .
This allows for changing the sampling distribution in the estimation of, for instance, KL divergence as described in equation 3 from p to m by re-writing the original integrand log pq as
[Equation]
We leverage this insight and propose the S CALED D ENSITY R ATIO E STIMATOR ( S DRE) method for density ratio estimation that is theoretically well-deﬁned, computationally more efﬁcient, does not suffer from distribution-shift issue and thereby, mitigates the density-chasm problem originating from both FOD and HOD discrepancies.
To this end, in S DRE, we ﬁrst scale the densities p and q by another probability measure m and estimate log pq as log pm − log qm .
If the scaling measure m is constructed such that p << m , q << m and its samples overlap with the samples from both p and q , then as shown in the next section, a single multi-class logistic regression trained to classify samples from p, q and m into 3 classes, can accurately recover pm and qm on samples from any of the 3 distributions.
Later, we show that the S DRE approach can also be applied to model-free DREs (Section 5.3).
4. M ULTI - CLASS C LASSIFICATION TO M ULTI - DISTRIBUTION DRE
Nock et al.
(2016) build upon Proposition 3 of Menon & Ong (2016) (equation 5) using their scaledBregmann theorem and state the following lemma.
Lemma 1 Given a class-probability estimator ˆ η : X (cid:55)→ [0 , 1] C − 1 , let the density ratio estimator
ˆ r ( x ) =
ˆ η ( x ) ˆ η C ( x ) .
Then for any convex differentiable f : [0 , 1] C − 1 (cid:55)→ R ,
[Equation]
Here, f † = g ( x ) f ( x g ( x ) ) and g ( x ) =
π C 1 − π C + ˜ π T x . π C is the prior probability of class C , η C ( x ) is the class probability of x under the class C and p C is the distribution of X indexed by the class C . ˆ r is a C − 1 dimensional vector of density ratios between the C − 1 densities ( p 1 , . . . p C − 1 ) and the reference density p C .
Similarly, ˆ η ( x ) is a C − 1 -dimensional vector of class probabilities.
Please note that the choice of the reference class C is arbitrary.
In simple terms, Lemma 1 shows that a multi-class logistic regression (among other multi-class classiﬁers) leads to a multi-class density ratio estimator that minimizes the sBD to the corresponding true density ratio on samples from the reference distribution.
We demonstrate the implication of Lemma 1 on S DRE in Figure 2b.
Here, we train S DRE on the same exact setup as TRE in Figure 2a where p = N ( − 1 , 0 . 1) and q = N (+1 , 0 . 2) .
To this end, we propose to use Cauchy (0 , 1) as the scaling measure.
This is because the samples from the heavy tailed distributions easily overlap with the highly concentrated samples of p and q as shown in
Figure 2b.
Then we train a single 3-way logistic regression to classify samples from p, q and m .
As clearly shown, S DRE is able to perfectly estimate log p/q on samples from all the three distributions without suffering from any distribution shift issues.
This empirically conﬁrms that, unlike TRE, the
S DRE estimator converges to the correct density ratio on the entire domain of m as long as m is such that both p and q are absolutely continuous to m .
S DRE can also use more than one scaling measures { m k } Kk =1 .
Not only does Lemma 1 hold true for any number ( K ) of distributions, but also Ma & Collins (2018) show that, when the number of samples is ﬁxed, increasing K leads to a more consistent estimator.
Please note, increasing the number of scaling measure does not change the estimation of r p/q .
In a multi-class logistic regression model trained on samples from p, q and m k , the ratio between p and q is estimated as the function of their logits h 1 , h 2 , h k as follows.
From Lemma 1, r p/m k = exp( h 1 ) exp( h k ) i.e. log r p/m k = log exp( h 1 ) exp( h k )
=
h 1 − h k .
Similarly, log r q/m k = log exp( h 2 ) exp( h k )
=
h 2 − h k .
This gives us the ﬁnal ratio as, log r p/q = log r p/m k − log r q/m k = h 1 − h 2 , that can be calculated similarly with respect to any of the K scaling measure.
To demonstrate that the improved performance of S DRE over TRE in the example from Figuredoes not only stem from the novel construction scheme for m , we also train S DRE using TRE’s waymark distributions as scaling measures.
This is illustrated in Figure 5c in Appendix B.
Unlike
TRE that does not estimate log pq accurately on samples from p, q, p 1 and p 3 , S DRE is able to estimate log pq on all 5 distributions signiﬁcantly better even with sub-optimal scaling measures.
4. C ONSTRUCTING M
In this work, we use three ways of constructing the scaling measure M .
Here we describe them brieﬂy.
Please see Appendix A for details.
Overlapping Distribution:
We introduce a novel scaling measure construction scheme under which, M is deﬁned as any distribution whose samples overlap with both p and q and p << m , q << m .
For example, it can be heavy-tailed distributions, such as Cauchy and Student-t, or simply normal and uniform distributions or their mixtures.
Linear Mixing:
In this construction scheme, distribution M is deﬁned as the empirical distribution of the samples constructed by linearly combining samples X p = { x ip } Ni =1 and
X q = { x iq } Ni =1 from distributions p and q respectively.
That is, m is the empirical distribution over the set X m =
{ x im | x im =
(1 − α ) x ip + αx iq , x p ∈
X p , x iq ∈ X q } Ni =1 .
This construction is related to the linear combination waymark of Rhodes et al.
We clarify the difference in Appendix
A. Dimension-wise Mixing:
In this construction scheme, that is borrowed from TRE as it is, M is deﬁned as the empirical distribution of the samples generated by combining different subsets of dimensions from samples from p and q .
4. S DRE VS TRE
S DRE differs from TRE in two major ways.
First, TRE instructs to create K , closely-packed waymark distributions p 1 , . .
, p K between p and q such that none of the K +1 BC-DREs that it proposes to train, suffers from the density chasm problem.
This approach is not only computationally more demanding as it requires training K + 1 models (with some parameter sharing), but also leads to the aforementioned distribution shift issue.
In contrast, S DRE instructs to introduce scaling densities { m k } Kk =1 such that their samples overlap with the samples from both p and q .
It then estimates log pq as log pm k − log qm k by training a single multi-class logistic regression model across all K +distributions.
This approach is not only computationally more efﬁcient, but is also theoretically guaranteed to not suffer from distribution shift as shown in the previous section.
Second, for K = 1 , TRE and S DRE are similar as they both can be seen as using telescoping of ratios to estimate r p/q .
However, they are not equivalent.
TRE proposes the following telescoping:
[Equation]
m/q and S DRE proposes to telescope in the following manner: log p/q = log p/m − log q/m .
As such, TRE and S DRE use two different sets ( dMdQ vs dQdM in the second term) of Radon-Nikodym Derivatives (RND) to re-deﬁne log p/q .
The consequence of this is that TRE is only deﬁned when p << m << q , whereas S DRE is deﬁned for any m such that, p << m and q << m .
This implies that only under the following two conditions TRE and S DRE telescopings are equivalent: (1) K = 1 and (2) m is such that p << m << q and q << m .
This condition holds when p, q and m have full support.
However this condition easily breaks if, for example, p and q are mixtures of ﬁnite support distributions except for the trivial case when support of m is exactly equal to the support of q .
We demonstrate this in Figure 6 in Appendix B with a concrete example.
Another important and practical implication of this distinction is that, since m only appears in the denominator in S DRE, it can also be used to improve model-free approaches to density ratio estimation, such as the ﬁxed-design estimators based on the maximum mean discrepancy measure (Sugiyama et al., 2012).
We show this in Section 5.3.
It is also worth noting that unlike TRE, S DRE is realized as a single multi-class logistic regression that is trained using a softmax cross-entropy loss.
This allows us to use the results from Ma & Collins (2018) who formally studied the distinction in the convergence of softmax cross-entropy and binary cross-entropy based NCE methods (e.g. sDRE vs TRE/BC-DRE). They showed that ranking-based noise contrastive estimation (Gutmann &
Hyv¨arinen, 2010) (multi-class objective) is consistent under a much weaker assumption compared to the binary version.
Speciﬁcally, in contrast to the binary version, the ranking based approach does not require the capacity to model the normalization constant of the distributions.
E XPERIMENTS
In this section, we provide a thorough empirical analysis on the accuracy and robustness of S DRE and demonstrate that S DRE outperforms all other baselines – BC-DRE, TRE (Rhodes et al., 2020), and F-DRE (Choi et al., 2021).
We start with a toy 1D Gaussian dataset and fully analyze the behavior of S DRE.
Then, we extend to high-dimensional complex tasks ranging from mutual information (MI) estimation to representation learning and show signiﬁcant improvements on all accounts.
Finally, we also show how S DRE can be applied to improve model-free DREs. [Table]
5. 1D E XPERIMENT
In the following 1D experiments, we use two 1D Gaussian distributions and consider the four setups described in Table 1.
The ﬁrst and second conﬁgurations (row-wise) are designed such that the distributions are only separated by HOD and FOD, respectively.
The other two conﬁgurations represent the cases in which the distributions are separated by both FOD and HOD.
In S DRE, we use the quadratic form p ( x ) =
wx 2 + b to model each of the densities and use Cauchy centered at 0 and linear mixing to construct m .
We use KL-divergence estimation (expectation of the log ratios on samples from p ) as the evaluation metric.
We provide the exact setup for S DRE in Table 3 in
Appendix C.
In all of the conﬁgurations, while the baseline methods, including TRE and F-DRE, fail to correctly estimate the KL divergences by a signiﬁcant margin, S DRE reliably estimates them as shown in
Table 1.
Figure 7 in Appendix C illustrates that S DRE accurately estimates the underlying density ratios on the entire support of m as well.
These results demonstrate that it is better to use a single scaling measure whose samples overlaps with those from both p and q , instead of using a chain of upto K = 28 closely-packed waymark distributions (see Table 3 in Appendix C for the exact
TRE conﬁgurations).
To provide further clarity into S DRE’s density ratio estimation behavior, we analyze the reliability of its log ratio estimates using the Bayesian setup in Appendix D. The high accuracy of S DRE’s KL divergence estimates can be attributed to the fact that the S DRE estimator is most certain and precise around the support of p , and the KL divergence estimates are computed with samples from p .
This is illustrated in Figure 8 of Appendix D.
5. H IGH D IMENSIONAL E XPERIMENT
Following from TRE, we use the MI estimation problem from Belghazi et al.
; Poole et al.
(2019) to evaluate S DRE on the more challenging higher-dimensional setup.
In this problem, we estimate the mutual information, between a standard normal distribution and a Gaussian random variable x ∈ R 2 d with a block-diagonal covariance matrix, where each block is 2 × 2 with 1 on the diagonal and ρ on the off-diagonal.
ρ is computed given the number of dimensions and target mutual information with I = − d 2 log( − ρ 2 ) .
Because this method of distribution construction only creates separation in HOD, we further adapt it by moving the means of the two distributions and create harder problems with both FOD and HOD.
In S DRE, we model the densities of p , q , and the scaling measures m k with a quadratic form p ( x ) =
x
T Wx + g φ ( x ) + b , where g φ is a linear layer.
We use linear-mixing to construct m k , where K = 3 or K = 5 .
We provide the exact conﬁgurations for S DRE in Table 4 in
Appendix E.
We describe the setups and report the results averaged across 3 runs with different random seeds in Table 2.
In all of the settings, S DRE outperforms all other baselines, conﬁrming the generalization of the results from the toy 1D experiments.
These results corroborate our conjecture that TRE only solves problems with HOD.
It is worth noting that, using only upto 5 scaling measures that are constructed using the linear mixing scheme, S DRE beats TRE substantially on cases with both FOD and HOD even though, TRE uses upto 15 waymarks that are also constructed using their linear mixing approach.
This demonstrates, as per Lemma 1, our proposal of using the multi-class logistic regression does, in fact, prevent distribution shifts issues present in TRE when both FOD and HOD are present.
F-DRE, in contrast, is only accurate in problems with FOD, but tends to underestimate the log ratios even more than BC-DRE on settings with MI greater than 20 .
Further details about this experiment such as how we chose m
and K can be found in Appendix E, where we also provide plots of estimated log ratio vs the ground truth log ratio and training curves. [Table]
[Figure] Robustness and Generalization beyond Gaussian Distribution: We rigorously test the robustness of S DRE by evaluating it on the following setups.
(1) p and q as 1D truncated normal distributions to evaluate it on another setup with distributions with ﬁnite support.
(2) Randomized mean parameters to ensure that the S DRE model does not beneﬁt from the symmetry of the distributions around 0 in Table 2.
(3) p and q as Student-T distributions with randomized means to test the performance of S DRE under model mismatch.
In all these settings, S DRE model is able to reliably estimate the true KL-divergence demonstrating that it is highly robust.
Please see Table 5 in Appendix G for results and additional details.
5. I MPROVING KMM- BASED DRE WITH S DRE
We demonstrate the applicability of S DRE in the ﬁxed design setup by applying it to the Kernel
Mean Matching-DRE method (KMM-DRE; Huang et al., 2006; Sugiyama et al., 2012) for KL estimation.
In ﬁxed design setup, the DRE method directly outputs the densities for a set of samples without building a model.
We consider two KL estimation experiments with increasing difﬁculty: (a) We estimate KL between N ( − d, 1) and N (+ d, 1) by varying d from 0 to 3 and (b) We estimate
KL between N (0 , s 2 ) and N (0 , 1 /s 2 ) by varying s from 1 to 3.
For each experiment, we report the following KL estimates: (1) KL by MC using true ratios, (2) KL by MC using KMM-DRE estimates, and (3) KL by MC using estimates from KMM-DRE with S DRE for three types of m construction: linear mixing, mixture, and Cauchy.
We conduct the same experiments for different bandwidth σ , the parameter of the Gaussian kernel used by the KMM estimator.
As it can be seen from Figure 14,
S DRE consistently improves KMM-DRE in both experiments with different difﬁculty by improving the baseline KMM-DRE KL divergence estimates by up to 11 nats in some cases.
We defer a detailed discussion of the results to Appendix H.
5. R EPRESENTATION L EARNING IN S PATIAL M ULTI O MNIGLOT
Following the setup from Ozair et al.
; Rhodes et al.
, we apply S DRE to mutual information estimation and representation learning in the SpatialMultiOmniglot problem.
The goal is to estimate the mutual information between u and v , where u is a n × n grid of Omniglot char- [Figure] acters from different Omniglot alphabets and v is a n × n grid containing the next characters of the corresponding characters in u .
After learning, we evaluate the representations from the encoder with a standard linear evaluation protocol (Oord et al., 2018).
For the model, as in TRE, we use a separable architecture commonly used in MI-based representation learning literature and model the densities with log p ( u, v ) =
g ( u ) T Wf ( u ) , where g and f are 14-layer convolutional ResNets
(He et al., 2015).
We construct the scaling distributions via dimension-wise mixing.
In this experiment, we only compare to the single ratio baseline and TRE because (Rhodes et al., 2020, Figure 4) already demonstrated that it signiﬁcantly outperforms both Contrastive Predictive Coding (CPC)
(Oord et al., 2018) and Wasserstein Predictive Coding (WPC) (Ozair et al., 2019) models of representation learning on the same exact task.
As can be seen in Figure 4a, S DRE outperforms both TRE and the single ratio baseline, exactly matching the ground truth MI.
This improvement in MI estimation is reﬂected in the representations.
Figure 4b illustrates that S DRE’s encoder learns representations that achieve ∼ 100% Omniglot character classiﬁcation for both d = n 2 = 4 , 9 .
On the other hand, the performances of the single ratio estimator and TRE (using the same exact dimension-wise mixing to construct waymark distributions) both degrade noticeably as the complexity of the task increases.
As such, TRE only reach up to 91% and 85% respectively for d = 4 and d = 9 .
All models were trained with the same encoder architecture to ensure fair comparison.
Please refer to Appendix F for further details and two additional experiments where we study the effects of varying K and the encoder design.
C ONCLUSION
In this work, we presented the alternative method of S DRE for density ratio estimation.
We demonstrated that it has better theoretical grounding and improves the state-of-the-art in more challenging, high-dimensional density ratio estimation problems that have both FOD and HOD type of discrepancies.
However, the presented framework has its limitations.
For example, while S DRE estimates the ratio fairly well, (1) we do not provide any bounds on the estimation and (2) m is not learned but provided to the model.
We hope to address these issues in future work.
Moreover, our representation learning results suggest that improving MI estimation correlates with better representation learning.
Therefore, another interesting direction for future work could be to explore if S DRE can improve MI-based representation learning methods, such as contrastive learning (e.g. Oord et al., 2018; Chen et al., 2020).
E THICAL I MPACT
While density ratio estimation does not directly have major ethical concerns, it can, however, be used towards improving other deep learning methods of representation learning, generative models, etc.
These work tend to have signiﬁcant ethical considerations.
For example, improved generative models using S DRE could potentially be used to advance the deep-fake methods that may be used with malicious intent.
However, we hope that other applications of our work, such as those of MI estimation in science domains, outweigh the drawback of its potential illicit usages.
R EPRODUCIBILITY S TATEMENT
We have taken efforts to clearly describe all the experimental setups in separate appendices, one for each of the experiments.
For the baseline methods, we have used models provided by the original authors.
Where possible we have used existing evaluation tasks and metrics to report the results.
Finally in order to ensure reproducibility of our own results we are releasing all the code at the time of submission.
While most of our code is setup as Jupyter notebooks for easy accessibility, we will work to add documentation to improve the readability further
A A PPENDIX C ONSTRUCTING M
We now elaborate on the three types of scaling measures that we used in this work
Overlapping Distribution:
S DRE estimator, log pq = log
pm − log mq is deﬁned when p << m and q << m .
Therefore, m needs to such that its support contains the supports of p and q .
Any distribution with full support such as the normal distribution trivially satisﬁes this requirement.
However, satisfying this requirement does not guarantee empirical overlap of the distributions p , q with m in ﬁnite sample setting.
In order to ensure overlap of samples between the two pairs of distributions we recommend the following:
• Heavy-tailed Distributions: Distributions such as, Cauchy and Student-t are better choice for M compared to the normal distribution.
This is because their heavier tails allow for easily connecting p and q with higher sample overlap when they are far apart (especially in the case of FOD).
• Mixtures: Another way to connect p and q using m such that they have their samples overlap, is to use the mixture distribution.
Here, we ﬁrst convolve p and q with a standard normal and then take equal mixtures of the two. • Truncated Normal: If p and q have ﬁnite support, one can also use a truncated normal distribution or a uniform distribution that at least spans over the entire support of q .
This is assuming that p << q .
Linear Mixing:
In this construction scheme, distribution M is deﬁned as the empirical distribution of the samples constructed by linearly combining samples X p = { x ip } Ni =1 and X q = { x iq } Ni = from distributions p and q respectively.
That is, m is the empirical distribution over the set
X m = { x im | x im = αx ip + (1 − α ) x iq , x p ∈ X p , x iq ∈ X q } Ni =1 .
This construction is related to the linear combination waymark of Rhodes et al.
In TRE the waymark distribution is deﬁned as the empirical distribution of the set X m =
x im | x im = √ 1 − α 2 x ip + αx iq , x p ∈ X p , x iq ∈ X q } N i =1 .
This weighting scheme skews the samples from the waymark distribution towards p .
Therefore, care needs to be taken when p and q are ﬁnite support distributions so that the samples from the waymark distributions do not fall out of the support of q .
Using either of the weighting schemes, one can construct K different scaling measures.
S DRE can either use these K scaling measures separately using a K+2-way classiﬁer or deﬁne a single mixture distribution using them as component distributions and train a 3-way classiﬁer.
We refer to this construction as Mixture of Linear Mixing.
Dimension-wise Mixing:
In this construction scheme, that is borrowed from TRE as it is, M is deﬁned as the empirical distribution of the samples generated by combining different subsets of dimensions from samples from p and q .
We describe the exact construction scheme from TRE below for completeness:
Given a d -length vector x and that d is divisible by l , we can write down x = ( x [1] , ...
x [ l ]) , where each x [ i ] has the length d/l .
Then, a sample from the k th scaling measure is given by: x ik = ( x iq [1] , ...x iq [ j ] , x ip [ j + 1] , ..., x ip [ l ]) , for j = 1 , ..., l ) , where x ip ∼ p and x iq ∼ q are randomly paired.
B A PPENDIX TRE VS S DRE
In Figure 2b in Section 3, we showed that in comparison to TRE that uses 3 intermediate waymark distributions, S DRE is able to signiﬁcantly better estimate the ground truth density ratio using just a single Cauchy scaling measure that we introduced.
To demonstrate that this improved performance does not only stem from the novel m , but also from the use of a single multi-class logisitic regression (afforded by the exact form of S DRE telescoping), we train S DRE using TRE’s waymark distributions as scaling measures.
This is illustrated in Figure 5c.
Unlike TRE that does not estimate log pq accurately on samples from p, q, p 1 and p 3 , S DRE is able to estimate log pq on all 5 distributions signiﬁcantly better even with sub-optimal scaling measures.
The impact of using the sub-optimal scaling measure in S DRE on the quality of density ratio estimation can be teased out by comparing to the results in Figures 5b where S DRE is trained using Cauchy for m .
Finite Support Distribution Example for K = 1 :
For the case, K = 1 , in Section 4.3 we theoretically showed that TRE and S DRE are not equivalent in general.
We now demonstrate this with a speciﬁc example in Figure 6.
Here we set p = 0 .
5 × T N ( − 1 , 0 .
1 , low = − 1 .
1 , high = − 0 .
9)+0 .
5 ×T N (1 , 0 . 1 , low = 0 .
9 , high = 1 . 1) and q = 0 .
5 ×T N ( − 1 , 0 .
, low = − 1 . 2 , high = 0 .
8) + 0 .
5 × T N (1 , 0 . 2 , low = 0 .
8 , high = 1) , as shown in Figure 6a.
We set both, the waymark distribution TRE and the scaling measure for S DRE to m =
T N (0 , 1 , low = − 1 . 2 , high = 1 . 2) using the proposed overlapping distribution construction.
As such, p << q << m and therefore,
TRE is undeﬁned for the second term as mq is not deﬁned for samples from m that are outside the support of q .
On the other hand, S DRE is well deﬁned as qm is ﬁnite and deﬁned over the entire support of m .
It can be clearly seen in Figure 6b that TRE estimator for mq blows up to very high values on samples from m where q does not have any support.
This, however, is not the case with
S DRE, which accurately estimates qm as 0 on samples outside the support of q as shown in Figure 6c.
Despite being undeﬁned, when used with our proposed m , TRE may still estimate log r p/q accurately on samples from p .
We conjecture that this is because, the numerical estimation of the dM/dQ is ﬁnite over the support of p .
As such, care needs to be taken when equating TRE to S DRE to ensure that the relevant RNDs are well deﬁned along with of the accuracy of the numerical estimation.
[Figure]
[Figure]
Figure q
TRE p k
S DRE m a
N (0 , 1 e − 6)
N (0 , 1)
Linear
Mixing with α = [6.10e-05, 0.0078, 0.13]
Mixture of Linear
Mixing with α = [6.10e-05, 0.0078, 0.13] b
N ( − 1 , 0 . 08)
N (2 , 0 . 15)
Linear
Mixing with α = [0.053, 0.11, 0.16, 0.21, 0.26, 0.31, 0.37, 0.42, 0.47, 0.53, 0.58, 0.63, 0.68, 0.74, 0.79, 0.84, 0.89, 0.95]
C (0 , 1) c
N ( − 2 , 0 . 08)
N (2 , 0 . 15)
Linear
Mixing with α =
C (0 , 1) d
N ( − 5 , 1)
N (5 , 1)
Linear
Mixing with α = [0.11, 0.22, 0.33, 0.44, 0.55, 0.66, 0.77, 0.88]
C (0 , 2)
C A PPENDIX 1D DENSITY RATIO ESTIMATION TASK
In Section 5.1 we used KL-divergence as the evaluation metric to assess the accuracy of the density ratio estimation of each of the models, BC-DRE, TRE, F-DRE and S DRE. Since KL-divergence only evaluate the density ratio over samples from p , here we provide the plots of the log density ratio for all the models over a larger interval to better capture their behavior. The conﬁgurations of the experiments are described in Table 3.
Column 1 (Figure Labels) corresponds to the respective columns in Figure 7.
The rows of Figure 7 represent the four models in this order: BC-DRE, TRE, F-DRE, and S DRE.
Except for the ﬁrst conﬁguration (Figure Label a), which was trained with 1,000 samples in total, all other conﬁgurations are trained with 100,000 samples in total for all four models.
As evident from the ﬁgure, S DRE is the only model that not only accurately estimates the KL divergences, but also perfectly estimates the log density ratios on the entire support of m .
D A PPENDIX U NCERTAINTY Q UANTIFICATION OF S DRE L OG - RATIO E STIMATES WITH H AMILTONIAN M ONTE C ARLO
In the 1D experiments, S DRE consistently led to highly accurate KL divergence estimates even in challenging settings where state-of-the-art methods fail.
To understand why S DRE gives such accurate KL estimates, we conduct an analysis on the reliability of its log-ratio estimates by analyzing the distribution of the estimates in a Bayesian setup, and study how it impacts the KL divergence estimation.
For this analysis, we use a classiﬁer with the standard normal distribution as the prior on its parameters.
The distribution of the log-ratio estimates is simply the distribution of the estimates from the classiﬁers with different posterior parameters, which are sampled.
We consider the setup where p = N ( − 1 . 0 , 0 . 1) and q = N (1 . 0 , 0 . 2) and draw samples from the posterior using an Hamil- [Figure] tonian Monte Carlo (HMC) sampler.
We then compute a set of samples of the log-ratio estimates from S DRE and estimate the mean and quantiles using these samples.
Figure 8 shows these results.
We ﬁnd that S DRE is accurate and manifests lowest uncertainty around the mean ( − 1 . 0 ) of p .
The uncertainty increases as we move away from the modes of distributions p and q .
Since KL divergence is the expectation of the log-ratio on samples from p and the high density region of p exactly matches the high conﬁdence region of S DRE, it is able to consistently estimate the KL divergence accurately even when p and q are far apart [Figure] [Table]
E A PPENDIX H IGH D IMENSIONAL E XPERIMENT
In Section 5.2, we showed that S DRE performs better than all baseline models when p and q are high dimensional Gaussian distributions.
Prior work of Rhodes et al.
(2020) has considered high dimensional cases with HOD only, whereas here we additionally consider cases with FOD and HOD to provide a more complete picture.
Our results show that S DRE outperforms all other methods on the task of MI estimation as the function of the estimated density ratio.
It is worth noting that S DRE uses only upto 5 scaling measures that are constructed using the linear mixing scheme and beats
TRE substantially on cases with both FOD and HOD, although TRE uses upto 15 waymarks also constructed using linear mixing approach.
This demonstrates, as per Lemma 1, our proposal of using the multi-class logistic regression does, in fact, prevent distribution shifts issues of TRE when both FOD and HOD are present and help estimate the density ratio more accurately.
We now describe the S DRE conﬁguration and other setup related details.
Scaling Measures:
For all the high dimensional experiments throughout this work, we construct m using the linear mixing scheme as described in Appendix A. Table 4 provides the number K of scaling measures along with the exact mixing weights for each of the 6 settings.
As a general principle, we chose these three sets of mixing weights so that their cumulative samples overlap with the samples of p and q similar to how the heavy tailed distribution worked in the 1D case.
Please note, unlike the 1D experiment, the HOD in these experiments is signiﬁcantly higher.
Therefore, using a heavy tailed distribution does not work.
Linear mixing p and q on the other hand, mixes ﬁrst and higher order statistics and therefore, populates samples that overlap with both p and q .
In some cases, we found that a mixture of linear mixing with K = 1 can also be used to estimate the density ratio.
However, this requires using a neural network-based classiﬁer and requires much more tuning of the hyperparameters.
For choosing K , we use a grid search based approach.
We monitor the classiﬁcation accuracy across all the K + 2 distribution.
If this accuracy is very high ( > 95 % for all classes), this implies that the classiﬁcation task is easy and therefore the DRE may suffer from the density chasm issue.
On the other hand, if the classiﬁcation accuracy is too low ( < 50 % for all classes), then again, the DRE does not estimate well.
We found that targeting an accuracy curve as shown in Figure 9 (last panel) leads to optimal density ratio estimation.
This curve plots the test accuracy across all the classes and, empirically when it stays between the low and the high bounds of (50%,95%), the DRE estimates the ratios fairly well.
The ﬁrst panel shows that S DRE estimates the ground truth ratio accurately across samples from all the K +2 distributions, the second panel shows that KL estimates of S DRE is close to the ground truth KL and the third panel shows that both test and training losses have converged.
Figure 10 shows another example for the case of randomized means.
While S DRE also manages to get the ground truth KL correctly and most of the ratio estimates are also accurate, it does, however, slightly overestimate the log ratio for some of the samples from p .
Additional Baseline Comparisons We also compare to Kato & Teshima (2021) in a highdimensional setting (row 2 of 2 with dimensionality of 40 and µ 1 = − 1 , µ 2 = 1 .
In this setting in which the ground truth MI is 100, while sDRE meaningfully estimates the MI as 119.96, the best model from Kato & Teshima (2021) only estimates it to be 1 . 60 .
F A PPENDIX S PATIAL M ULTI O MNIGLOT E XPERIMENT
SpatialMultiOmniglot is a dataset of paired images u and v , where u is a n × n grid of Omniglot characters from different Omniglot alphabets and v is a n × n grid containing the next characters of the corresponding characters in u .
In this setup, we treat each grid of n × n as a collection of n 2 categorical random variables, not the individual pixels.
The mutual information I ( u, v ) can be computed as: I ( u, v ) =
(cid:80) n 2 i =1 log l i , where l i is the alphabet size for the i th character in u .
This problem allows us to easily control the complexity of the task since increasing n increases the mutual information.
[Figure] For the model, as in TRE, we use a separable architecture commonly used in MI-based representation learning literature and model the densities with log p ( u, v ) =
g ( u ) T Wf ( u ) , where g and f are 14-layer convolutional ResNets
(He et al., 2015).
We construct the scaling distributions via dimension-wise mixing – exactly the way that TRE does.
To evaluate the representations after learning, we adopt a standard linear evaluation protocol to train a linear classiﬁer on the output of the frozen encoder g ( u ) to predict the alphabetic index of each character in the grid u .
A DDITIONAL E XPERIMENTS
In addition to the experiments in the main text, we run two more experiments with SpatialMultiOmniglot to test the effect of the size of K and the effect of using the same encoder for g and f (i.e, modeling the densities with the form log p ( u, v ) =
g ( u ) T Wg ( v )) .
Varying K : We test the effect of changing K in the d = 4 experiment.
For K = 1 , we aggregate all the dimension-wise mixed samples into 1 class, whereas for K = 3 , we separate them into their respective classes (corresponding to the number of dimensions mixed).
We illustrate this effect in 11.
In line with the ﬁnding of Ma & Collins (2018), increasing the number of K not only helps the model reach the ground truth MI, but also the quality of representations worsens from 100% to 86 .
7% test classiﬁcation accuracy.
Single Encoder Design: Furthermore, we test the contribution of using two different encoders f and g instead of one.
As seen in Figure 12, in both cases of d = 4 , 9 , the two models reach slightly different but similar MI estimates, but, interestingly, do not differ at all in the test classiﬁcation accuracy.
Empirically, we also found that using one encoder helps the model converge to much faster.
Overall, this experiment demonstrates that using two different encoders does not necessarily work to our advantage.
G A PPENDIX R OBUSTNESS AND M ODEL M IS M ATCH A NALYSIS
In order to evaluate the robustness of S DRE method, we use the following setups.
Breaking Symmetry:
In the experiment setup for the high dimensional experiments as described in Section 5.2, the means of the Gaussian distribution p and q were symmetrical in majority of the cases.
In order to ensure that this symmetry around zero did not provide an advantage to the S DRE method, we also evaluate it on Gaussian p and q with randomized means for the two settings shown in rows 2 and 3 of Table 5.
Beside the mean parameters, we do not change any other conﬁguration in the experiment or the model.
S DRE manages to estimate the ground truth KL divergence accurately in this case, therefore establishing that it did not beneﬁt unfairly from the symmetry of distributions around zero in the high dimensional experiments.
Model Mismatch:
In rows 3, 5 and 6 of Table 5, we evaluate S DRE by replacing one or both distributions p and q from Gaussian in the case above with Student-t distribution of the same scale with randomized means.
We set the degrees of freedom to 10 .
These experiments test how S DRE performs when there is a model mismatch, i.e. how S DRE performs using the same quadratic model that was used when p and q were set to be Gaussian with lighter tails.
As can be seen, S DRE seems to be very robust in these cases and is able to accurately estimate the ground truth KL.
p and q with Finite Support: Finally, we test S DRE on another case where p and q are ﬁnite support distributions that have both FOD and HOD.
For this, we set them to truncated normal, as shown in row 1.
This setting is similar to the 1D example with Gaussian p and q that we have used throughout the manuscript.
We also set m to be truncated normal with scale set to 2 in order to allow it to have overlap with both p and q .
As expected, S DRE not only manages to correctly estimate the ground truth KL, but also correctly decays its log ratio estimates when the samples do not fall under the support of p as shown in Figure 13. Dim p q m
GT-KL
S DRE
Truncated Normal loc=-1, scale=0.
Truncated Normal loc=1, scale=0.
Truncated Normal loc=-1, scale= (-1.1,1.2) 50. Normal loc=R(-.5,.5), cov= 2 × 2 BD
Normal loc=R(-.5,.5), cov= I
LM 54.
Normal loc=-1, cov= 2 × 2 BD
Mixture of
Normal loc=0.9, cov= I and Normal loc=0.9, cov= I
LM 105.
Student T loc=
R(-.5,.5), scale= 2 × 2 BD, df=
Student T loc=
R(-.5,.5), scale= I , df=
LM 53.
Normal loc=R(-1,1), cov= 2 × 2 BD
Normal loc=R(-1,1), cov= I
LM 110.
Student T loc=R(-1,1), scale= 2 × 2 BD, df=
Student T loc=R(-1,1), scale= I , df=
LM 103.
Normal loc=0, cov= 2 × 2 BD
Student T loc=0, scale= I , df=
LM 82.
H A PPENDIX I MPROVING KMM- BASED DRE WITH S DRE
Another advantage of S DRE is that it is applicable to DRE methods with a ﬁxed design setup, i.e. a
DRE method that directly outputs the densities for a set of samples without building a model; TRE is not applicable to DRE methods in this setup as it models the ratios directly.
The ﬂexibility of our model derives from the fact m is in the denominator of both ratios.
An example of DRE methods with a ﬁxed design setup is the kernel mean matching (KMM) based DRE method (KMM-DRE;
[Figure]
The main advantage of this method is that the ratio estimators can be computed in closed form and therefore are computationally more
We demonstrate the applicability of S DRE in the ﬁxed design setup by applying S DRE to the
KMM-DRE method for KL estimation.
In order to show the beneﬁt of S DRE, we consider two
KL estimation experiments with increasing difﬁculty: (a) We estimate KL between N ( − d, 1) and
N (+ d, 1) by varying d from 0 to 3 and (b) We estimate KL between N (0 , s 2 ) and N (0 , 1 /s 2 ) by varying s from 1 to 3.
For each experiment, we report the following KL estimates: (1) KL by MC using true ratios, (2) KL by MC using KMM-DRE estimates, and (3) KL by MC using estimates from KMM-DRE with S DRE for three types of m construction: linear mixing, mixture, and Cauchy.
1 We conduct the same experiments for different bandwidth σ , the parameter of the
Gaussian kernel used by the KMM estimator.
Figure 14 summarizes the results.
Overall, S DRE signiﬁcantly improves the baseline KMM-DRE KL divergence estimates by up to 11 nats.
As it can be seen from the ﬁgures, S DRE consistently improve KMM-DRE in both experiments with different difﬁculty.
In Figure 14a, it can be seen S DRE with mixture and Cauchy as M gives
KL estimates closer to the true value even when d = 2 . 5 , while KMM-DRE fails even for d = 1 . 0 .
In Figure 14b, we can see S DRE with mixing M helps calibrate the KL estimates from KMM-DRE:
For s ≤ 1 . 75 , KMM-DRE overestimates the KL and S DRE gives estimates closer to the true KL by giving a lower estimate; for s ≥ 1 . 75 , KMM-DRE underestimates the KL and S DRE gives estimates closer to the true KL by giving a higher estimate.
It is also worth noting that different For KMM-DRE, we solve the constrained optimization problem using the solvers from JuMP.jl, via Den- We set
[Table]