Expose an API for the accumulated/total gradients for shared parameters
Nobody has claimed this yet.
- Dominant language
- Julia
- Stars
- 96
- Forks
- 30
- PR merge metrics
- No merged PRs in 30d
Description
### Motivation and description
It would be useful to have a public API for the two steps within `update!` separately:
https://github.com/FluxML/Optimisers.jl/blob/4ff61fca27c31f6d6bbc6bac19019b1de3634fd7/src/interface.jl#L70-L79
In the total gradient is typically more useful than the individual gradient contributions. The use case where this came up was for tracking parameter gradient norms: near convergence the shared gradients may sum to zero, but generally each contribution will be non-zero, so summing the component norms will give you the wrong idea about the gradient size.
### Possible Implementation
A trivial implementation would just refactor `update!` into two functions:
```julia
function update!(tree, model, grad, higher...)
# First walk is to accumulate the gradient. This recursion visits every copy of
# shared leaves, but stops when branches are absent from the gradient:
grads = total_gradients(tree, model, grad, higher...)
# Second walk is to update the model. The params cache indexed by (tree,x),
# so that identified Leafs can tie isbits parameters, but setup won't do that for you:
return update!(tree, model, grads)
end
function total_gradients(tree, x, x̄s...)
grads = IdDict{Leaf, Any}()
_grads!(grads, tree, x, x̄s...)
return grads
end
function update!(tree, model, grads::IdDict)
newmodel = _update!(tree, model; grads, params = IdDict())
tree, newmodel # note that tree is guaranteed to be updated. Also that it's not necc a tree.
end
```
But anything along these lines would be great.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/interface.jl around lines 70–79 and trace how update! performs its gradient-accumulation and model-update walks. Review the proposed total_gradients, _grads!, and _update! flow, then define the public separation while preserving shared-parameter behavior. Done means callers can obtain accumulated gradients independently and still update the model.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100