Concrete Score Matching: Generalized Score Matching for Discrete Data

Introduction

When estimating a statistical model, the representation of the underlying probability distribution has profound implications on the downstream modeling task. Likelihood-based model families such as Variational Autoencoders (VAEs) , autoregressive models , normalizing flows , and diffusion models are forced to either approximate the intractable normalizing constant or utilize restrictive model architectures; implicit models try to capture the underlying sampling process by relying on unstable adversarial training procedures. On the other hand, representing the distribution as the gradient of the log probability density function—also known as the (Stein) score—allows us to circumvent such issues. This is key to score matching’s success on a wide range of continuous data modalities . In fact, its recent resurgence has led to significant advances in machine learning applications of density estimation, such as image generation and audio synthesis , among others.

However, score matching approaches that rely on modeling the gradient of the data distribution are inherently designed for continuous data; such methods hinge on the existence of the gradient, which is undefined for discrete domains . Given the ubiquitous nature of structured, discrete data in our world today such as graphs, text, genomic sequences, images, we desire an approach that would allow us to expand the successes of score-based generative models into discrete domains.

To address this challenge, we first introduce the Concrete score: a generalization of the (Stein) score that is amenable to both continuous and discrete data typesWe borrow the terminology from “concrete mathematics” .. Given a predefined neighborhood structure, the Concrete score leverages structural information in the data to construct surrogate gradient information about the discrete space by considering the similarity between two neighboring examples. More precisely, the Concrete score of any input is defined by the rate of change of the probabilities with respect to directional changes of the input. In direct analogy to the continuous (Stein) score, the intuition is that the Concrete score should remain small when two examples are “close” to one another; when the two examples are very different, the Concrete score will be large. We find that measuring such changes by the Euclidean distance recovers the (Stein) score in continuous domains, while using the Manhattan distance leads to our novel score function in discrete domains.

Using this definition of the Concrete score, we then introduce a framework to learn such scores from samples called Concrete Score Matching (CSM). To successfully scale our method to high dimensions, we also propose an efficient training objective as well as several variations of our approach that improve performance in practice. Empirically, we demonstrate the efficacy of CSM on a wide range of density estimation tasks on a mixture of synthetic, tabular, and high-dimensional image datasets, and demonstrate that it performs favorably relative to existing baselines for modeling discrete data.

The contributions of our work can be summarized as follows:

We propose the Concrete score, a generalization of the (Stein) score to discrete domains. We construct the Concrete score such that it leverages structural information in the data to capture surrogate “gradient” information over discrete spaces.

We introduce a novel score matching framework for estimating such scores from samples called Concrete Score Matching (CSM). We then outline several connections between CSM and existing (continuous) score matching techniques, such denoising score matching (DSM).

We propose efficient learning objectives that allow our method to scale gracefully to high-dimensional datasets, and show how to port over the recent successes in continuous score matching approaches to discrete domains.

Preliminaries

where $\textrm{tr}(\cdot)$ denotes the matrix trace, and the optimal score model satisfies ${\bm{s}}_{\theta^{*}}({\mathbf{x}})\approx{\bm{s}}({\mathbf{x}})$. While the trace term in Eq. 1 remains problematic—it requires an expensive evaluation of the Hessian of the log-density function—recent works leverage the Skilling-Hutchinson trace estimator or directional derivatives to efficiently approximate the training objective.

Despite their key role in successfully scaling score-based generative models to high-dimensional datasets, we note two critical limitations of existing score matching techniques: (1) ${\mathcal{X}}$ must be continuous; and (2) the density function $p_{\rm{data}}({\mathbf{x}})$ must be differentiable. This poses a significant challenge in discrete domains, where both requirements fail to hold.

The Concrete Score

The above limitations prevent the direct application of score matching techniques to discrete data. We thus propose the Concrete score, a generalization of the (Stein) score for discrete settings. The key intuition is that although the gradient is undefined in discrete spaces, we can still construct a surrogate for the gradient by leveraging local directional changes to the input. We do so by exploiting special neighborhood structures in the data, and elaborate upon the necessary conditions on these structures to guarantee that our surrogate gradient indeed recovers a valid score function that (1) completely characterizes the distribution, and (2) is amenable for parameter estimation.

To be more precise, let $p_{\rm{data}}(\mathbf{x})$ be the data distribution over ${\mathcal{X}}$. We denote ${\mathcal{N}}:{\mathcal{X}}\to{\mathcal{X}}^{K}$ as the function mapping each example ${\mathbf{x}}\in{\mathcal{X}}$ to a set of neighbors, such that $\mathcal{N}(\mathbf{x})=\{\mathbf{x}_{n_{1}},...,\mathbf{x}_{n_{k}}\}$ and $K=|{\mathcal{N}}(\mathbf{x})|$ We consider a fixed number of neighbors $K$ for each ${\mathbf{x}}$ to simplify our notation, but note that this can be generalized to a variable number of neighbors for every ${\mathbf{x}}$.. This neighborhood induces a particular graphical structure onto the support of $p_{\rm{data}}({\mathbf{x}})$, which we call the "neighborhood-induced graph", that will play a key role in constructing the surrogate gradient. We provide a formal definition below.

Let $p_{\rm{data}}(\mathbf{x})$ be the data distribution over ${\mathcal{X}}$ and ${\mathcal{N}}$ be the function mapping each node $\mathbf{x}\in{\mathcal{X}}$ to its set of neighbors. The neighborhood-induced graph ${\mathcal{G}}$ is the directed graph which results from adding a directed edge from ${\mathbf{x}}$ to each node in its neighborhood set $\mathbf{x}_{n}\in{\mathcal{N}}({\mathbf{x}})$, for all ${\mathbf{x}}\in\rm{supp}(p_{\rm{data}}({\mathbf{x}}))$.

