[Feature Request] Add support for custom conflict-resolution hooks in gradient aggregation

[powerofaisinstudy-debug]

It would be great to have a clear hook/extension point in TorchJD that allows users to pass custom gradient transformation or conflict-resolution logic directly during the optimization step.

Currently, implementing experimental or non-standard gradient aggregation strategies requires either modifying core library functions or wrapping existing optimizers in a non-standard way.

When experimenting with multi-task learning models, researchers often need to test custom heuristics for resolving conflicting gradients (e.g., dynamic gradient scaling, specialized projection masks, or combining multiple conflict-resolution algorithms).

Having a clean, modular API hook for custom gradient transformations would make TorchJD much more flexible for prototyping new multi-task optimization research without needing to fork the repository.

cc @PierreQuinton @ValerianRey @ppraneth

Ideally, the API could accept a custom callable or transformer module:

import torch
import torchjd

def custom_gradient_transform(matrix: torch.Tensor) -> torch.Tensor:
    # Custom transformation / projection logic on the task gradient matrix
    return updated_matrix

# Pass custom transformation logic into the optimizer or step wrapper
optimizer = torchjd.optim.MinimalCostBinding(
    model.parameters(), 
    transform_fn=custom_gradient_transform
)

This would make it easier to benchmark custom multi-task algorithms against standard methods like PCGrad or CAGrad.

Happy to help draft a PR or discuss the design further if this fits into the project roadmap!

[PierreQuinton]

You could create a new Aggregator wrapper:

class CustomAggregatorWrapper(Aggregator):
  def __init__(self, aggregator: Aggregator) -> None:
    self.aggregator = aggregator

  def forward(self, matrix: Matrix, /) -> Tensor:
    return self.aggregator(custom_gradient_transform(matrix))

But in principle, we could make Aggregator composable:

class Aggregator:

  ...

  def compose(function: Callable[[Matrix], Matrix]) -> Aggregator:
    return Composition(self, function)

  __lshift__ = compose


class Composition(Aggregator):
  def __init__(self, aggregator: Aggregator, function: Callable[[Matrix]. Matrix]) -> None:
    self.aggregator = aggregator
    self.function = function

  def forward(self, matrix: Matrix, /) -> Tensor:
    intermediate_matrix = self.function(matrix)
    return self.aggregator(intermediate_matrix)

This would then be used with

my_aggregator << my_custom_transformation_on_matrices

[powerofaisinstudy-debug]

This composability model using << is super clean and elegant, @PierreQuinton! Making matrix transformation a first-class citizen on the Aggregator interface keeps everything modular and intuitive.

I'd be glad to open a PR for this! I can implement the Composition class, update Aggregator.__lshift__, and add unit tests covering custom matrix transformations to ensure high coverage.

Let me know if you'd like me to start working on it!

[ValerianRey]

Aggregators are nn.Modules internally so they already support all hooks that nn.Modules support. Is there anything you want to do that you cannot currently do?

[powerofaisinstudy-debug]

Hey ! While standard nn.Module hooks can intercept inputs, using them for clean matrix transformations requires writing boilerplate hook registration logic. The << composition syntax provides a declarative, readable pipeline (Aggregator() << Transform()) so users don't have to manually manage hook handles and state during multi-task experiments.

cc @PierreQuinton @ValerianRey