What is update_indexed and when should we use it?

[Zhaoxian-Wu]

I'm trying to implement new optimization algorithms with aihwkit, but I got a little confused on the original implementation of the native optimizer. I noticed in the Step method implementation in AnalogOptimizerMixin, there are two ways to update the weight:

if analog_ctx.use_indexed:
    for x_input, d_input in zip(
        analog_ctx.analog_input, analog_ctx.analog_grad_output
    ):
        analog_tile.update_indexed(x_input, d_input)
else:
    x_input = cat(
        analog_ctx.analog_input, axis=-1 if analog_tile.in_trans else 0
    )
    d_input = cat(
        analog_ctx.analog_grad_output, axis=-1 if analog_tile.out_trans else 0
    )
    analog_tile.update(x_input, d_input)

I have read the documents and there is no related material to explain the difference between these two methods.

I wonder

1. what the difference is between using update_indexed and update (source code of update_indexed)?
2. When should the use_indexed be set as True?
3. Moreover, what are the exact formats of x_input and d_input?
4. Why should we use the concatenate function cat() here?

I would appreciate it if somebody could solve my confusion.

[maljoras]

In general, this is internal code and the user should not worry about that unless you want to implement your own update functionality. Just leave it as it is, typically the use_indexed input argument is just for internal use and debugging.

That said, the functionality of the indexed update is internal C++ code that is optimized, as it avoids the "unfold" operation for the convolution internally by accessing the input tensor with indices. This uses less memory typically and is faster in certain cases (for 2D convs). However, for smaller DNNs and modern GPUs the difference is not large so either version can be used. However, for 1D and 3D convolutions are not supported by torch's unfold function. Here the indexed update needs to be used.

If you want to avoid overriding the indexed version for a custom tile for conv2d, you can set the use_indexed argument to the AnalogConv2d to false. But it should not modified dynamically on the fly after the layer is constructed, which might lead to some issues.

In general the code you are showing is from the optimizer. I would advise against editing or changing anything there since this makes the non-torch custom update functionality possible, you might just break everything if you do adapt it. Note that for in-memory training the update is changed to using stochastic pulse trains for instance. The cat is used to concatenate multiple forward passes that might have been called before the optimizer is called. Concatenating it to batches enhances the performance.

[Zhaoxian-Wu]

I see. The reason I am quite interested in this part is that I would like to customize my own optimizer on analog training and see whether it would be able to achieve better performance.

It seems that this part involves many C++ details. I wonder whether you have considered encapsulating them further into code at a lower abstraction level in the future version? Since I believe in the future, there will be more researchers paying attention to the optimization algorithm on analog devices when the analog training has a broader impact and this detailed code may get them confused :D