Add Expected Attention with Stats

[alessiodevoto]

PR description

So far Expected Attention only supports non QK norm models (no support for Gemma and Qwen) and it requires computing mean and covariance at inference time. This PR adds the option to load precomputed stats if they are available, or compute them on the fly if not. There are two main contributions:

• Expected attention with stats
• Change to the base press to have it work with subpresses that require post_init_from_model.

Checklist

• Tests are working (make test)
• Code is formatted correctly (make style, on errors try fix with make format)
• Copyright header is included
• All commits are signed-off using git commit -s
• (new press) mypress_press.py is in the presses directory
• (new press) MyPress is in __init__.py
• (new press) README.md is updated with a 1 liner about the new press in the Available presses section
• (new press) New press is in the default_presses list in tests/default_presses.py
• (new press) A docstring is provided that follows the same structure as the existing ones

[maxjeblick]

Thanks a lot for opening the PR!

I've left some comments; the most important one is probably the one about how to compute mean query and mean covariance.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxjeblick]

/ok to test

[copy-pr-bot]

/ok to test

@maxjeblick, there was an error processing your request: E1

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/

[maxjeblick]

/ok to test 24ef932

[alessiodevoto]

I have refactored the code significantly (sorry 😅), before computing and uploading stats to the hub and adding the press to tests and press list I would like to hear your opinion :)

Major changes:

• updated ea code and made it more modular (this way it is easier to subclass)
• instead of having a separate file just for computing stats, I put everything inside the press file. This way it stays self-contained and cleaner.

Let me know wdyt @maxjeblick 🙂

[maxjeblick]

Hi!
Refactoring (moving to another file) looks good.
The code itself needs some clean up, see my other comments above.
Id' also prefer having a main function with fire.Fire instead of argsparse.

[alessiodevoto]

Thanks Max!
No problem I will update later as well.

• For fire instead of argparse -> Ok.
• As for the comments, I think I addressed all of them in this PR. The only difference that I wanted to double check is whether we should have an external file for stats computation or it is better to keep it self contained like for duo attention.

[maxjeblick]

As for the comments, I think I addressed all of them in this PR. The only difference that I wanted to double check is whether we should have an external file for stats computation or it is better to keep it self contained like for duo attention.

I had one comment (it is collapsed above) about the if-else logic returning collected_queries, mean_query, cov_query. I think this can be simplified to always return those values.

The folder structure looks fine, it is nice that it is self-contained.

[alessiodevoto]

Got it, I simplified the if else logic. Also, there was a bug in the computation (I was computing the stats but forgot to save them), now it is fixed.

[alessiodevoto]

/ok to test e4157b2

[alessiodevoto]

Update:

• computed stats both for mock testing model and Qwen (for future)
• added tests (they download stats from the hub)

[alessiodevoto]

/ok to test 7f93961

[alessiodevoto]

Had a dependency error because of "fire", fixed. All tests pass now 👌

[alessiodevoto]

Forgot to clean some comments, now it's should be /ok to test bc27dd8

[alessiodevoto]

/ok to test bc27dd8

[maxjeblick]

LGTM, great work!