How to sample N-dimensional discrete data with K classes per dimension?

1

How can I create a torch.distribution class that samples an N-dimensional discrete space of K classes each?

I initially thought of using multinomial (docs: here), but that would require initializing an array, which scales exponentially with respect to the number of dimensions (i.e. the sampling array would be K^N in size).

If I fix this distribution to be equal in probability per class per dimension, is there an easy way to generate samples (as well as the probability, which should just be (1/K)^N for uniformly samples discrete values of K classes in N dimensions)

Any help is welcome!

2

Hi Alpha!

If I understand what you are looking for, I think multinomial() will
work just fine for you. As I understand it, you are not looking for any
kind of correlation between class values in different dimensions, so
if this is the case, you should be able to “factor” the random sampling
across dimensions.

Here is an example:

>>> import torch
>>> print (torch.__version__)
2.3.0
>>>
>>> _ = torch.manual_seed (2024)
>>>
>>> k = 5
>>> n = 10
>>>
>>> probs = (torch.ones (1, k) / k).expand (n, k)   # uniform class probabilities for all dimensions
>>>
>>> num_samples = 2
>>> torch.multinomial (probs, num_samples = num_samples, replacement = True)
tensor([[1, 3],
        [1, 3],
        [4, 1],
        [1, 2],
        [2, 0],
        [0, 1],
        [4, 3],
        [3, 1],
        [4, 2],
        [2, 4]])

Each column in the result tensor is a sample – in this example, there
are two samples. As you note, there are k**n distinct (vector-valued)
values for a sample and when feeding uniform probabilities to
multinomial(), each such value will be sampled with probability
1 / k**n.

Note that this illustration does not depend on the probabilities being
uniform across classes or dimensions. For arbitrary probabilities, you
do need to materialize the full [n, k] probability matrix, but that scales
as n*k, rather than exponentially.

If this is not what you are looking for, could you post some concrete
code that generates your desired result (even if it doesn’t scale well)?

Best.

K. Frank

1 Like
4

Hi @KFrank,

This looks like what I was looking for. My idea case would be to sample an arbitrary N-dim discrete data distribution (via some kind of Normalizing Flow perhaps), but I need a base distribution like the one you’ve provided.

But the only issue is that Normalizing Flows are designed for continuous distributions (rather than discrete distributions, so I’m not sure if it’s feasible to use this as a base distribution for a flow).