Extent of purpose for quantization-aware-training for weights in aihwkit?

[arseniivanov]

Hello!

I am from a Edge AI/TinyML background where quantization-aware-training(QAT) is often used to gain network robustness when running integer-based nets.

I have found that aihwkit supports ADC and DAC conversions through the rpu_config such as:

rpu_config.forward.inp_res = 1/64.  # 6-bit DAC discretization.
rpu_config.forward.out_res = 1/256. # 8-bit ADC discretization.

This could been seen as a "quantization", however the quantization here affects the in/out signals themselves, and not the weights, which limits the effectiveness of QAT using this functionality in my opinion as we do not get the weight robustness.

In "Hardware-aware training for large-scale and diverse
deep learning inference workloads
using in-memory computing-based accelerators", the authors write:

The successes of QAT cannot be directly trans-
lated onto AIMC, since the MVM approximations arise from fundamentally different con-
cepts. In AIMC, weights are represented by conductances that are physical properties of
NVM devices.

However, when looking at paper such as "Fully hardware-implemented memristor
convolutional neural network", they describe their memristor programming process such as:

During the corresponding experiment, the conductance was
programmed from 2.5 μS (0.5 μA at a 0.2-V read pulse) up to 20 μS (4 μA
at a 0.2-V read pulse) with a constant step of 2.5 μS. The equivalent
15-level weight of the memristor pair was thus referred to the 15 indi-
vidual differential conductance values that were uniformly distributed
from negative 17.5 μS (2.5 μS–20 μS) to positive 17.5 μS (20 μS–2.5 μS).
Moreover, the effect of read disturbance on the 15-level conductance
weights after applying 106 read pulses (0.2 V) is investigated ...

Here it seems like discrete steps were used in programming the crossbar array, so in my mind, there would be a benefit to exploring QAT for training of this kind of AIMC network.

My questions are:

Is there any other information/research done in this area?

Is there a purpose for QAT in this framework?

Bonus/Curious question:

If implementing QAT for aihwkit, is it reasonable to extend AnalogLayer/AnalogLayerBase to use the set_weights/get_weights such as:

    def forward(self, x_input: Tensor) -> Tensor:
        weight, bias = self.get_weights()
        self.set_weights(self.quantize(weight), self.quantize(bias) if bias is not None else None)
        return super().forward(x_input)

Or are there inherent noises/imperfections added in the AnalogLayer/AnalogLayerBase similar to how it's done for AnalogSGD?

[maljoras]

Hi @arseniivanov,
many thanks for raising this interesting point. In the first paper, hardware-aware training is done with an analog representation of the weights in mind, that is the weight is directly encoded into the conductance values, and only having one pair of resistive elements (one for positive and one for negative parts of the weight). Thus, it makes no sense to use QAT methods since there is no quantization of the weights in this case. There are only noise and limited weight ranges, since the resistive value writing and read out is subject to noise and non-idealities as described in the first paper. Therefore this is very different to the QAT situation in digital, where weights are quantized and the computation is fully deterministic. In the second paper the assumptions are indeed different, since a resitistive element is taken for each bit of a binary word (bit slicing of the weight). So it is indeed quantized in some way. However, note that here also the weight written is noisy and thus the non-determinsitic part still remains in discrepancy of the situation of the typical QAT in digital. If you do not assume any noise and just use digital QAT, it will likely not be the best model. Note that e.g. the digital quantized model could learn to distinguish perfectly between two close numbers e.g. 244 and 243, which however would be not possible in the noisy case, where a larger signal-to-noise ratio for the decision boundary is needed. Often, for that reason, a very heavily digitally quantized model has poor accuracy results on analog.

That being said a digitally quantized model might be useful for the quantized AIMC approach in the second paper as a starting case for hardware-aware training, where it could be then made noise robust.

In terms of training quantized weight with the AIHWKIT, one could do it in different ways. One could be to build the same forward pass as expected for the particular hardware (e.g. the second paper ) something like:

    def forward(self, x_input: Tensor) -> Tensor:
        # [.. ] init y_q tensor with correct size 
        for  significance_factor, analog_tile in zip(self.sig_factors, self.analog_tile_array):
            y_q.add_(significance_factor * analog_tile.forward(x_input))
        return y_q

Where the analog tiles are instantiated for each bit of the weight in the init. Then one might be able to use noise-aware training with this model.

However, this might be somewhat complicated and would might want to directly train with QAT and noise. That can actually be done by using DISRETIZE_AND_NORMAL or DOREFA. However, the support for this kind of quantization aware training is limited and other packages on QAT might yield better results. You can take a look at the tutorial we recently wrote, to see how one can use hardware-aware training with the AIHWKIT, see here.

Thus, a third way would be to use a state-of-the-art QAT network obtained from some other specialized package on QAT and then convert the resulting model to analog using the convert_to_analog tool. Then one could apply realistic AIMC inference noise to this weight. However, note that by default the "analog" approach to the weight representation is used, that is, only using one pair of conductances. If you would use the bit-splitting, then you would need to do the splitting as indicated above and then load the weights as wished. Inference noise can be added to that as well.

[arseniivanov]

Ah, I see. Thank you for taking time and a thorough reply. I see that there are so many approaches to weight representation on AIHW.

My goal is to achieve exactly this:

That being said a digitally quantized model might be useful for the quantized AIMC approach in the second paper as a starting case for hardware-aware training, where it could be then made noise robust.

I guess its similar to the Precision of bit slicing with in-memory computing based
on analog phase-change memory crossbars paper, I will read that one more thoroughly.

Thank you for providing alternative ways to implement this. I have only read the readthedocs for the concepts so far, so I will read the tutorial to understand the limits and possibilities of the framework. For me it sounds like the first approach would be the best, and if done correctly may be useful to future research. I will assess the complexity and feasibility for me to do this. Good day to you!