Repository navigation
[Feature Request] Add support for custom conflict-resolution hooks in gradient aggregation #775
Description
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!
Activity
PierreQuinton commentedon Sep 12, 2026
You could create a new
Aggregatorwrapper:But in principle, we could make Aggregator composable:
This would then be used with
powerofaisinstudy-debug commentedon Sep 12, 2026
This composability model using
<<is super clean and elegant, @PierreQuinton! Making matrix transformation a first-class citizen on theAggregatorinterface keeps everything modular and intuitive.I'd be glad to open a PR for this! I can implement the
Compositionclass, updateAggregator.__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 commentedon Sep 14, 2026
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 commentedon Sep 14, 2026
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