Make RelPosEncXL able to encode the sign of the relative position - #3073

What does this PR do?

RelPosEncXL.make_pe() fills the negative half of the relative positional embedding with sin(+p) instead of sin(-p). cos is even so the cosine terms were already right, but sin is odd, so the rows for relative distance -d come out bit-identical to the rows for +d.

Both halves are read. rel_shift maps row p of pos_embs onto the (query, key) pairs at relative distance d = (klen - 1) - p, so rows past the centre are the keys sitting to the right of the query. Probing it with a one-hot matrix, klen = 5:

row 3 (d=+1) -> i=1:j=0, i=2:j=1, i=3:j=2, i=4:j=3
row 4 (d= 0) -> i=0:j=0, i=1:j=1, i=2:j=2, i=3:j=3, i=4:j=4
row 5 (d=-1) -> i=0:j=1, i=1:j=2, i=2:j=3, i=3:j=4

Rows 3 and 5 are equal on develop, so a non-causal RelPosMHAXL adds the same positional bias to a key one step ahead as to a key one step behind. Running a constant sequence through it, so that the content term is flat and only the positional term shapes the map, gives an attention row that is exactly symmetric about the query:

[0.179007, 0.191125, 0.158071, 0.1226, 0.158071, 0.191125]

ESPnet's RelPositionalEncoding negates the position for both terms, and Transformer-XL's PositionalEmbedding takes a signed pos_seq, so both are odd in the sine terms.

Refs #3070, reported and diagnosed by @coollip.

Why the default is unchanged

Every checkpoint SpeechBrain has released was trained against the symmetric embedding, so switching it silently would degrade all of them. The corrected formulation is behind use_legacy_symmetric=False and the default output is bit-identical to develop:

>>> torch.equal(develop_pe, patched_pe)   # emb_dim=64, seq_len=50, default args
True

To put a number on what flipping the default would cost, speechbrain/asr-streaming-conformer-librispeech decoded in full-context greedy mode over the first 500 utterances of LibriSpeech test-clean, on CPU:

use_legacy_symmetric=True    WER 2.12   224 edits / 10561 words
use_legacy_symmetric=False   WER 2.90   306 edits / 10561 words

That is a compatibility cost and not evidence that the symmetric encoding is better, since the checkpoint was trained with it. I have no way to train a Conformer from scratch to compare the two properly, so this PR makes no claim that the corrected encoding trains better, only that it is the one Transformer-XL and ESPnet define.

Models on the hub built with attention_type: RelPosMHAXL and causal: False, which is what would need retraining or pinning if the default were flipped:

speechbrain/asr-branchformer-large-tedlium2
speechbrain/asr-conformer-loquacious
speechbrain/asr-conformer-transformerlm-ksponspeech
speechbrain/asr-conformer-transformerlm-librispeech
speechbrain/asr-conformersmall-transformerlm-librispeech
speechbrain/asr-streaming-conformer-gigaspeech
speechbrain/asr-streaming-conformer-librispeech

plus 27 recipe hparams using RelPosMHAXL, none of which set causal: True. Causal encoders are unaffected either way, the negative rows are masked out and the module output is bit-identical.

Verification

The new test compares make_pe() against a scalar reference written from the definition, in the same shape as the existing test_rope_rotate. On develop:

$ pytest tests/unittests/test_attention.py -q
E   TypeError: RelPosEncXL.__init__() got an unexpected keyword argument 'use_legacy_symmetric'
5 failed, 17 passed in 1.55s

With the patch, and with the surrounding modules:

$ pytest tests/unittests/test_attention.py tests/unittests/test_conformer.py \
    tests/unittests/test_streaming.py tests/unittests/test_transformer_src_tgt_masks.py -q
28 passed in 1.19s

$ pytest tests/unittests -q
4 failed, 743 passed, 6 skipped in 208.93s

$ pytest --doctest-modules speechbrain/nnet/attention.py speechbrain/lobes/models/transformer/ -q
31 passed in 3.29s

The 4 unittest failures are test_RNN, test_SpectrogramDrop, test_fallback_when_all_fail and test_run_shell. They fail identically on an unpatched checkout here (torch 2.13, python 3.12, macOS), so they are unrelated.

Open question

The flag is only reachable by constructing RelPosEncXL directly. Making it settable from a recipe needs a pass-through in TransformerInterface and TransformerASR, and flipping the default instead would need the affected hyperparams pinned the way #2604 was handled. I kept this PR to the one module so the choice stays with you, and I am happy to follow up with either.