An important point is that the neighborhood structure can be asymmetric, since the neighborhood-induced graph is a directed graph. This implies that there may exist cases where ${\mathcal{N}}({\mathbf{x}})=\{{\mathbf{x}}_{1}\}$ does not necessarily imply that ${\mathcal{N}}({\mathbf{x}}_{1})=\{{\mathbf{x}}\}$. We provide an intuitive example below.

When $\mathcal{X}=\{\mathbf{x}_{0},\mathbf{x}_{1},\ldots,\mathbf{x}_{5}\}$, and the neighborhood structure $\mathcal{N}$ is defined as: $\mathcal{N}(\mathbf{x}_{0})=\{\emptyset\}$ and $\mathcal{N}(\mathbf{x}_{i})=\{\mathbf{x}_{0}\}\textrm{ }\forall i=1,\ldots,5$, the neighborhood-induced graph is the star graph in Figure 1.

There exist a wide range of neighborhood structures, all of which yield different neighborhood-induced graphs. We visualize five common structures in Figure 1.

Such graphs provide insight into the kinds of neighborhood structures that are amenable to our framework. One necessary characteristic of such graphs is connectedness. In fact, the five common graphical structures (among others) shown in Figure 1 are all weakly connected graphs. We emphasize that connectedness does not take the directionality of the underlying graph ${\mathcal{G}}$’s edges into account—it does not matter whether a node ${\mathbf{x}}$ has an incoming or outgoing edge.

With the above definitions in place, we can now define the surrogate gradient as the local differences between a set of examples (similar to directional derivative in the continuous case). Using this notion of the “gradient”, we construct the Concrete score as the rate of change of the probabilities with respect to these local directional changes in the input ${\mathbf{x}}$ as defined below.

Although we have defined the Concrete score, we have yet to justify whether it is indeed a suitable score function for estimating $p_{\rm{data}}({\mathbf{x}})$. The necessary condition, which we call completeness , ensures that the Concrete score preserves enough information about $p_{\rm{data}}({\mathbf{x}})$ such that we can successfully learn $p_{\rm{data}}({\mathbf{x}})$ from data samples using score matching. We make this statement more precise in the following theorem.

Let $p_{\rm{data}}(\mathbf{x})$ be a (discrete) data distribution. Denote ${\bm{c}}_{p_{\rm{data}}}(\mathbf{x};\mathcal{N})$ as the Concrete score of $p_{\rm{data}}({\mathbf{x}})$ with neighboring structure $\mathcal{N}$, and ${\bm{c}}_{\theta}(\mathbf{x};\mathcal{N})$ the Concrete score for a distribution $p_{\theta}(\mathbf{x})$ parameterized by $\theta\in\Theta$. When the graph induced by the neighborhood structure $\mathcal{N}$ is connected, $c_{\theta}(\mathbf{x};\mathcal{N})={\bm{c}}_{p_{\rm{data}}}(\mathbf{x};\mathcal{N})$ implies that $p_{\theta}(\mathbf{x})=p_{\rm{data}}(\mathbf{x})$ $\forall\mathbf{x}\in{\mathcal{X}}$.

If two nodes ${\mathbf{x}}$ and ${\mathbf{x}}^{\prime}$ in a neighborhood-induced graph ${\mathcal{G}}$ share an edge, then their density ratio $p_{\rm{data}}({\mathbf{x}}^{\prime})/p_{\rm{data}}({\mathbf{x}})$ can be uniquely identified using ${\bm{c}}_{p_{\rm{data}}}(\mathbf{x};{\mathcal{N}})$. Thus when ${\mathcal{G}}$ is connected, the density ratio between any node pairs in ${\mathcal{G}}$ is uniquely identified given ${\bm{c}}_{p_{\rm{data}}}(\mathbf{x};{\mathcal{N}})$. Therefore, knowing the density ratio between any two points uniquely identifies $p_{\rm{data}}(\mathbf{x})$. ∎

We provide the full proof in Appendix A. We emphasize here that the connectedness of the neighborhood-induced graph (a kind of regularity condition on ${\mathcal{N}}$ and $p_{\rm{data}}(\mathbf{x})$) was crucial for demonstrating the completness of the Concrete score.

Note that Theorem 1 only requires the neighborhood-induced graph to be connected. This implies that given a connected graph, we can augment it with additional edges—the new graph will still remain complete, but contain additional information in the augmented neighborhood structure that may help parameter estimation in practice. This is a phenomenon that we observe empirically in Section 5.3.

2 Connection to Stein Scores and Existing Score Matching Techniques

The form of the Concrete score in Eq. 2 suggests a natural connection with the Stein score in continuous domains. In the following proposition, we illustrate the connection between the two when we define the neighborhood structure to be similar to the grid in Figure 1. In particular, as the distance between neighboring nodes shrinks to zero, our Concrete score converges to the Stein score up to a multiplicative constant.

We note that while from this perspective the Concrete score can be seen as a finite-difference approximation to the continuous (Stein) score, CSM is different from finite difference score matching (FD-SM) . We introduce a new family of score functions generalizable to discrete domains in the form of $(p({\mathbf{x}}+{\mathbf{v}})-p({\mathbf{x}}))/p({\mathbf{x}})$, where ${\mathbf{v}}$ is the direction from ${\mathbf{x}}$ pointing to its neighbor with length $\delta$. On the other hand, FD-SM still attempts to match the Stein score, and approximates directional derivatives with finite difference using two neighbors in the form of $\log p({\mathbf{x}}+{\mathbf{v}})-\log p({\mathbf{x}}-{\mathbf{v}})$. Nevertheless, the differences between the two forms are $o(\delta)$ as $\delta\to 0$.

