FluxML / FluxML/Optimisers.jl

Expose an API for the accumulated/total gradients for shared parameters

Open
#211 0 comments 0 reactions 0 assignees View on GitHub

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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.