feat(aggregation): Add ExcessMTL

[KhusPatel4450]

feat(aggregation): Add ExcessMTLWeighting

Implements ExcessMTLWeighting from Robust Multi-Task Learning with Excess Risks (He et al., ICML 2024).

At each forward call, per-task excess risks are estimated via a second-order Taylor approximation (Equations 6-7) using an AdaGrad-style diagonal Hessian accumulated across all calls. Task weights are then updated via an exponentiated gradient step (Equation 9) and normalised to the probability simplex.

Design notes

• State: two registered buffers, _grad_sum ([m, n], accumulates squared gradients) and _weights ([m], current task weights). Both move with .to(device) and appear in state_dict().
• Warmup (n_warmup_steps, default 0): during warmup, weights stay uniform and gradient statistics are collected. On the first post-warmup call, the average excess risk over the warmup period is saved as a normalisation baseline (initial_w), following Appendix C.1. Setting n_warmup_steps=0 matches the official implementation and LibMTL behaviour (first call's excess used as baseline directly).
• Normalisation convention: weights initialised to [1/m, ..., 1/m] and always sum to 1, following the paper (vs. LibMTL's sum-to-m).
• _n_steps: stored as a registered buffer (scalar torch.long) so warmup progress survives checkpointing. Zeroed in-place in reset() to preserve device placement.

References

• Paper: https://proceedings.mlr.press/v235/he24n.html
• Official implementation: https://github.com/uiuctml/ExcessMTL
• LibMTL: https://github.com/median-research-group/LibMTL/blob/main/LibMTL/weighting/ExcessMTL.py

[ValerianRey]

Tyvm for the PR! This is only a partial review, but it's gonna be easier to review the rest after those things are adressed I think.

[ValerianRey]

Ty for the updates! This is much cleaner! I'll make the thorough round of review soon

[ValerianRey]

Very good job ty!

It's honestly one of the hardest implementations we've ever made!

Got a few things to fix still, and then we can merge.