Minibatch Standard Deviation Layer

Question

I'm reworking some of the GANs I originally made in TensorFlow2 to see if I can improve performance in Mathematica, and have been stuck on how to create a custom Minibatch Standard Deviation Layer. I'm trying to implement it to stabilize the training process and reduce instances of Mode Collapse. (More information on its purpose (with Python) can be found here.) Basically it takes the mean of the samples and figures out the standard deviation within each batch. The generator is penalized if the standard deviation is too low.

I thought the new FunctionLayer in 12.2 would be useful for this purpose but this is as far as I've gotten so far. Some how I need to thread that minibatch layer back into the network.

discriminator = NetGraph[{
   ConvolutionLayer[16, 5, "Stride" -> 3, "Input" -> {3, 4000}],
   ElementwiseLayer["ELU"],
   ConvolutionLayer[32, 3, "Stride" -> 1],
   ElementwiseLayer["ELU"],
   ConvolutionLayer[64, 2, "Stride" -> 1],
   ElementwiseLayer["ELU"],
   minibatch,
   FlattenLayer[],
   100,
   Ramp,
   1,
   Tanh,
   SequenceLastLayer[]
   },
  {1 -> 2 -> 3 -> 4 -> 5 -> 6 -> 8 -> 9 -> 10 -> 11 -> 12 -> 13, 
   6 -> 7}
]

minibatch = NetChain[{
   AggregationLayer[Mean, {2}],
   FunctionLayer[StandardDeviation]
   }]

GAN = NetInitialize@NetGANOperator[{generator, discriminator}]

BATCHSIZE = 64;
noise := RandomVariate[NormalDistribution[0, 0.2], 400]

trained = NetTrain[
  GAN,
  <|"Sample" -> RandomSample[trainingData, BATCHSIZE], 
   "Latent" -> Table[noise, BATCHSIZE]|>,
  BatchSize -> BATCHSIZE
  ]

The Wolfram framework always seems to apply the whole net to one example. How can I apply a function to all samples in the round and then append it to the array? Even if it's just a reference, I'm pretty sure I can follow it through.

Updated 2020-12-18 Along with the bounty I have also added a 3 samples of the training data that can be seen on a public gist of the shape {3,300}.

Updated 2021-01-02 Corrected code.

Answer 1

I am not sure this is correct, but here's my attempt - perhaps it can get you started and/or inspire others to give a more correct answer.

My understanding is that we want access to the standard deviation of some features across the batches during training. BatchNormalizationLayer should have access to this internally - even though the only exposed ports are Input and Output. In-fact the Properties & Relations section under Documentation gives us a hint, that it actually computes:

batchNormFunction = 
  Function[Block[{sd = Sqrt[#MovingVariance + #Epsilon]},
    (#2 * #Scaling / sd ) + (#Biases - (#Scaling * #MovingMean)/sd)]];

As such, I think we can (ab)use the layer to extract the batch standard deviation as follows:

params = <|"Scaling" -> 1, "Biases" -> 0, "MovingMean" -> 0, 
   "MovingVariance" -> mv, "Epsilon" -> 0|>;
input/batchNormFunction[params, input]

Out[84]= Sqrt[mv]

Note this will fail if either "MovingVariance" or the input is zero. So you probably want to add a small amount of "Biases" and "Epsilon" to stabilize it. Here I define a 'frozen' BatchNormalizationLayer with no "Biases" and the default value of "Epsilon" (0.001) as follows:

bn = BatchNormalizationLayer["Scaling" -> 1, "Biases" -> 0, 
  "MovingMean" -> 0, "Input" -> {512, 4, 4}, 
  LearningRateMultipliers -> <|"Biases" -> 0, "MovingMean" -> 0, 
    "MovingVariance" -> 1, "Scaling" -> 0|>]

and define our minibatch layer using the NetGraph:

minibatch = 
 NetGraph[<|"bn" -> bn, "average" -> AggregationLayer[Mean, All], 
   "rep" -> ReplicateLayer[{1, 4, 4}], "cat" -> CatenateLayer[], 
   "divide" -> ThreadingLayer[Divide]|>, {NetPort["Input"] -> 
    "bn", {NetPort["Input"], "bn"} -> 
    "divide" -> "average" -> "rep", {NetPort["Input"], "rep"} -> 
    "cat"}]

Note that {512,4,4} above comes from the particular GAN I was playing around with, which I suspect will be easy enough to change. For reference here is the the GAN network - it seems to initialize properly, although I haven't trained on anything.

leaky = ParametricRampLayer["Slope" -> 0.2, 
   LearningRateMultipliers -> 0];
conv[channels_, kernel_] := 
 Sequence[ConvolutionLayer[channels, kernel, 
   PaddingSize -> (kernel + 1)/2 - 1], leaky]
deconv[channels_, kernel_] := 
 Sequence[DeconvolutionLayer[channels, kernel, 
   PaddingSize -> (kernel + 1)/2 - 1], leaky]
enc = NetEncoder[{"Image", {128, 128}}];
downSample[] := PoolingLayer[2, 2, "Function" -> Mean]
upSample[] := ResizeLayer[{Scaled[2], Scaled[2]}]

discriminator = 
 NetChain[{conv[16, 1], conv[16, 3], conv[32, 3], downSample[], 
   conv[32, 3], conv[64, 3], downSample[], conv[64, 3], conv[128, 3], 
   downSample[], conv[128, 3], conv[256, 3], downSample[], 
   conv[256, 3], conv[512, 3], downSample[], minibatch, conv[512, 3], 
   ConvolutionLayer[512, 4], leaky, LinearLayer[{}]}, "Input" -> enc]

generator = 
 NetChain[{DeconvolutionLayer[512, 4], leaky, deconv[512, 3], 
   upSample[], deconv[256, 3], deconv[256, 3], upSample[], 
   deconv[128, 3], deconv[128, 3], upSample[], deconv[64, 3], 
   deconv[64, 3], upSample[], deconv[32, 3], deconv[32, 3], 
   upSample[], deconv[16, 3], deconv[16, 3], 
   DeconvolutionLayer[3, 1]}, "Input" -> {512, 1, 1}]

NetInitialize[NetGANOperator[{generator, discriminator}]]