# Custom SGD optimizer in Mathematica neural network framework?

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?

• The Method option of NetTrain has a sub-option "LearningRateSchedule". My approach to your question is to define some "global" variables, so I can change them inside TrainingProgressFunction and use them inside "LearningRateSchedule". BTW nice question! – Silvia Aug 14 at 9:22
• Thanks, I just wanted to avoid tf. "LearningRateSchedule" only changes step size, TrainingProgressFunction is to monitor progress ... to build custom optimizer I need to directly get gradients and manually update weights - the only way I see is as in my answer below, but it seems extremely slow, this NetReplace seems like creating new network (~8 times per sample) - there is missing some NetUpdate for that, preferably allowing to directly work on vector of all Weights and Biases. A colleague is doing it in tf (nasty but you can modify used function), so I am leaving it for now. – Jarek Duda Aug 14 at 11:34
• I think this is an important question. I'm not sure my answer addressed it, if not please do let me know (though I'll be travelling next 3 days so may not be able to update promptly.) – Silvia Aug 14 at 19:15
• @Silvia, thank you, but you are using SGD with modified schedule, while I wanted to implement a completely different new custom optimizer (momentum method enhanced with online parabola model for smarter step choice: arxiv.org/pdf/1907.07063 ). To implement a custom optimizer a minimum is being able to ask about gradient and manually update weights and biases. And it can be done with NetPortGradient[] and NetReplacePart[] as below, but it is quite inefficient as the latter seems to make a new copy of entire network, while only update is needed here. – Jarek Duda Aug 14 at 19:41

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


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"}];