3 Inference with Concrete Scores

To perform inference, we note that there exist several ways to define Markov chain Monte Carlo (MCMC) samplers that only rely on the Concrete score—we highlight the simplest setting of Metropolis-Hastings for clarity of exposition. We first note that, by definition:

This implies that we can obtain the density ratios between the neighboring pairs $p_{\rm{data}}(\mathbf{x}_{n_{1}})/p_{\rm{data}}(\mathbf{x})$ and $p_{\rm{data}}(\mathbf{x})/p_{\rm{data}}(\mathbf{x}_{n_{1}})$ given the Concrete score. Given a proposal distribution $q$, our sampler will accept the proposed update $\mathbf{x}^{\prime}$ with acceptance probability $A({\mathbf{x}}^{\prime}|{\mathbf{x}})=\min\left(1,\frac{p_{\rm{data}}({\mathbf{x}}^{\prime})q({\mathbf{x}}|{\mathbf{x}}^{\prime})}{p_{\rm{data}}({\mathbf{x}})q({\mathbf{x}}^{\prime}|{\mathbf{x}})}\right)$. Specifically, we can choose $q$ to sample uniformly among ${\mathcal{N}}(\mathbf{x})$ given a particular $\mathbf{x}$. When the neighborhood structure is symmetric, this selection leads to a simplified acceptance probability $A({\mathbf{x}}^{\prime}|{\mathbf{x}})=\min\left(1,\frac{p_{\rm{data}}({\mathbf{x}}^{\prime})}{p_{\rm{data}}({\mathbf{x}})}\right)$. We note that when the underlying graph is connected, the chain is aperiodic (due to the presence of a rejection step), irreducible, and positive recurrent (since there is a finite number of discrete states), so the Markov chain is guaranteed to converge in the limit to the model distribution.

This means that we can also compute tighter bounds on log-likelihood estimates obtained by un-normalized probability models trained via CSM. In particular, we can compute both a lower bound estimated via Annealed Importance Sampling (AIS) and a conservative upper bound estimated via the Reverse Annealed Importance Sampling Estimator (RAISE) . This is because both AIS and RAISE allow us to approximate the intractable normalizing constant by constructing a sequence of intermediate distributions between our estimated target distribution and another proposal distribution, and we know that the Concrete score can be repurposed to obtain the necessary density ratios between neighboring pairs of distributions along this sequence.

Learning Concrete Scores with Concrete Score Matching

In the following theorem, we demonstrate that Eq. 4 is indeed a principled learning objective: the estimated $\theta^{*}$ is a consistent estimator for the true underlying distribution $p_{\rm{data}}({\mathbf{x}})$.

In the limit of infinite data and infinite model capacity, the optimal $\theta^{*}$ that minimizes Eq. 4 recovers the true Concrete score, or satisfies ${\bm{c}}_{\theta^{*}}(\mathbf{x};{\mathcal{N}})={\bm{c}}_{p_{\rm{data}}}(\mathbf{x};\mathcal{N})$.

Theorem 1 implies that the estimated (underlying) $p_{\theta^{*}}(\mathbf{x})$ from $c_{\theta^{*}}({\mathbf{x}};{\mathcal{N}})$ satisfies $p_{\theta^{*}}(\mathbf{x})=p_{\rm{data}}(\mathbf{x})\textrm{ }\forall\mathbf{x}\in{\mathcal{X}}$. Next, we note that Eq. 4 can be simplified to an objective that we can tractably optimize by removing its dependence on the unknown ${\bm{c}}_{p_{\rm{data}}}({\mathbf{x}};{\mathcal{N}})$, similar in spirit to Eq. 1. Although the added flexibility of the score model may potentially lead to situations where it does not correspond to a valid probability distribution—similar to how continuous score models do not necessarily form a conservative vector field—this does not appear to hurt CSM’s performance empirically.

Optimizing Equation 4 is equivalent to optimizing:

where $\mathcal{N}(\mathbf{x})=\{\mathbf{x}_{n_{1}},...,\mathbf{x}_{n_{k}}\}$ is the set of neighbors of $\mathbf{x}$.

2 Efficient Training via Monte Carlo

In practice, Eq. 5 is still inefficient; it involves a summation over all $|{\mathcal{N}}({\mathbf{x}})|$ neighbors for ${\mathbf{x}}$, which can be expensive for high-dimensional datasets. Therefore, we propose several modifications to the original training objective that we found to be crucial for scaling up and improving our approach.

We begin with the leftmost term ${\mathcal{J}}_{1}$. First, we approximate the outer expectation w.r.t. $p_{\rm{data}}(\mathbf{x})$ with Monte Carlo samples from the empirical data distribution. Next, rather than evaluating the objective for all $|\mathcal{N}(\mathbf{x})|$ neighbors, we sample a neighbor ${\mathbf{x}}_{n_{i}}$ uniformly at random among ${\mathcal{N}}({\mathbf{x}}_{i})$ for every ${\mathbf{x}}_{i}$ in a mini-batch. We then upweight the output from the model using this sample by $|{\mathcal{N}}({\mathbf{x}})|$. We provide the unbiased training algorithm for the leftmost term in Algorithm 1. We note that is similar in spirit to sliced score matching , an approximation technique to make continuous score matching scalable to higher dimensions.

