Questions about HuBERT/Wav2Vec2 pre-training

🐛 Describe the bug

[This is not really a bug report, more a request for clarification/discussion...]

I'm training a HuBERT BASE-sized model, first iteration targeting 100 clusters of MFCC, on a custom 1000-hr dataset (i.e. similar size to Librispeech). See the intriguing plots below:

Note that the triangles in the Tensorboard plots indicate there were NaN values. I'm perplexed as to why that leads the loss curve and unit prediction accuracies to have a "hiccup" around step 180K, but then seem to recover by step 200K. My hypothesis is that it's related to the special treatment of feature penalty and layer normalization in the BASE-sized models.

As noted in the Wav2Vec2 paper:

For the smaller Librispeech dataset, we regularize the model by applying an L2 penalty to the activations of the final layer of the feature encoder and scale down the gradients for the encoder by a factor of 10. We also use a slightly different encoder architecture where we do not use layer normalization, and instead of normalizing the raw waveform, the output of the first encoder layer is normalized.

Inspecting the HuBERT code, it's worth clarifying that the L2 "feature penalty" is in fact always included in the loss function, and it is scaled up by a hardcoded factor of 10x -- regardless of dataset or model size, and also irrespective of any masking -- although the 10x downscaling of the feature encoder gradients and the specialized layer normalization are only enabled in configurations for BASE-sized models. So I think a perhaps more straightforward interpretation is that the feature penalty is effectively unscaled for BASE models, and 10x upscaled for LARGE/XLARGE models. Am I reading that correctly?

So my first question: what changes might be suggested to avoid the loss "hiccup" that I've observed? Should I try adjusting the feature penalty scale and/or enabling standard layer normalization in the BASE model configuration?

My second concern is the drop in unmasked accuracy, which seems to start declining somewhat prior to the peak learning rate warmed up by step 20K. I suspect this is because the implementation of the HuBERT loss function does not give any weight to the unmasked logits. The HuBERT paper explored weightings of 0.0, 0.5, 1.0 and found that it was generally best to give zero weight to the unmasked loss component, especially when the targets are relatively low quality in terms of phonemic correlation. However, I wonder: might it be worthwhile to consider some small but non-zero weighting, say 0.1, for the unmasked loss, to prevent the under-fitting dip seen in these plots?

I also wonder about the length normalization when combining masked and unmasked losses. It seems that the current TorchAudio implementation will first combine these weighted masked and unmasked losses (which are summed from masked and unmasked logits of different lengths, depending on the masking parameterization), and add the feature penalty (averaged over all frames, irrespective of masking; this is later scaled by the length of masked logits) before later normalizing the overall summation of losses by the length of the masked logits. By contrast, the fairseq implementation would normalize by the sum of lengths of the masked plus unmasked logits (i.e. the full sequence length) if the weight of the unmasked loss is non-zero. Should the TorchAudio implementation be updated to match the fairseq implementation?

Moreover, I wonder: would it be a sensible improvement to instead normalize the masked and unmasked losses by their respective lengths prior to their weighted summation, and to also compute the feature penalty with respect to the weighting of the masked and unmasked losses (e.g unmasked feature frames should not contribute to the penalty if the unmasked weight is zero)? The advantage of this is that the weighting becomes decoupled from the effect of the masking parameterization and thus it's easier to tune this hyperparameter independently.

Versions

I'm using a slightly modified local fork of the main branch. The principal change is to refactor the training_step to (re-)enable automatic_optimization=True in Lightning (specifically for the gradient accumulation functionality, see #2918), rather than having a manual backward step.