Neural network layer for vector normalization

Question

We have a neural network that returns a vector $v$ with length $n$. However, these results $v_i$ must be normalized to unity so that their sum $\sum_{i=1} ^{n} v_i=1$. It must be done inside the network to make its training go well.

Which kind of layer can do something like:

1. Take a vector $v$ with length $n$
2. Calculate the sum $S$ of all components $v_i$ of the vector $v$ $$S=\sum_{i=1} ^{n} v_i$$ (its total)
3. Divide all the components $v_i$ of the vector $v$ by the calculated sum $S$ $$v_i '=v_i/S $$
4. Return the new vector $v'=(v_i ')_{i=1} ^n$ with the same length $n$.

The sum can be obtained by SummationLayer[] operating on the entire vector $v$. Division could be obtained by ElementwiseLayer[(#/S)&]. But how do you deliver the divisor (sum) $S$ from SummationLayer to ElementwiseLayer?

Maybe there is another way to achieve this normalization ...

Answer 1

You can build this function yourself using SummationLayer, ReplicateLayer and ThreadingLayer.

v = {1, 2, 4, 5, 6}

normalizeLayer = NetGraph[{
   SummationLayer[],
   ReplicateLayer[Automatic],
   ThreadingLayer[#1/#2 &]
   }, {NetPort@"Input" -> 1 -> 2, {NetPort@"Input", 2} -> 3}]

normalizeLayer[v]

{0.0555556, 0.111111, 0.222222, 0.277778, 0.333333}

That is, sum the input, replicate the sum Automatically to have the same dimensions as the input, and divide the input by it.

(In fact, the ReplicateLayer[Automatic] gets its dimensions from ThreadingLayer - Wolfram Language takes care of resolving these dimensions for you. The ThreadingLayer knows that it needs dimensions that are the same as its first input, and asks the ReplicateLayer to replicate its input to those dimensions. It's very handy.)

SoftmaxLayer is very related here:

SoftmaxLayer[][v]

{0.00440888, 0.0119846, 0.0885547, 0.240717, 0.654335}

But as you can see, the results are slightly different - SoftmaxLayer doesn't do exactly what you're looking for (it normalizes the exponential of the input, rather than the raw input). In both cases the output sums to 1.

To build a SoftmaxLayer from scratch, you can modify the above normalizeLayer very slightly, to get the sum of the exponential and dividing the exponential of the input by it:

ourSoftmaxLayer = NetGraph[{
   ElementwiseLayer[Exp@# &],
   SummationLayer[],
   ReplicateLayer[Automatic],
   ThreadingLayer[Exp@#1/#2 &]
   }, {NetPort@"Input" -> 1 -> 2 -> 3, {NetPort@"Input", 3} -> 4}]

ourSoftmaxLayer[{1, 2, 4, 5, 6}]

{0.00440888, 0.0119846, 0.0885547, 0.240717, 0.654335}

Note that there is a small error between our softmax and the SoftmaxLayer output:

ourSoftmaxLayer[{1, 2, 4, 5, 6}] - SoftmaxLayer[]@{1, 2, 4, 5, 6}

{-4.65661·10^-10, 0., -7.45058·10^-9, -1.49012·10^-8, -5.96046·10^-8}

I attribute this to some rounding errors.

We can visualize the difference between the normalization function in the question and a softmax layer:

BarChart[Transpose@{normalizeLayer[v], SoftmaxLayer[][v]},
   ChartLabels -> {"norm", "soft"}]

In a future version of Wolfram Language (perhaps 12.1), the upcoming NetFunction functionality may make this slightly easier, allowing you to do something along the lines of NetFunction[u, (#/Total@u &) /@ u], but we shall see.