Next, we turn to the rightmost term ${\mathcal{J}}_{2}$. Similar to ${\mathcal{J}}_{1}$, we can estimate ${\mathcal{J}}_{2}$ using an unbiased estimator based on samples as demonstrated in Algorithm 2. We define ${\mathcal{N}}^{-1}(\mathbf{x}^{\prime})=\{(\mathbf{x},i)|{\mathcal{N}}(\mathbf{x})_{i}=\mathbf{x}^{\prime}\}$ as the “reverse neighborhood” set, where an element $(\mathbf{x},i)\in{\mathcal{N}}^{-1}(\mathbf{x}^{\prime})$ indicates that $\mathbf{x}^{\prime}$ is the $i$-th neighbor of $\mathbf{x}$. There could be multiple $\mathbf{x}$ with the same $i$ in ${\mathcal{N}}^{-1}(\mathbf{x}^{\prime})$ as in the case of the star graph. Computing and storing ${\mathcal{N}}^{-1}$ as a hash table mapping $\mathbf{x}$ to the set of ${\mathcal{N}}^{-1}(\mathbf{x})$ has a time and space complexity of at most $O(\sum_{\mathbf{x}}|{\mathcal{N}}(\mathbf{x})|)$ (i.e., the number of edges). For special structures (e.g., grid), we can design specific algorithms that bypass this process. We provide more details in Appendix C.

We leverage the grid structure in our experiments, where the neighbors of ${\mathbf{x}}$ are defined to be ${\mathcal{N}}({\mathbf{x}})=\{{\mathbf{x}}\pm{\mathbf{e}}_{d}\}_{d=1}^{D}$. Algorithm 1 and Algorithm 2 allow us to efficiently train the model by sampling a dimension $d\in[D]$ and flipping the bit for the $d$th entry of ${\mathbf{x}}$ for each ${\mathbf{x}}_{n}$.

3 Denoising Concrete Score Matching

Experimental Results

In this section, we evaluate the performance of CSM on synthetic, tabular, and discrete image datasets on a variety of sampling and density estimation tasks. We provide additional details on specific experimental settings and hyperparameter configurations in Appendix C.

In our experiments, we compare CSM to two relevant baselines for modeling discrete data: Ratio Matching and Discrete Score Matching with Marginalization (Discrete Marginalization). For conciseness, we provide additional details and derivations for their training objectives in Appendix D.

Ratio Matching. Similar to score matching, Ratio Matching leverages the fact that ratios of probabilities are independent of the intractable normalizing constant (due to cancellation). The method then seeks to match the ground truth density ratios $\frac{p_{\rm{data}}({\mathbf{x}})}{p_{\rm{data}}({\mathbf{x}}_{-d})}=\frac{q_{\theta}({\mathbf{x}})}{q_{\theta}({\mathbf{x}}_{-d})}$ where ${\mathbf{x}}_{-d}$ denotes the vector ${\mathbf{x}}$ with the $d$th entry bit-flipped (e.g. from 0 to 1).

Discrete Score Matching with Marginalization (Discrete Marginalization). The Discrete Marginalization baseline is another way of estimating discrete probability distributions with score matching as in . However, we note that this approach is difficult to scale because it requires us to marginalize over the data dimension, which may also cause instabilities during training.

2 1-D Discrete Data

In this experiment, we consider a 16-category 1-D data distribution as shown in Figure 2. We parameterize our Concrete score model ${\bm{c}}_{\theta}({\mathbf{x}},{\mathcal{N}})$ using a shallow MLP model with Tanh activations, and use the Cycle neighborhood structure during training. We demonstrate that the learned Concrete score model can faithfully capture the true data distribution when trained with CSM. As shown on the left side of Figure 2, we find that CSM generates samples (blue) using MH that almost perfectly matches the samples drawn from $p_{\rm{data}}({\mathbf{x}})$ (green).

We also observe in Appendix A.3 that the Concrete score on $p_{\rm{data}}(\mathbf{x})$ can recover the Stein score of a data distribution perturbed with triangular noise. Even though the Concrete score model is trained on discrete data (green), this connection allows us to sample with Langevin dynamics using the recovered Stein score. On the right side of Figure 2, we show how the samples generated using Langevin dynamics (smoothed orange) closely match the triangular noise-perturbed data distribution (gray). Then, we leverage a closed-form denoising formula for the perturbed distribution (Appendix A.3) to recover the clean data distribution (green) from samples obtained via Langevin dynamics (smoothed orange). The resulting denoised samples (orange) are labeled as Denoised in Figure 2. We provide more details and discussion in Section A.3.

3 Sampling with Toy 2-D Multi-Class Datasets

Next, we consider three 2-D toy benchmark datasets with multiple modes and discontinuities as commonly used in the density estimation literature . We quantize the data into $91\times 91$ bins to obtain the discrete training data for the experiment (see Figure 3). We compare our method against the two baselines: (1) Ratio Matching ; and (2) Discrete Marginalization . In particular, we found that the original equations in were incorrect, and provide results using the correct training objectives (denoted as Ratio-fixed and Marginal-fixed) in Figure 3. We provide a step-by-step derivation addressing this issue with the correct expressions for the training objectives in Appendix D.

For a fair comparison, we use the same model architecture and training configurations across all methods. We use the grid neighborhood structure for training CSM. As shown in Figure 3, the samples from CSM best capture the shape of the underlying data distributions across all three datasets (leftmost column). Baselines implemented using the original equations in indeed demonstrate poor performance on the density estimation task (Ratio and Marginal columns) in Figure 3.

4 Likelihood Evaluation on Discrete Tabular Data

