Dimension reduction via score ratio matching
Abstract
We propose a method to detect a low-dimensional subspace where a non-Gaussian target distribution departs from a known reference distribution (e g., a standard
Gaussian).
We identify this subspace from gradients of the log-ratio between the target and reference densities, which we call the score ratio .
Given only samples from the target distribution, we estimate these gradients via score ratio matching, with a tailored parameterization and a regularization method that expose the lowdimensional structure we seek.
We show that our approach outperforms standard score matching for dimension reduction of in-class distributions and that several benchmark UCI datasets in fact exhibit this type of low dimensionality
Introduction and motivation
Dimension reduction methods are ubiquitous in large-scale data science, statistical modeling, and machine learning.
The computational burdens of many common analyses and algorithms may scale poorly with the dimension of the problem, and accurately capturing complex dependence structure in high-dimensional problems may require massive sample sizes.
Many such tasks involve characterizing an unknown target distribution π given only a representative set of samples { x i } ni =1 ∼ π .
Both generative modeling and density estimation are examples of these tasks [9, 15].
When one has access to the (unnormalized) target density and its gradients, dimension reduction methods are well-explored [19, 5, 20, 1, 2].
One such method, certified dimension reduction [20], uses gradients of the log-density to expose a low-dimensional subspace where the target distribution departs most strongly from a reference distribution.
In the context of Bayesian inference, this structure has been used to accelerate MCMC sampling methods [3, 4] or flow-based variational approximations [2], with error guarantees.
Here, the target is the posterior distribution and a natural choice for the reference is the prior distribution.
In this work, we develop an analogous dimension reduction method for when one is given a set of samples { x i } ni =1 ∼ π , but the target density is unavailable.
We propose an algorithm for uncovering this low-dimensional structure based on score ratio matching (§3), rather than direct score matching.
Our algorithm employs a training objective, a network parameterization, and a regularization penalty all tailored to our dimension reduction goal.
We demonstrate that our score ratio matching method better reveals low-dimensional structure compared to standard score matching, and that several common benchmark datasets in fact exhibit this kind of low-dimensional structure (§4).
Background
Score matching overview Score matching has recently appeared as a powerful framework with applications to generative modeling [16, 17, 10, 6] and Bayesian inference [21, 13].
The core task is approximating the score function ∇ x log π ( x ) with a neural network s θ ( x ) , referred to as the score network [17].
We direct readers to [16, 17] for an overview of score matching, especially the derivation of the objective function for learning the score, network training strategies, and its application to generative modeling using Langevin sampling.
Low-dimensional subspace hypothesis We begin by defining a class of target distributions that depart from a known reference distribution only in a few directions.
Definition 1.
Let ρ be a chosen reference density on R d .
Given a unitary matrix U ∈ R d × d and an integer r ≤ d , let D r ( U ) be the set of distributions with densities of the form
[Equation]
for some f : R r → R > 0 , where U r ∈ R d × r contains the first r columns of U .
Such structure can be exploited in several ways.
Let U =
[ U r U ⊥ ] be unitary and denote x r =
U ⊤ r x and x ⊥ =
U ⊤⊥ x .
Then, the target distribution rotated into the basis of U decomposes as
[Equation]
where ν r ( x r ) :
= f
( x r ) ρ r ( x r ) and ρ ⊥| r is known (and can be evaluated and simulated).
Hence, generative modeling and density estimation tasks are reduced to learning the r -dimensional distribution ν r using projected samples { U ⊤ r x i } ni =1 .
Leveraging such structure requires identifying both the basis
U and the minimal reduced dimension r needed to accurately approximate π within the set D r ( U ) .
Methods that exploit this structure are popular in Bayesian inference, where π is the posterior distribution.
There it is typical to take ρ to be the prior distribution; f is then an approximation to the likelihood function.
This assumption is natural in situations where we expect the data to only be partially informative of the parameter.
Such structure has been leveraged in both MCMC methods [5, 3] and variational inference [2].
In this work, we take ρ to be the standard Gaussian density N (0 , I d ) and identify a subspace where π departs from ρ .
We note that our results can be generalized to other reference distributions; see Appendix A. A similar structure was exploited for diffusion models in [11] by hand-selecting the low-dimensional subspace.
Here, we find this subspace by measuring the error incurred from approximating π with some π r ∈ D r ( U ) for a given basis U and reduced dimension r .
The following results provide a strategy to find a suitable U and r based on an upper bound for the KL divergence from π to its closest approximation within D r ( U ) .
Proposition 1 (Modified Proposition 2.12 of [20]) .
Let ρ be the standard Gaussian density, and let
[Equation]
be the so-called diagnostic matrix of size d × d .
Then,
For any r ≤ d and unitary matrix U ∈ R d × d there exists π r ∈ D r ( U ) such that
[Equation]
Let ( λ i , u i ) ∈ R ≥
× R d be the i -th eigenpair of the eigenvalue problem Hu i
= λ
i u i , λ 1 ≥ · · · ≥ λ d and take U =
[ u 1 , . . . ,
u d ] to be the matrix containing the eigenvectors of H .
Then, E r ( U ) is minimized for any r ≤ d and there exists π r ∈ D r ( U ) such that
[Equation]
The results of Proposition 1 have several practical implications.
First, given the ability to compute the diagnostic matrix H , one obtains an upper bound on the error (in KL divergence) induced by approximating π with a distribution in the class D r ( U ) , for any choice of U .
Second, a natural choice for U is the eigenbasis of H , and one may choose r based on the decay of the eigenvalues of H ; that is, given a KL error tolerance ε > 0 , one can pick r so that 12 ( λ r +1 + · · · + λ d ) < ε . Indeed, if the rank of H is r , then π ∈ D r ( U ) and thus the marginal of x ⊥ must be standard Gaussian.
Score ratio matching
We now describe how to approximate the score ratio function ∇ x log ( π ( x ) /ρ
( x )) using score ratio matching , enabling us to perform analogous dimension reduction given only samples
x i from π .
A naïve strategy would be to directly approximate the score of π and use it to compute the score ratio as ∇ x log ( π ( x ) /ρ
( x ))
= ∇ x log π ( x ) − ∇ x log ρ ( x ) .
Instead, we take a different approach that leverages the (possible) low-dimensional structure of π in Definition 1.
We approximate the score ratio directly using a score ratio network s θ : R d → R d by minimizing an objective function that does not require access to the score of the target density π .
This is made possible by the following proposition, whose proof is in Appendix B.
Proposition 2.
Let s θ be differentiable.
Then we have the following equivalence of objectives:
[Equation]
where C is a constant that only depends on the densities π and ρ .
In practice, we replace the expectation above with the empirical sum over the dataset to obtain a key term of our optimization objective,
[Equation]
Under our hypothesis on the target density in Definition 1, we expect the score ratio, rather than the score itself, to be well approximated by a ridge function [14], i.e., a function that is constant for x ∈
Im ( U ⊥ ) .
Next, we describe a parameterization for s θ ( x ) and a regularization method that are tailored to learning this low-dimensional structure.
Score-ratio network parameterization and regularization For π r ∈ D r ( U )
, the score ratio takes the specific form
[Equation]
We see that the range of the score ratio lies within the subspace spanned by U r .
We encode this observation into our score ratio network in two ways.
First, we parameterize the network as
[Equation]
where V ∈ R d × d , is the first and last layer’s weight matrix and ˜ s θ : R d → R d is a typical score network as described in [17].
This parameterization enforces that if V converges to a low (effective) rank during optimization, the range of s θ ( x ) is restricted accordingly.
We also use a regularization technique that helps reveal low-dimensional structure when it is present.
If Definition 1 holds, then we expect V to have (numerical) rank r .
As r is unknown, we penalize the nuclear norm of V V ⊤ , as commonly used for low-rank matrix estimation [8].
This leads to the final objective function
[Equation]
where ∥ · ∥ ∗ is the nuclear norm.
1: Input : Target data { x i } ni =1 ∼ π , and user tolerance ε >2: Center the mean and scale data by the Cholesky factor of the empirical precision matrix.
s θ F ( s θ ) the score-ratio approximation s θ ( x ) .
4: Estimate diagnostic H
[-=
-]=
s θ x i s θ x i )
λ i ,
i ) ≥
R d . u n ] and pick r so that (cid:98) E r ( U ) =
12 ( λ r +1 + · · · + λ d ) < ε
Numerical results
We now present several numerical experiments showing that: (i) our score ratio method more accurately captures the relevant low-dimensional subspace in a toy problem where this structure is known to be present, and (ii) this structure is present in several datasets from the UCI repository [7].
Details on the score network parameterization and training procedure are in Appendix C.
Embedded banana distribution Consider the following “embedded banana” distribution, where the data-generating process is defined by
[Equation]
and x = Ry , where R ∈ R 10 × 10 is a random rotation matrix that is sampled by computing the QR factorization of a matrix with standard Gaussian entries.
In this case, we have π ∈ D r =2 ( R ) .
Hence, we expect our algorithm to find the subspace spanned by the two leading columns of R .
In this example, we compute the score ratio analytically and define a consistent estimator for the true diagnostic matrix H .
In Figure 1a we plot the error bound E r ( U ) =
12 tr( U r U ⊤ r H ) for three different bases U : (1) the eigenbasis of the true diagnostic matrix; (2) the eigenbasis of the diagnostic matrix computed with our score ratio approximation; and (3) the eigenbasis of the diagnostic matrix computed with a standard score approximation (as described at the beginning of §3).
For our method, we see that E r ( U ) sharply drops at r = 2 to less than 10 − 2 .
We also see that our method yields considerably lower errors at each r compared to standard score matching.
For a visual representation of the results, Appendix C shows a scatter plot of additional held-out samples from π (which were not used during training) and the samples rotated into our discovered basis U when taking d = 3 .
UCI datasets We now report results for several datasets from the UCI repository [7], commonly used as benchmarks for density estimation: POWER, GAS, and MINIBOONE.
Since we take the reference distribution to be standard normal, before applying our algorithm we whiten the data; see Appendix A. For these datasets we do not have access to an analytic score function, and thus can only report the estimated KL error bounds, (cid:98) E r ( U ) , for each example.
In Figure 1b–d, we see that low-dimensional structure seems to be present in each dataset, via the rapid decay in the error bound for small r .
For example, an r = 1 dimensional subspace for the POWER dataset yields an approximation error on the order of 10 − 1 .
[Figure]
Discussion
We have proposed a dimension reduction methodology based on score matching.
Our framework identifies a subspace that best captures the departure of the data-generating distribution from a reference distribution.
While such methods are well studied in the context of Bayesian inference, our approach brings the benefits of dimension reduction with error guarantees to settings where only samples are available.
To aid in finding low dimensional structure, we introduced a network parameterization that exploits the score ratio’s gradient structure, coupled with a low-rank matrix recovery technique.
Future work will utilize the proposed framework to accelerate and improve the accuracy of density estimation and generative modeling by exploiting the discovered low-dimensional structure.
We also plan to investigate how the intrinsic dimension of the problem affects the number of samples needed to learn the score ratio and its hyperparameters. Acknowledgments and Disclosure of Funding This work was supported in part by the DOE Office of Advanced Scientific Computing Research (ASCR) M2dt MMICC center.
A Choice of reference distribution and data pre-processing
In this work we choose the reference distribution ρ to be a standard Gaussian, though many other choices are possible. The results in Proposition 1 depend on ρ satisfying a log-Sobolev inequality, which allows for uniform, multivariate Gaussian, and mixture-of-Gaussian reference distributions, among others; see [20] for more details and examples. In practice, ρ should be taken to be comparable to the samples in location and scale.
To implement the algorithm presented in Section 3, it is only required that the score of the reference be known and easily computable.
Since here the reference distribution chosen here is standard normal, we whiten the data before applying our algorithm to datasets from the UCI repository, i.e., we standardize the data to have zero mean and identity covariance matrix.
Specifically, we shift the data by its sample mean and transform it by the square root of a sample estimate of the precision matrix.
This ensures that the data can be meaningfully compared with the standard normal.
B Proof of Proposition
Let
[Equation]
(cid:18)
[Equation]
( x ) ρ ( x ) (cid:19)(cid:13)(cid:13)(cid:13)(cid:13)  .
We expand the squared norm to obtain
J ∗ ( s θ ) = 1 2 E π (cid:34) s θ ( x ) ⊤ s θ ( x ) + ∇ x log (cid:18) π ( x ) ρ ( x ) (cid:19) ⊤ ∇ x log (cid:18) π ( x ) ρ ( x ) (cid:19) − 2 s θ ( x ) ⊤ ∇ x log (cid:18) π ( x ) ρ ( x ) (cid:19)(cid:35) .
Note the second term, ∇ x log (cid:16) π ( x ) ρ ( x ) (cid:17)
⊤ ∇ x log (cid:16) π ( x ) ρ ( x ) (cid:17) , does not depend on the network parameters, and thus need not be included in our optimization objective.
The following steps rewrites the third term into quantities we can evaluate:
[Equation]
This leads to the final result that
[Equation]
[Figure]
For each numerical example, we used 10 4 training samples.
We use the Adam [12] optimizer with learning rate 5 × 10 − 3 and batch size 1000 to train the network for 500 epochs.
We take the nuclear norm regularization parameter to be λ = 0 .
8 /d . As discussed in [18], directly evaluating the trace operator in the objective function F is prohibitively expensive for even moderate dimensions d , and so we also make use of the sliced-score matching method proposed in that work with 100 projections.
The network ˜ s θ had 1 fully connected hidden layer for the embedded banana example, and 2 hidden layers for the UCI datasets with ReLU activation functions and width 128 .