Back Propagation for SVD
- 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