We then evaluate the performance of CSM on density estimation tasks for a wide range of tabular (discrete) datasets drawn from both the Twenty Datasets and the Amazon Baby Registries benchmarks . In order to evaluate likelihoods, we directly parameterize the density with a discrete autoregressive model (MADE ), but train the model using the baseline approaches and CSM. For a fair comparison, we use the same model architecture and experimental configurations across all methods. Similar to our previous experiments, we use the grid neighborhood structure for training CSM. As shown in Table 1, our approach (CSM) demonstrates favorable performance relative to the Ratio Matching and Discrete Marginalization baselines. We provide additional experimental results in Appendix B.1 and more experimental details in Appendix C.

5 Generative Modeling with High-dimensional Images

For our final experiment, we demonstrate that we can scale CSM to achieve good performance on complex, high-dimensional binarized image datsets. We experiment with the MNIST dataset, which has 784 dimensions. We directly parameterize the Concrete score function with a U-Net , and train the model using the CSM based on a grid neighborhood structure. Similar to , we perturb the image with different levels of discrete (Categorical) noise, and train the models at different noise levels with annealing. After the Concrete score model is trained, we sample from the model using Metropolis-Hastings as discussed in Section 3.3 and follow a similar sample initialization process as . We provide more details in Appendix C. We present the samples from our model in Figure 4(b). We observe that our CSM approach with Metropolis-Hastings is able to generate samples that look similar to the digit examples in the training set (Figure 4(a)).

Related Work

Score Matching. Our work builds on the body of literature on score matching and score-based generative modeling . Notably, we build upon the generalization of score matching to discrete data . Our score function can be viewed through the lens of generalized score matching as in , where the Concrete score serves as a particular instantiation of the linear operator in their framework. However, we provide a novel perspective in terms of constructing the Concrete score via surrogate gradient information over a structured discrete space, and provide efficient training objectives for CSM. Our work is also closely related to score matching with finite difference (FD-SM) as discussed in Section 3.2. However, rather than approximating the gradient with finite differences for continuous data, we leverage the first-order forward difference to construct our Concrete score function in discrete domains. In addition, CSM bears similarities to a concurrent work on scaling up Ratio Matching for discrete energy-based models (EBMs) called RMwGGIS—our approach can be viewed as a different way to train discrete EBMs (via the Concrete score).

Generative Modeling for Discrete Data. CSM also provides another way to train generative models for discrete data. A related work is , which trains a score-based model on distributions over graphs. However, they perturb the adjacency matrices with Gaussian noise and use the continuous variant of DSM. Although there exist several model families for learning binary and Categorical probability distributions such as normalizing flows , Sum-Product Networks (SPNs) , denoising diffusion probabilistic models , discrete (latent) EBMs , and Generative Flow Networks (GFNs) , there does not yet exist a discrete score-based model that can scale to high-dimensional discrete datasets. Finally, our approach bears similarities with Gibbs with Gradients (GWG) . In particular, GWG augments existing MCMC samplers with gradient information by leveraging local structure to estimate likelihood ratios for transitioning to the next state. However, GWG utilizes gradients from the probability mass functions of the underlying discrete distributions, in contrast to the way in which we construct the Concrete score function.

Conclusion

We introduced Concrete Score Matching (CSM), a novel framework for learning discrete probability distributions via score matching. We proposed to leverage particular kinds of structural information in the data to construct surrogate gradient information about the discrete space, and used this to define a valid score function. We also introduced several modifications to the original training objective that allowed CSM to scale gracefully to high-dimensional datasets, and demonstrated that CSM performs well on a variety of sampling and density estimation tasks.

However, this work is not without limitations. Since CSM depends on the neighborhood structure, certain types of graphs may work better for score matching in practice than others. Additionally, our efficient training objectives may suffer from high variance when scaling to high dimensions for particular kinds of neighborhood-induced graphs, such as the Star graph. For future work, although we fixed the number of neighbors $K$ for each ${\mathbf{x}}$ in our experiments, it would be interesting to adaptively determine the optimal number of neighbors to use for each ${\mathbf{x}}$ during training. We also believe that CSM can be generalized to leverage more complex neighborhood structures than those we explored in the paper. Additionally, empirically investigating the performance of D-CSM with Langevin dynamics relative to CSM would be interesting future work.

This work introduces a novel score function—the Concrete score—that is defined over discrete spaces, as well as a corresponding score matching (CSM) framework that can scale to high-dimensional discrete datasets. This leads to empirical performance improvements over a range of density estimation and sampling tasks, and does not have a direct consequence on societal issues. However, we note that CSM could serve as the basis for developing more powerful generative models of structured, discrete data. Although this could lead to tangible benefits (e.g. improved generative modeling of text data), we should be mindful to take the usual precautions required for generative modeling research (e.g. against the development of deepfakes).

Acknowledgements and Funding Disclosure

We thank the anonymous reviewers for insightful discussions and feedback. KC is supported by the Qualcomm Innovation Fellowship and the Two Sigma Diversity PhD Fellowship. This research was supported by NSF (#1651565), AFOSR (FA95501910024), ARO (W911NF-21-1-0125), ONR, DOE, CZ Biohub, and Sloan Fellowship.

References

Checklist

Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes]

Did you describe the limitations of your work? [Yes] Section 7

Did you discuss any potential negative societal impacts of your work? [Yes] Section 7

Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]

If you are including theoretical results…

Did you state the full set of assumptions of all theoretical results? [Yes] Sections 3-4.1

Did you include complete proofs of all theoretical results? [Yes] Appendix A

Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [No] We will release the code publicly upon publication.

Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Section 5 and Appendix C

Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No]

Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] Appendix C

