Custom SGD optimizer in Mathematica neural network framework?

Question

I have a new approach to SGD optimizer, and thought to try to test it in Mathematica. It uses gradients to simultaneously maintain online parabola model for smarter choice of step size - I only need to ask for gradients and be able to manually update parameters.

Is it doable within Mathematica neural network framework? I see NetPortGradient, but how to modify parameters?

As it seems a common research direction, maybe there is some simple example?

Answer 1

I'm not sure I understand OP's question fully and/or correctly, but according to my understanding it's a nice question that I asked myself before, so here I want to share the way I did when I wanted to try my own optimization strategy.

---

According to Deep Learning - Algorithm 8.1, SGD updates parameters by

$$\boldsymbol\theta \leftarrow \boldsymbol\theta - \epsilon_k \, \hat{\boldsymbol g}\;\text{,}$$

where $\epsilon_k$ are the learning rate schedule.

Now in Mathematica's NetTrain, when option

Method -> { "SGD", "LearningRate" -> r, "LearningRateSchedule" -> f }

is used along with option

LearningRateMultipliers -> { "layer1" -> λ1, "layer2" -> λ2, ... }

, I believe basically we have

$$\epsilon_k := r\,\lambda_k\,f(\#_\text{current batch},\#_\text{total batch})\;\text{.}$$

So as long as our strategy follows Algorithm 8.1 in general, any customization should be done against $\epsilon_k$, thus we can fully customize it by specifying r, f and λs.

The important thing to remember about options like "LearningRateSchedule" and TrainingProgressFunction is they are Functions, so they get re-evaluated every time NetTrain invoking them. So other than their explicit inputs, like $\#_\text{current batch}$ and $\#_\text{total batch}$ for f, we can inject any variables/values as long as they are available inside scope of those functions. The most naive way to achieve that is by global (to NetTrain) variables. Through that "injection", we can do arbitrary computation between any two batches despite the designed purposes of TrainingProgressFunction etc.

Another thing to note is TrainingProgressFunction is NOT just to "monitor progress". As inside its scope we have access to lots of "runtime properties" like BatchLossList, "Gradients" etc., and we already mentioned that we can do arbitrary computing inside that scope, so surely we can do any analysis against the provided runtime properties, encode the result to some global variables, which then can be extracted inside "LearningRateSchedule" to adjust $\epsilon_k$ accordingly.

So here is the propose:

1. Setup some global variable, say we name it multiplier.
2. Inside TrainingProgressFunction, modify multiplier according to arbitrary analysis we like against, say, #BatchLossList.
3. Inside Method, we use multiplier like "LearningRateSchedule" -> Function[multiplier]

Demo

An example from NetTrain's doc:

data = Flatten[Table[{x, y} -> Exp[-Norm[{x, y}]], {x, -3, 3, .005}, {y, -3, 3, .005}]];
net = NetChain[{32, Tanh, 1}];

trained1 is the result of the usual way:

trained1 = NetTrain[net, data, MaxTrainingRounds -> 24, BatchSize -> 1024, Method -> "SGD"];

trained2 is the result of our customized SGD, where we repeatedly fit a linear model against latest BatchLossList, increase/decrease our multiplier if the trend of recent loss is too "flat"/"steep":

DynamicModule[{\[ScriptK]Val, \[ScriptQ]Val, multiplier = 1, reset = 1, btV, btlV},
    Module[{maxround = 24, \[ScriptB], \[ScriptK], \[ScriptQ], regLen = 100, multiplierMin = 0.01},
        PrintTemporary[{
                    ToString[Abs@#, StandardForm] & /@ {"\[ScriptK]Val", "\[ScriptQ]Val"}
                    , Dynamic[Abs@#] & /@ {\[ScriptK]Val, \[ScriptQ]Val}
                    } // Grid[#, Frame -> All] &]
        ; PrintTemporary[Row[{"multiplier: ", Dynamic[multiplier]}]]
        ; PrintTemporary[Grid[{{"Batch", "Length[BatchLossList]", "reset"}, {Dynamic[btV], Dynamic[btlV], Dynamic[reset]}}, Frame -> All]]
        ;
        trained2 = NetTrain[net, data
                , MaxTrainingRounds -> maxround, BatchSize -> 1024
                , Method -> {"SGD", "LearningRateSchedule" -> (multiplier &)}
                , TrainingProgressFunction -> {
                        Function[
                            btV = #Batch; btlV = Length[#BatchLossList]
                            ; If[And[btlV > reset + regLen
                                    , #BatchLossList[[-regLen ;;]] // RightComposition[
                                            Log
                                            , FindFit[#, \[ScriptQ] \[ScriptI]^2 + \[ScriptK] \[ScriptI] + \[ScriptB], {\[ScriptB], \[ScriptK], \[ScriptQ]}, \[ScriptI]] &
                                            , ({\[ScriptK]Val, \[ScriptQ]Val} = {\[ScriptK], \[ScriptQ]} /. #) &
                                            ]
                                    ; Abs[\[ScriptQ]Val] < 10^-4
                                    ]
                                ,
                                Which[
                                    Abs[\[ScriptK]Val] > 2 10^-3
                                    , multiplier = Clip[.75 multiplier, {multiplierMin, 1}]
                                    ; reset = btlV
                                    ,
                                    Abs[\[ScriptK]Val] < 1 10^-3
                                    , multiplier = Clip[1.5 multiplier, {multiplierMin, 1}]
                                    ; reset = btlV
                                    ]
                                ]
                            ]
                        ,
                        "Interval" -> Quantity[regLen, "Batches"]
                        }
                ]
        ]
    ]

For this simple task, the two results are basically the same:

With[{y = 0},
    Plot[{Log@Abs[trained1[{x, y}] - Exp[-Norm[{x, y}]]], Log@Abs[trained2[{x, y}] - Exp[-Norm[{x, y}]]]}, {x, -3, 3}]
    ]

Answer 2

While not very fast or convenient, it is doable as described in example from https://reference.wolfram.com/language/tutorial/NeuralNetworksIntroduction.html#1033731519 :

gradient = trainingNet[<|"Input" -> ___ , "Target" -> 2|>, NetPortGradient[{"lenet", 1, "Biases"}]]

oldBias = Normal@NetExtract[trainingNet, {"lenet", 1, "Biases"}];

newBias = oldBias - 0.1*gradient;

trainingNet2 = NetReplacePart[trainingNet, {"lenet", 1, "Biases"} -> newBias]