kontrabas:

But now (while I was writing this) I found another problem - while I’m calculating loss from this “turned on” output layers from model - in the step loss.backward() and optimizer.step() it backpropagates through all of the output layers, not only this “turned on” for which data has passed, isn’t it? Is there possibility to backpropagate and optimize only through this output layers of model (“heads”) which are turned on while training for each sample?


The computation graph will only include the layers (and their parameters), which were used during the forward pass, so that only these parameters will get valid gradients if you call backward().
The parameters of the unused layers should get any gradients. Could you check it please?



 kontrabas:

if I train my model using samples and without normalization, I can inference with batches (without batch normalization), can’t I?


I don’t completely understand, how your decide which “heads” are used for the current input.
If you use some specific statistics for the current activation, you could probably apply these check for the complete batch and send each sample to the corresponding head.