If you are using existing assets (e.g., code, data, models) or curating/releasing new assets…

If your work uses existing assets, did you cite the creators? [Yes] Section 5

Did you mention the license of the assets? [No] We only used publicly available datasets for our experiments.

Did you include any new assets either in the supplemental material or as a URL? [N/A]

Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [N/A]

Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]

If you used crowdsourcing or conducted research with human subjects…

Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]

Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]

Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]

Appendix

Appendix A Proofs of Theoretical Results

In this section, we provide proofs and additional discussions of the theoretical results.

If two nodes ${\mathbf{x}}$ and ${\mathbf{x}}^{\prime}$ in a neighborhood-induced graph ${\mathcal{G}}$ share an edge, then their density ratio $p_{\rm{data}}({\mathbf{x}}^{\prime})/p_{\rm{data}}({\mathbf{x}})$ and $p_{\rm{data}}({\mathbf{x}})/p_{\rm{data}}({\mathbf{x}}^{\prime})$ can be uniquely identified using ${\bm{c}}_{p_{\rm{data}}}(\mathbf{x};{\mathcal{N}})$ based on the definition:

For simplicity, we name the elements in $\mathcal{X}$: $\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{N}\}$. Note that when ${\mathcal{G}}$ is a weakly connected graph, any node pair $(\mathbf{x}_{1},\mathbf{x}_{n})$ is connected via a path in the graph. Denoting the path between these two nodes as $\mathbf{x}_{1}\to...\to\mathbf{x}_{n}$, we can obtain the density ratio for any neighboring pairs on the path using the definition of ${\bm{c}}_{p_{\rm{data}}}(\mathbf{x};{\mathcal{N}})$. Therefore the density ratio $p_{\rm{data}}({\mathbf{x}}_{n})/p_{\rm{data}}({\mathbf{x}}_{1})$ can be computed using the products of the density ratios for all neighboring data pairs on the path. This implies that when ${\mathcal{G}}$ is weakly connected, we can uniquely recover the density ratios {$p_{\rm{data}}({\mathbf{x}}_{1})/p_{\rm{data}}({\mathbf{x}}_{1}),p_{\rm{data}}({\mathbf{x}}_{2})/p_{\rm{data}}({\mathbf{x}}_{1}),\ldots,p_{\rm{data}}({\mathbf{x}}_{N})/p_{\rm{data}}({\mathbf{x}}_{1})\}$, which uniquely determine $p_{\rm{data}}(\mathbf{x})$. Therefore, knowing the density ratio between any two points uniquely identifies $p_{\rm{data}}(\mathbf{x})$.

It is easy to see the optimal $\theta^{*}$ that minimizes the following equation

satisfies ${\bm{c}}_{\theta^{*}}(\mathbf{x};{\mathcal{N}})={\bm{c}}_{p_{\rm{data}}}(\mathbf{x};\mathcal{N})$ almost everywhere.

A.2 Denoising Concrete Score Matching

The model ${\bm{c}}_{\theta}$ that minimizes the least squares

A.3 Connection to distribution perturbed with triangular noise

We define the PDF of the $D$-dimensional triangular distribution with lower limit -1, upper limit 1, and mode 0 as the following:

Using the fact that $\mathcal{T}_{D}$ is independent among dimensions, we have

Plugging in Equation 16 and Equation 17, we have

A.4 Connection to Stein Scores

In this section, we study a special case of the Concrete score, as described in Proposition 1, to highlight its connection to both the continuous (Stein) score and existing score matching methods. Given a $D$-dimensional continuous data distribution $p(\mathbf{x})$ and $\delta>0$, we define a particular neighborhood structure $\mathcal{N}({\mathbf{x}})=\{{\mathbf{x}}_{n_{i}}\}_{i=1}^{D}$ where $\mathbf{x}_{n_{i}}=p(\mathbf{x}+\delta{\mathbf{e}}_{i})$ and ${\mathbf{e}}_{i}$ is the standard (one-hot) basis vector with the $i$-th element ${\mathbf{e}}_{i}=1$. Then:

From Eq. 21, we can make two observations. First, we see that \bigg{[}\frac{p(\mathbf{x}+\delta{\mathbf{e}}_{1})-p(\mathbf{x})}{\delta},...,\frac{p(\mathbf{x}+\delta{\mathbf{e}}_{D})-p(\mathbf{x})}{\delta}\bigg{]}^{T} approximates the directional derivative of $p(\mathbf{x})$ via a first-order forward difference operation. As a result, the scaled Concrete score function $\frac{{\bm{c}}_{\theta}(\mathbf{x},{\mathcal{N}})}{\delta}$ converges to ${\bm{s}}({\mathbf{x}})$ in the limit of $\delta\to 0$:

which is precisely the definition of the Stein score.

Appendix B Additional Experimental Results

We present additional results from Section 5.4, where we trained a MADE model using standard maximum likelihood. We use the training setting as detailed in Appendix C. As shown in Table 2, this maximum likelihood baseline serves as an upper bound on performance across all score-based methods. We note that as in continuous (Stein) score matching, log-likelihoods and score matching losses are not always correlated even though they both theoretically converge to the optimal solution given infinite model capacity . This is due to practical constraints such as model mis-specification or optimization challenges. We believe that this is also the case for discrete score matching, which explains why directly optimizing likelihood outperforms all other approaches in Table 2.

B.2 Neighborhood Structure Specification in Practice

In theory, our approach can use any neighborhood structure as long as the neighborhood-induced graph is connected (see Theorem 1). In practice, the choice of neighborhood structure can affect performance due to challenges in optimization and proper model specification. We provide an empirical analysis to build intuition on a series of 1-D synthetic datasets, and believe that the same intuition can generalize to higher dimensions.

