Add cross-attention to output hypotheses

[mgaido91]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

The PR adds the encoder-decoder cross-attention to the output hypotheses returned by ASR models.

Collection: ASR

Changelog

• Returns the cross-attention scores in the output of the greedy generator
• Returns the cross-attention scores in the output of the beam search generator

Usage

• You can potentially add a usage example below

from nemo.collections.asr.models import ASRModel
from nemo.collections.asr.models.aed_multitask_models import MultiTaskTranscriptionConfig
from nemo.collections.asr.parts.submodules.multitask_decoding import MultiTaskDecodingConfig

model = ASRModel.from_pretrained(model_name="nvidia/canary-1b-v2")
decoding_config = MultiTaskDecodingConfig()
decoding_config.return_xattn_scores = True
model.change_decoding_strategy(decoding_config)
config = MultiTaskTranscriptionConfig(
    batch_size=4,
    return_hypotheses=True,
    num_workers=0,
    verbose=False,
    prompt={'source_lang': 'en', 'target_lang': 'en'},
    enable_chunking=False
)
output = model.transcribe("/Users/mgaido/Downloads/vp-test/aa.wav", override_config=config)
assert output[0].xatt_scores is not None

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

@nithinraok @andrusenkoau

Additional Information

• Related to Encoder-decoder attention extraction in ASR transcribe #14961 .

[nithinraok]

Thanks Marco. great work. Added comments. Also,
Could you add an option something like preserve_xattn_scores, so when enabled through

decoding_cfg = MultiTaskDecodingConfig(
    strategy="beam",  # or "greedy"
    preserve_xattn_scores=True,
)

only store and return xattn_scores (to save memory by default)

[nithinraok]

Shouldn;t shape is List[BxHxT1xT2] . Also best to add: this is used with AED models

[mgaido91]

this is for a single hypothesis, so there is no B... So this would be List[HxT1xT2]. If you prefer I can rename HxUxT to HxT1xT2. I will also add the indication about AED models.

[nithinraok]

what about condition when new_xattn_scores_list is None? cat would fail

[mgaido91]

if new_xattn_scores_list is None, xatt_scores_list will stay None, so we never enter in this if

[nithinraok]

I meant to ask for each step, but probably its fine.

[mgaido91]

I got it but if xatt_scores_list is not None, then new_xattn_scores_list must be not None, otherwise there is something weird happening (ie. attention scores are returned for some tokens and for some others are None in the same generation). I'd rather add an assert on new_xattn_scores_list, WDYT?

[nithinraok]

could you update here as well and also include in returns tuple

[nithinraok]

@mgaido91 could you update this one as well

[mgaido91]

@nithinraok thanks for your review. I do have a question, though. You said to add a param in MultiTaskDecodingConfig, however IIUC we should rather focus on TranscribeConfig, which is what is set through the transcribe API. Changing the MultiTaskDecodingConfig might be tricky for a user. Also, there are some config tricks required in the TrascribeConfig (i.e. enable_chunking=False) to enable getting the attention. So if we add a param in MultiTaskDecodingConfig, it might be quite complex for a user to get the attention, as it would require changing two different configs. WDYT?

[nithinraok]

@nithinraok thanks for your review. I do have a question, though. You said to add a param in MultiTaskDecodingConfig, however IIUC we should rather focus on TranscribeConfig, which is what is set through the transcribe API. Changing the MultiTaskDecodingConfig might be tricky for a user. Also, there are some config tricks required in the TrascribeConfig (i.e. enable_chunking=False) to enable getting the attention. So if we add a param in MultiTaskDecodingConfig, it might be quite complex for a user to get the attention, as it would require changing two different configs. WDYT?

Very good point.

Well, first, the option need to be passed to MultiTaskDecodingConfig for sure. To make it independent of the behavior of transcribe inference or not. But however, I get your point.
Two options:

• If the option needs to be through TranscribeConfig then we need to run model.change_decoding_strategy(updated_cfg) to change the default decoding behavior (which is return xattn_scores=None) inside .transcribe() each time. Which I don;t think is a good idea.
• 2nd options is we create decoding cfg. like:

multitask_decoding = MultiTaskDecodingConfig()
multitask_decoding.strategy = "greedy"
multitask_decoding.return_xattn_scores = True

and call

asr_model.change_decoding_strategy(multitask_decoding)

before performing .transcribe()

Lets go with later option for now. I will keep thinking about this as this needs to be changed across. Are there any other options like these you are interested on to be changed through .transcribe()?

[mgaido91]

Lets go with later option for now. I will keep thinking about this as this needs to be changed across. Are there any other options like these you are interested on to be changed through .transcribe()?

Ok, I will work on this in the next days. Maybe it would be worth adding some checks and logs to guide the user though. I will try to come up with a proposal for that while working on this. Thanks.

[andrusenkoau]

Hi @mgaido91, thank you for the great work! I have almost no questions. Let's wait your final changes with decoding config.

[andrusenkoau]

Could you add the check for both return_xattn_scores options (as above) here?

[mgaido91]

yes, sure, I did not do it to minimize the CI cost, I am updating it, thanks!

[mgaido91]

Thanks @andrusenkoau !

Let's wait your final changes with decoding config.

I already made them. You can find the return_xattn_scores in the decoding config. I also added a warning message to instruct the users. Is there something missing that you have in mind?

PS I do not see why [CICD NeMo / Nemo_CICD_Test (pull_request)](https://github.com/NVIDIA-NeMo/NeMo/actions/runs/20991982507/job/60354816538?pr=15229) is failing, I cannot see meaningful logs, maybe I just do not know where to look at as I am not that familiar with NeMo CI. It is also weird as it passed before I updated the branch with the master. May you point me to the logs?

[mgaido91]

@nithinraok I think I addressed your comments, may you please take another look at this?

[nithinraok]

Thanks Marco. Minor comment which you might have missed earlier. LGTM otherwise. Thanks for the PR, great work!

[mgaido91]

@nithinraok @andrusenkoau thank you for your guidance and reviews!