# Random sampling of a discrete distribution in NetGraph

Within a NetGraph I would like to use the output of a softmax layer, interpreted as a probability distribution, to sample a "next token". Since this is done inside an RNN structure, I cannot do this outside of the NetGraph. However, I have not found any reference as to how this might be done inside a NetGraph.

I would like to construct a NetGraph layer which takes as input an array p={p1,p2,...,pn} where Sum[p]=1, so that we can regard p as a probability distribution on Range[n]. The output of the layer should be an integer in Range[n] constructed by randomly sampling Range[n] according to the probabilities {p1,p2,...,pn}.

• Welcome to mathematica.stackexchange.com! It's not clear to me how this is related to random number generation in NetGraphs. Can you clarify with some more details in your question - especially with examples, if you can? – Carl Lange May 7 at 15:43
• Don't worry! As it is, though, your title doesn't seem to mesh with the text of your question, and more context would be useful for any potential answerer. In any case you might want to look at the NetDecoder documentation, specifically for Tokens. Or perhaps the documentation for SoftmaxLayer. – Carl Lange May 7 at 16:17
• I would like to construct a NetGraph layer which takes as input an array p={p1,p2,...,pn} where Sum[p]=1, so that we can regard p as a probability distribution on Range[n]. The output of the layer should be an integer in Range[n] constructed by randomly sampling Range[n] according to the probabilities {p1,p2,...,pn}. – Steven Hutt May 7 at 16:20
• I can see how a "Token" NetDecoder would do the job. As I am building an RNN structure using the NetFoldOperator, I will check whether NetDecoder can be part of the network being "folded". Thanks! – Steven Hutt May 7 at 16:23
• I don't think it will. Please add more to your question so I, and other answerers, have a better idea of what you're trying to do. – Carl Lange May 7 at 16:25

I figured out how to do this, although it's a bit hacky. It relies on the following fact. Suppose {g1,g2,...,gn} are n samples from a Gumbel distribution. Let z = {Log[p1]+g1, Log[p2]+g2, ..., Log[pn]+gn}. Then argmax(z) is a random sample from Range[n] according to the probabilities {p1,p2,...,pn}. Note also that I wasn't able to generate the Gumbel samples inside the net, so I have to feed them in as an input. Obviously every time you call the net, you need to pass in a new sample.

I borrowed the argMaxLayer from here (thanks Carl!). I had to modify slightly to output an integer rather than one-hot.

Here's the full code :

argMaxLayer[n_] := NetGraph[{
"max" -> AggregationLayer[Max, 1],
"repl" -> ReplicateLayer[Automatic],
"switch" -> ThreadingLayer[If[#1 != #2, 0., 1.] &],
"argMax" -> ConstantTimesLayer["Scaling" -> Range[n]],
"sum" -> AggregationLayer[Total, 1]},
{NetPort["Input"] -> "max",
"max" -> "repl",
{NetPort["Input"], "repl"} -> "switch",
"switch" -> "argMax",
"argMax" -> "sum"}]

n = 4;
input = {0.25, 0.25, 0.25, 0.25};
gumbel = -Log[-Log[RandomVariate[UniformDistribution[{0, 1}], 4]]];
data = <|"Input" -> input, "Gumbel" -> gumbel|>;
sampleDiscDistLayer =
NetGraph[{"Plus" -> Plus, "Log" -> Log, "ArgMax" -> argMaxLayer[n]},
{NetPort["Input"] -> "Log",
{"Log", NetPort["Gumbel"], "Log"} -> "Plus",
"Plus" -> "ArgMax"}];
sampleDiscDistLayer[data]


Repeated calls will uniformly sample integers from Range[n]. Just change the input probabilities to get different discrete distributions on Range[n].

(And apologies for answering my own question as someone else - the UX for this site is very confusing. Think I deleted my profile or something weird)