Does the Binarize() function use STE？

[leejiajun]

Does the Binarize() function use STE？
I haven't seen the STE algorithm in this whole project.

[mohdumar644]

I want to know the same.

Probably the Binarize operation on the layer weights is not recorded (the weights are leaf nodes?). During backward pass, the gradient is calculated w.r.t. the binary weights. This same gradient is used to update the real weights, hence implicitly STE.

[haichaoyu]

Want to know the same +1

[Ashokvardhan]

@mohdumar644, @haichaoyu, @leejiajun : I went through the code of binarized_modules.py and in the class BinarizeLinear, we can see that the line self.weight.data=Binarize(self.weight.org) modifies the values of the weights to be binary. This means that both during the forward and backward passes, the binary weights are used. However for the gradient update we need real weights which are only available in self.weight.org. However I am not sure how these real weights are used by the optimizer.

[mohdumar644]

weights are binarized in forward pass using self.weight.data=Binarize(self.weight.org).

prior to optimizer step, real weights are copied back.

[mohdumar644]

btw, a more understandable approach to constructing binarized/quantized neural networks in PyTorch can be found here with a custom Quantizer/Binarizer module having explicit STE in the backward pass.

[Ashokvardhan]

@mohdumar644 thanks for the links! In the attached implementation too, the class BinarizeLinear uses the function QuantizeWeights which is very similar to this PyTorch implementation. They also use the copying trick. So I don't see a clear difference in the implementation except for the Quantizer module which has own forward and backward functions.

[Ashokvardhan]

weights are binarized in forward pass using self.weight.data=Binarize(self.weight.org).

prior to optimizer step, real weights are copied back.

I see. I was only looking at the file main_mnist.py which didn't have this copying trick and hence got confused. So it seems the current implementation for MNIST is wrong because they are updating the binary weights instead of the real weights. what do you think?

[haichaoyu]

Hi @Ashokvardhan , thanks for your insights. I just focused on the the resnet code for cifar10 (resnet_binary.py). In fact, I have two questions about the code.

• One is the hardtanh here. This function is called between bn and next conv layer. hardtanh makes output between -1 and +1. This makes backward of binarization inside each conv and linear does not take effects (gradients are always 1). My solution is to just remove the tanh layer and use an explicit STE layer as you mentioned above (The difference is that I used gate function in backward). After these modifications, performance remains the same.

• The second possible problem is that the Binarization operation (inside each conv and linear) are directly applied to Tensor.data instead of Tensor. As a result, binarization is not recorded in gradient graph. An explicit STE on Tensor instead of Tensor.data may solve this problem.

Please correct me if there is any problem.

[mohdumar644]

@mohdumar644 thanks for the links! In the attached implementation too, the class BinarizeLinear uses the function QuantizeWeights which is very similar to this PyTorch implementation. They also use the copying trick. So I don't see a clear difference in the implementation except for the Quantizer module which has own forward and backward functions.

My bad actually since I am a bit rusty - the links I shared only change the way the intermediate activations are binarized/quantized, not the weights - for the weights the way is the same i.e. an extended Linear module. The code I shared just avoided having to input.data = Binarize(input.data) inside this extended class.

weights are binarized in forward pass using self.weight.data=Binarize(self.weight.org).
prior to optimizer step, real weights are copied back.

I see. I was only looking at the file main_mnist.py which didn't have this copying trick and hence got confused. So it seems the current implementation for MNIST is wrong because they are updating the binary weights instead of the real weights. what do you think?

Did you look at this line?

[mohdumar644]

• One is the hardtanh here. This function is called between bn and next conv layer. hardtanh makes output between -1 and +1. This makes backward of binarization inside each conv and linear does not take effects (gradients are always 1). My solution is to just remove the tanh layer and use an explicit STE layer as you mentioned above (The difference is that I used gate function in backward). After these modifications, performance remains the same.

You basically integrated Hardtanh into Quantizer/Binarizer function by using a clamper/gate in the backward. Good!

[Ashokvardhan]

@mohdumar644 thanks for the links! In the attached implementation too, the class BinarizeLinear uses the function QuantizeWeights which is very similar to this PyTorch implementation. They also use the copying trick. So I don't see a clear difference in the implementation except for the Quantizer module which has own forward and backward functions.

My bad actually since I am a bit rusty - the links I shared only change the way the intermediate activations are binarized/quantized, not the weights - for the weights the way is the same i.e. an extended Linear module. The code I shared just avoided having to input.data = Binarize(input.data) inside this extended class.

weights are binarized in forward pass using self.weight.data=Binarize(self.weight.org).
prior to optimizer step, real weights are copied back.

I see. I was only looking at the file main_mnist.py which didn't have this copying trick and hence got confused. So it seems the current implementation for MNIST is wrong because they are updating the binary weights instead of the real weights. what do you think?

Did you look at this line?

Aah, didn't notice this. Thanks! Btw, this file main_binary.py assumes that one has access to the datasets for cifar-10, cifar-100, imagenet. Do you know in what format they are needed to run this code? In case you already happened to run these codes, can you let me know about this? Thanks!

[mohdumar644]

As per this,

• the cifar-x datasets (and even the mnist dataset when needed) will be automatically downloaded and extracted in a certain folder the first time the script is run, and later reused.
• the imagenet is much larger (>100 gb) and will be needed to be downloaded and setup in the relevant directory yourself

[Ashokvardhan]

Hi @Ashokvardhan , thanks for your insights. I just focused on the the resnet code for cifar10 (resnet_binary.py). In fact, I have two questions about the code.

• One is the hardtanh here. This function is called between bn and next conv layer. hardtanh makes output between -1 and +1. This makes backward of binarization inside each conv and linear does not take effects (gradients are always 1). My solution is to just remove the tanh layer and use an explicit STE layer as you mentioned above (The difference is that I used gate function in backward). After these modifications, performance remains the same.

Sorry, I didn't understand this. can you please explain in detail?

• The second possible problem is that the Binarization operation (inside each conv and linear) are directly applied to Tensor.data instead of Tensor. As a result, binarization is not recorded in gradient graph. An explicit STE on Tensor instead of Tensor.data may solve this problem.

Please correct me if there is any problem.

I realized that there is no problem with their implementation. The reason they didn't have to use STE() is because they are modifying only the Tensor.data which won't be recorded for gradient computation. This is exactly what we need because we are using the gradients on the binary weights to update the real weights directly.

[haichaoyu]

Regarding the first point:

• According to the paper (Equation 4), the gradient of STE is 0 when abs(weight) > 0 and 1 when abs(weight) < 1. If activations are passed through torch.HardTanh before STE, abs of all activations are less than 1. Thus gradients of STE are always 1?

Actually, both of my concerns are summarized in this issue title by leejiajun. Does the network use STE with gate gradient? (Equation 4 in the paper)

[itayhubara]

By using the HradTanhH you ensure that the gradients of STE is 0 when abs(activation) > 1 and 1 when abs(activation) < 1. Note that the STE is used over the activations (not the weights) to avoid running back-propagation through the sign function.