First, we consider the 1-D distribution in Figure 5, where the data distribution does not contain any low density regions (regions with close to zero density). We parameterize the unnormalized probability distribution with softmax logits, and train the model using CSM. For the training objective, we explore 4 different neighborhood structures (3 different cycles as well as a fully-connected graph) as shown in Figure 5. In this setting, we observe that different neighborhood structures yield similar performances as evaluated by log-likelihood (see Figure 5).

Next, we consider a different 1-D data distribution which contains low density regions. As shown in Figure 6, we observe that the neighborhood structure in this setting plays a critical role in the final log-likelihoods: the complete graph in Case 2 outperforms all other structures. Interestingly, when the low-density regions are in between sets of high-density modes (as in Case 3 in Figure 6), we find that the model can still perform comparably relative to Case 2.

This observation is similar to that of Stein score matching . Therefore, we believe that a similar noise annealing procedure followed by denoising CSM may potentially alleviate the effects of poorly chosen neighborhood structures for data distributions with low density regions. We will leave this investigation for future work. Such empirical results highlight that the best-performing structure in practice is often data-dependent. That is, a neighborhood graph which respects the particular dataset structure will perform the best out of all possible graph structures (analogous to the way in particular types of inductive biases are useful for solving specific tasks). One interesting avenue for future research would be to draw inspiration from and investigate whether it is possible to construct an “optimal” neighborhood structure for a given dataset. In this work, we uniformly sample a set of neighbors from ${\mathcal{N}}({\mathbf{x}})$ when computing the CSM objective in practice (Algorithms 1 and 2), which can lead to issues with high variance during training. We hypothesize that an optimal proposal distribution over ${\mathcal{N}}({\mathbf{x}})$ that minimizes variance may lead to insights on what the optimal neighborhood structure should be.

B.3 Denoising Concrete Score Matching

We present additional experiments on denoising concrete score matching (D-CSM) as introduced in Section 4.3. In particular, we consider two 2-D toy benchmark datasets as commonly used in the density estimation literature . We quantize the data into $91\times 91$ equally-distanced bins to obtain the discrete training data (see Clean data in Figure 7). Similar to our previous experiments, we use the grid neighborhood structure as detailed in Section 5.3 for training via D-CSM.

Appendix C Experimental Details

In this section, we provide more details for each of the experimental settings. We will release our code upon publication.

We discuss the experimental setting for Section 5.2. We use a 16-category 1-D data distribution with class labels from 0 to 15, and normalize the data to $$before feeding it into the model. We use a 3-layer shallow MLP for the score network with 100 hidden units and Tanh activations. We use the model to directly parameterize the Concrete score and define the neighborhood structure to be a Cycle structure as shown in Figure 8. We train the model using the Adam optimizer with learning rate$0.001$ for 100,000 iterations. The Langevin samples and denoised samples in Figure 2 are obtained using the approach detailed in Section A.3.

C.2 2-D Multi-Class Datasets

For the 2-D experiments, we consider three 2-D toy benchmark datasets with multiple modes and discontinuities as commonly used in the density estimation literature . We quantize the data into $91\times 91$ equally-distanced bins to obtain the discrete training data (see Figure 3). We use the grid neighborhood structure for training via CSM (see Figure 8). For a fair comparison, we use the same model architecture and training configurations across all methods. We directly parameterize the unnormalized probability distribution with softmax logits, but train the model using the baselines and CSM. Since there are $91\times 91$ possible classes, our model directly parameterizes an unnormalized distribution with $91\times 91$ possible values. We initialize the model to have a uniform probability across classes and train the model using the Adam optimizer with learning rate 0.0005. To sample from the model, we can either use likelihood sampling (by normalizing the model) or Metropolis-Hastings. We empirically observe that both approaches work well.

C.3 Discrete Tabular Data

For the tabular experiments, we consider tabular (discrete) datasets drawn from both the Twenty Datasets and the Amazon Baby Registries benchmarks . In order to evaluate likelihoods, we directly parameterize the probability distribution with a discrete autoregressive model (MADE ), but train the model using the baseline approaches and CSM. The MADE model uses 2 residual blocks with latent dimension 100 and Tanh activations. We train the model for 50000 iterations using the Adam optimizer with learning rate 0.0005. We use a grid neighborhood structure for CSM (see Figure 8). For a fair comparison, we use the same model architecture and experimental configurations across all methods. Similar to our previous experiments, we use the grid neighborhood structure for training via CSM (see Figure 8).

C.4 Discrete Image Data

We use Metropolis-Hastings for sampling from the model after training: in particular, we initialize the sample to be drawn from uniform binary noise then use the model corresponding to noise level $\sigma_{7}$ to perform Metropolis-Hastings updates. After that, we use the output from the model to initialize the input for sampling from the model with noise level $\sigma_{6}$. We then repeat the procedure and initialize the model corresponding to the previous noise level with the output from the next noise level—a process similar to . We observe that each noise level requires around 100-800 steps to obtain reasonable results.

C.5 Training for Special Neighborhood Structures

In this section, we provide additional details on special neighborhood structures where we can design even more efficient training procedures than using Algorithm 1 and Algorithm 2. For simplicity, we name the elements in $\mathcal{X}$: $\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{N}\}$.

When the neighborhood-induced graph is a Chain (Figure 8), the second term ${\mathcal{J}}_{2}$ in Equation 5 simplifies to the following via reparameterization:

Cycle.

Similarly, when the neighborhood-induced graph is a Cycle (Figure 8), the second term ${\mathcal{J}}_{2}$ becomes:

