Random sampling of a discrete distribution in NetGraph

Question

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}.

Answer 1

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)