Some problems with XVector implementation

[gau-nernst]

System Info

NA

Who can help?

@sanchi

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

While working on #25471, I spot a few problems with the current XVector implementation

1. Statistical pooling is using for loops for batched computation

transformers/src/transformers/models/wav2vec2/modeling_wav2vec2.py

Lines 2436 to 2438 in fe3c8ab

	for i, length in enumerate(tdnn_output_lengths):
	mean_features.append(hidden_states[i, :length].mean(dim=0))
	std_features.append(hidden_states[i, :length].std(dim=0))

This can be avoided by using the attention mask. However, since TDNN layers (basically 1D-conv) also change the effective attention mask, we need to update the attention mask again.

2. Std uses unbiased estimate (correction=1), while from what I know, statistical pooling typically uses biased estimate.

See that PyTorch uses correction=1 by default (unbiased=True in previous versions): https://pytorch.org/docs/stable/generated/torch.std.html

The original papers (x-vector: https://www.danielpovey.com/files/2018_icassp_xvectors.pdf, stats pooling: http://danielpovey.com/files/2017_interspeech_embeddings.pdf) don't specify the equations, but most other papers show equations without the correction factor: https://arxiv.org/pdf/1803.10963.pdf.

The official implementation is included in Kaldi, but I don't understand C++ and their codebase well enough to understand what's going on. https://github.com/kaldi-asr/kaldi/blob/71f38e62cad01c3078555bfe78d0f3a527422d75/src/nnet3/nnet-general-component.cc#L808-L821

If we use HF to train a model, it wouldn't matter much. But if we use this to port weights from other implementation, this can be a problem.

3. TDNN layer is 1D-conv. Using unfold and linear is wasteful. We can directly use 1D-conv instead.

NeMo's TDNN: https://github.com/NVIDIA/NeMo/blob/ab749e4401a3dcdfa5ea969347aaee20b7947c7c/nemo/collections/asr/parts/submodules/tdnn_attention.py#L138

Expected behavior

NA

[sgugger]

cc @sanchit-gandhi

[sanchit-gandhi]

Hey @gau-nernst! Thanks for the great write-up and for the code references!

Going through your points one-by-one:

1. Statistical pooling is using for loops for batched computation: would you like to open a PR to update this to use the attention mask as suggested? Think this should be fairly straightforward to add
2. Std uses unbiased estimate (correction=1): we could add the 'correction' as a config attribute should we wish to port a model from another library to HF. This way, we could use the weights in the HF implementation of Wav2Vec2. However, unless there is an external model that we want to port, I think we should leave it as is for the sake of simplicity
3. Using unfold and linear is wasteful: indeed, I agree! It's quite difficult to change the modelling code to add nn.Conv1D layers without breaking backwards compatibility for users who are running the old style unfold + linear, so here I think it's worth considering what the overhead is of using less efficient layers versus doing an implantation overhaul to use nn.Conv1D and maintain backwards compatibility.

Overall, the TDNN variant of Wav2Vec2 is not super used, but it's crucial that we keep the numerical / parameter compatibility with the existing model to prevent un-expected breaking changes. If you have the time to look into this more deeply I'd be happy to review any PRs, and as always help with any questions / queries!

[gau-nernst]

@sanchit-gandhi Thank you for the reply.

1. Yes, I can make a PR for this.
2. I agree with you, for the sake of backward compatibility and simplicity, we should leave it as it is.
3. In terms of speed, I don't think it will affect much since most of the computation is in the backbone anyway. For weights compatibility, sadly the shape of Conv1d weight is different from Linear weight in current implementation - (out_dim, in_dim, kernel_size) vs (out_dim, in_dim * kernel_size). One possible solution is to keep nn.Linear(), but call F.conv1d() together with reshaping kernel.weight, so that we can keep backward compatibility and enjoy speed up. I will look more into this.

[sanchit-gandhi]

Your solution to 3 sounds very interesting! Super keen to take a look at a PR if this is something you want to pursue!

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.