Feature: Heterogeneous Normalized Attention

[JakobEliasWagner]

Feature: Heterogeneous Normalized Attention

Description

This pull request introduces the implementation of the Heterogeneous Normalized Attention mechanism as described in the paper Hao et al., 2023.

The heterogeneous normalized attention block calculates the attention scores in these steps:

1. normalize the query and key sequence first

$$\tilde{q}_i = Softmax(q_i)$$

$$\tilde{k}_i = Softmax(k_i)$$

2. calculate the attention score without softmax

$$z_t = \sum_i \frac{\tilde{q}_t \tilde{k}_i}{\sum_j \tilde{q}_t \tilde{k}_j}v_i$$

This implementation is linear with respect to the sequence length.

We added a masking mechanism to the vanilla implementation suggested by Hao et al.

Which issue does this PR tackle?

• Heterogeneous normalized attention is not implemented.

How does it solve the problem?

• Implements HeterogeneousNormalizedAttention, a linear attention implementation.
• Implements masking for HeterogeneousNormalizedAttention.

How are the changes tested?

• Added 6 unit tests covering: initialization, shape projection, gradient flow, zero inputs, masked forwards, and correctness by masking a known tensor.

Checklist for Contributors

• Scope: This PR tackles exactly one problem.
• Conventions: The branch follows the feature/title-slug convention.
• Conventions: The PR title follows the Bugfix: Title convention.
• Coding style: The code passes all pre-commit hooks.
• Documentation: All changes are well-documented.
• Tests: New features are tested and all tests pass successfully.
• Changelog: Updated CHANGELOG.md for new features or breaking changes.
• Review: A suitable reviewer has been assigned.

Checklist for Reviewers:

• The PR solves the issue it claims to solve and only this one.
• Changes are tested sufficiently and all tests pass.
• Documentation is complete and well-written.
• Changelog has been updated, if necessary.

[samuelburbulla]

LGTM, you're the expert

[samuelburbulla]

Math does not render well in docs