FluxML / FluxML/Tracker.jl

Back Propagation for SVD

Open
#97 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
54
Forks
37
Avg merge
5h 49m
Merged PRs (30d)
1

Description

Here I present the correct (but poor) implementation of BP for SVD, this implementation changes the original `svd` interfaces a bit, hoping someone can help improve it.

```julia
using LinearAlgebra
using Flux
using Flux.Tracker: @grad, data, track, TrackedTuple
import Flux.Tracker: _forward
import LinearAlgebra: svd

"""stablized back propagation function for svd"""
function svd_back(U, S, V, dU, dS, dV)
NS = length(S)
S2 = S.^2
Sinv = 1 ./ S
F = S2' .- S2
@. F = F/(F^2 + 1e-12)

UdU = U'*dU
VdV = V'*dV

Su = (F.*(UdU-UdU'))*Diagonal(S)
Sv = Diagonal(S) * (F.*(VdV-VdV'))

U * (Su + Sv + Diagonal(dS)) * V' +
(I - U*U') * dU*Diagonal(Sinv) * V' +
U*Diagonal(Sinv) * dV' * (I - V*V')
end

svd(a::TrackedArray) = track(svd, a)
# I doubt the macro `@grad` interface is less intuitive than `_forward`
function _forward(::typeof(svd), a)
U, S, V = svd(data(a)) # making `svd` return value SVD, making Julian's life shorter.
# returning a list won't work, one will get 0 gradient
# [U|>param, S|>param, V|>param], -> (svd_back(U, S, V, dU, dS, dV),)
(U, S, Matrix(V)), Δ -> (svd_back(U, S, V, Δ...),)
end

# This is a use case
M, N = 4, 6
K = min(M, N)
A = param(randn(M, N))
res = svd(A)
# implement `Base.iterate(res::TrackedTuple) = ?` can make it prettier
U, S, V = res[1], res[2], res[3]

dU, dS, dV = randn(M, K), randn(K), randn(N, K)
Tracker.back!(res, (dU, dS, dV))
Tracker.grad(A)
```
## Why we use `Matrix(V)` here?
We see this line in file `src/tracker/scalar.jl` is called
```
track(f::Call, xs::Tuple) = TrackedTuple(xs, Tracked{typeof(xs)}(f, zero.(xs)))
```
One should notice function `zero` can change `type` sometimes!
Here, SVD returns V as Adjoint, `zero(Adjoint)` will get `Array`!

# Gocha!

## Some aspects can be improved
* Returning [U, S, V] in Flux should not cause gradient tracking failure.
* Tracker should be able to propagate over dagger?
* Return value checking for `_forward` is nessesary, so that readable error message can be throwed.
* `@grad` is an arguably useful interface
* Julia should remove over designed outputs for linear algebra functions like `svd`, I didn't see many benefits of such design.
* `zero` and `one` should never change type, here, it should be considered as a bug.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with the SVD example in the issue and read src/tracker/scalar.jl, especially track and TrackedTuple, along with the shown _forward and @grad paths. The issue lists several possible goals rather than one defined change, so first establish which behavior is in scope; done should include a reproducible gradient check for that selected goal and a clear test or error expectation.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
22/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.