Thus for both Chains and Cycles, the term ${\mathcal{J}}_{2}$ in Equation 5 can be evaluated efficiently using samples from $p_{\rm{data}}({\mathbf{x}})$.

Grid.

When the neighborhood-induced graph is a Grid (Figure 8), the second term ${\mathcal{J}}_{2}$ can be approximated efficiently by decomposing each dimension into smaller Cycles. Specifically, to estimate ${\mathcal{J}}_{2}$, we can uniformly sample a dimension $d$, and then use the same reparameterization approach used for the Cycle structure at each dimension.

Appendix D Additional Related Works and Corrections

In this section, we provide additional details about the baseline methods: in particular, we found that the original equations in were incorrect. We elaborate upon them below.

We assume with slight abuse of notation that ${\mathcal{X}}\in\{0,1\}^{D}$, and $p_{\rm{data}}({\mathbf{x}})$ is the unknown discrete data distribution. Ratio matching was originally proposed as an alternative approach to score matching for learning discrete binary probability distributions from samples. Similar to score matching, it leverages the fact that ratios of probabilities are also independent of the intractable normalizing constant (due to cancellation), and seeks to match the ground truth density ratios $\frac{p_{\rm{data}}({\mathbf{x}})}{p_{\rm{data}}({\mathbf{x}}_{-d})}=\frac{q_{\theta}({\mathbf{x}})}{q_{\theta}({\mathbf{x}}_{-d})}$ where ${\mathbf{x}}_{-d}$ denotes the vector ${\mathbf{x}}$ with the $d$th entry bit-flipped (e.g. from 0 to 1). Concretely, ratio matching minimizes the following objective:

As the original ratio matching approach is proposed for binary data, extends ratio matching to multi-class data. Following the notation in , let $\mathbf{x}^{\backslash i}$ denote the vector formed by dropping the $i$-th element (which we call $\mathbf{x}_{i}$) from $\mathbf{x}$. To make things concrete, this means that $p({\mathbf{x}})=p(\mathbf{x}^{\backslash i},\mathbf{x}_{i})$. We also let $q_{\theta}(\xi_{i}|\mathbf{x}^{\backslash i})$ denote the conditional probability for the $i$-th index of $\mathbf{x}$ taking the value $\xi_{i}$, and let $q_{\theta}(\sim\xi_{i}|\mathbf{x}^{\backslash i})$ denote the conditional probability for the $i$-th index of $\mathbf{x}$ not taking the value $\xi_{i}$.

The generalized ratio matching objective proposed in is:

It is easy to see that the optimal $\theta^{*}$ for Equation 27 can be achieved independent of $p_{\rm{data}}(\mathbf{x})$, which is problematic. We also empirically show that model trained with Equation 27 cannot learn meaningful features based on the data (see Ratio in Figure 3).

As an alternative, we provide a corrected multi-class ratio matching objective. Similar to , we assume that all the probabilities are non-zero. For simplicity, we use $p(\mathbf{x})$ to denote $p_{\rm{data}}(\mathbf{x})$. The objective for multi-class ratio matching can achieved by optimizing:

We provide the samples trained with Equation 28 in Figure 3 (Ratio-fixed). We observe that the samples appear to be more reasonable then those from a model trained with Equation 27 (Ratio).

D.2 Discrete Marginalization

Given a multi-class discrete data distribution $p_{\rm{data}}(\mathbf{x})$, discrete marginalization proposes to learn the data distribution by minimizing

where $\mathcal{M}_{i}$ is the marginalization operator defined as $\mathcal{M}_{i}:\mathcal{F}^{1}\to\mathcal{F}^{1}:f(\mathbf{x})\to\int_{\mathbf{x}_{i}}f(\mathbf{x})d\mathbf{x}_{i}$ or $\sum_{\mathbf{x}_{i}}f(\mathbf{x})$ in the discrete case. Thus $\frac{\mathcal{M}_{i}p}{p}=\frac{1}{p_{\rm{data}}(\mathbf{x}_{i}|\mathbf{x}^{\backslash i})}$ and $\frac{\mathcal{M}_{i}q_{\theta}}{q_{\theta}}=\frac{1}{q_{\theta}(\mathbf{x}_{i}|\mathbf{x}^{\backslash i})}$.

The authors further simplified Equation 30 to

However, we observe that this is a typo in Equation 31—the optimal $\theta^{*}$ can be achieved while being independent of $p_{\rm{data}}(\mathbf{x})$. In our experiments, we also observe that model trained with Equation 31 cannot learn meaningful representations of the data (see Marginal in Figure 3).

In the following, we provide the corrected simplified version of Equation 30.

Optimizing Equation 30 is equivalent to optimizing

In the last equation, we used the fact that $\sum_{\mathbf{x}}p_{\rm{data}}(\mathbf{x})\mathcal{M}_{i}(\frac{\mathcal{M}_{i}q_{\theta}}{q_{\theta}})=\sum_{\mathbf{x}}p_{\rm{data}}(\mathbf{x})\sum_{\xi_{i}}(\frac{\mathcal{M}_{i}q_{\theta}}{q_{\theta}})=\sum_{\mathbf{x}}p_{\rm{data}}(\mathbf{x})\sum_{\xi_{i}}\frac{1}{q_{\theta}(\xi_{i}|\mathbf{x}^{\backslash i})}$, which directly follows from Lemma 4 in .

We provide the samples trained with Equation 32 in Figure 3 (Marginal-fixed). We observe that the samples appear to be more reasonable then those from a model trained with Equation 31 (Marginal).