How to use custom optimizer for analog update

[afmsaif]

Let's say I have a custom optimizer named Adabelief. Now I want to use this optimizer to update weights and biases of analog layers. How can I do it? One option I found that I can AnalogOptimizer like optimizer = AnalogOptimizer(Adabelief, model.parameters(), lr=0.5). In the documentation it is mentioned that "This class wraps an existing Optimizer, customizing the optimization step for triggering the analog update needed for analog tiles." Is it a correct approach to implement a custom optimizer in analog training? How this AnalogOptimizer class is actually working?

[maljoras]

Hi @afmsaif ,
thank you for the question. Using the AnalogOptimizer approach would only use this optimizer for the digital parts of the parameters, such as digital layers, layer norm affine transforms etc as well as in case of hardware-aware training for an analog inference deployment (where the backward pass and update is thought to be perfect but noise is injected in the forward pass only; see InferenceRPUConfig).

However, if you are considering "truly" analog in-memory training, that is having some of the weight updates done with specialized tile operations in memory (such as using stochastic pulse trains to implement in-memory gradient updates), it is not controlled by the AnalogOptimzer. Instead, this kind of in-memory update behavior is implemented in the analog_tile.update method, that is then only called from the AnalogOptimizer. If you want to customize this update behavior of individual tiles, you can derive a new class from an AnalogTile or CustomTile and override that method. This is exemplified in the CustomTile (see here) . An example how to do that is also (somewhat) described in the recent tutorial paper (see here Section VII.C).