feat(aggregation): Add GradVac aggregator

[rkhosrowshahi]

Summary

Adds Gradient Vaccine (GradVac) from ICLR 2021 as a stateful Aggregator on the full task Jacobian.

Behavior

• Per-block cosine statistics and EMA targets \bar{\rho}, with the closed-form vaccine update when \rho < \bar{\rho}.
• group_type: 0 whole model (single block); 1 all_layer via encoder (leaf modules with parameters); 2 all_matrix via shared_params (one block per tensor, iteration order = Jacobian column order).
• DEFAULT_GRADVAC_EPS and configurable eps (constructor + mutable attribute).
• Autogram not supported (needs full rows and per-block inner products). Task shuffle uses torch.randperm; use torch.manual_seed for reproducibility.

Files

• src/torchjd/aggregation/_gradvac.py, export in __init__.py
• docs/source/docs/aggregation/gradvac.rst + index toctree
• tests/unit/aggregation/test_gradvac.py

Verification

• ruff format / ruff check on touched paths
• ty check on _gradvac.py
• pytest tests/unit/aggregation/test_gradvac.py tests/unit/aggregation/test_values.py -W error
• Sphinx HTML build for docs

[rkhosrowshahi]

Hey @ValerianRey :
Could you please add cc: feat and package: aggregation labels to this PR?

[ValerianRey]

I think the link check action failed because of bad luck, re-running the job fixed it.

[ValerianRey]

Thanks a lot for the PR, I'll look at it tonight or tomorrow!

[ValerianRey]

/opencode:Plan can you please check that the actual implementation (from line 219 to line 251 of _gradvac.py) is equivalent to the implementation in https://github.com/median-research-group/LibMTL/blob/main/LibMTL/weighting/GradVac.py? Please report anything that differs.

[opencode-agent]

Let me analyze the differences between the two implementations:

Comparison Summary

After comparing lines 219-251 of _gradvac.py with the LibMTL reference implementation, here are the findings:

✅ Equivalent Behaviors

1. Core Algorithm: Both implement the same GradVac algorithm:

   • Clone the gradients to pc_grads
   • Iterate over task pairs (i, j) where i ≠ j
   • Random shuffle of j for each i
   • Compute cosine similarity rho_ijk per block
   • Apply vaccine update when rho_ijk < bar (EMA target)
   • Update EMA: bar = (1-beta)*bar + beta*rho_ijk

2. Weight Formula: Mathematically identical:

   • LibMTL: ||g_i||*(ρ̄*√(1-ρ²)-ρ*√(1-ρ̄²))/(||g_j||*√(1-ρ̄²)+ε)
   • TorchJD: norm_i*(bar*sqrt_1_rho2-rho_ijk*sqrt_1_bar2)/(norm_j*sqrt_1_bar2+eps)

3. Group Types: Both support the same three grouping strategies (whole_model, all_layer, all_matrix)

🔍 Differences Found

Aspect	LibMTL	TorchJD (this PR)
RNG for shuffling	random.shuffle() (Python stdlib)	torch.randperm() (PyTorch)
Epsilon	Hardcoded 1e-8	Configurable eps attribute
Numerical safety	No clamp before sqrt()	.clamp(min=0.0) before sqrt()
Block indexing	Runtime sum(self.k_idx[:k])	Precomputed offsets list
Whole model case	Special case with k_idx=[-1]	Natural via sizes=(n,)
Return value	batch_weight (unused)	pc_grads.sum(dim=0)
State management	self.step counter	_state_key tuple

📝 Notes

1. The clamp operation in TorchJD ((1.0 - rho_ijk * rho_ijk).clamp(min=0.0).sqrt()) is a numerical safety improvement not present in LibMTL - it prevents potential NaN from negative values due to floating-point errors.

2. The random number generator difference is documented in the TorchJD docstring: "Task-order shuffling uses the global PyTorch RNG".

3. Both implementations correctly handle the asymmetric nature of the algorithm (updating pc_grads[i] based on comparisons with grads[j]).

Conclusion: The implementations are algorithmically equivalent. TorchJD includes minor improvements (configurable eps, numerical clamp) and follows the codebase's design patterns.

opencode session  |  github run

[ValerianRey]

/opencode:Plan Please review this PR in detail.

[ValerianRey]

Really good work, tyvm! This seems almost ready to merge IMO. Just a bunch of nitpicks that should be really easy to fix, and 1 somewhat important implementation question (iterating over groups in outer loop vs inner loop).

I would also argue that it seems that this could be implemented as a Weighting internally (because we actually act on norms and cosine similarities between gradients, which is what the gramian contains). Also, it's possible to keep track of norms and cosine similarities between projected gradients even if we don't have those gradients, just by making some operations on the gramian. This is what we did to implement PCGrad as a Weighting.

For example, imagine you have g1 and g2 be two gradients. From the gramian, you know ||g1||, ||g2|| (the sqrt of the diag elements), and g1 . g2 (an off-diag element), so you can deduce cos(g1, g2) from that.

If you compute g1' = g1 + w * g2, you can also directly deduce the norm of g1':
||g1'||² = ||g1||² + w² ||g2||² + 2w g1 . g2 (all elements of the right handside are known).

Similarly, you can compute g1' . g2 = (g1 + w * g2) . g2 = g1 . g2 + w g1 . g2.

So even after projection, you still know the dot products between all of your gradients, meaning that you still know the "new" gramian.

I didn't think through it entirely but at a first glance it seems possible to adapt this as a weighting, because of that. The implementation may even be faster actually (because we have fewer norms to recompute). But it may be hard to implement, so IMO we should merge this without even trying to implement it as a Weighting, and we can always improve later. @PierreQuinton what do you think about that?

[rkhosrowshahi]

Opencode's review was quite low quality, but it mentioned something that I missed: we need a test for GradVac in tests/unit/aggregation/test_values.py.

Similarly, i'd like to have GradVac added to tests/plots/interactive_plotter.py.

Thanks. I added the GradVac to the code and improved the code a bit to be more user-friendly. See the PCGrad and GradVac in the plot, find the same aggregated gradient. If you liked the changes, I can add to the commit as well.

[ValerianRey]

I think this is ready to merge, except for some plotting things. Can we remove the changes to the plotter and make plotter improvements in a different PR (except adding GradVac to the list of aggregators in the plotter)? I see a few issues in the plotter changes, and I'd rather merge this PR now and make the rest of the changes in a different PR. @rkhosrowshahi

BTW the link check action will fail because the links I added in the readme point to some documentation that will only be built after we merge this.