+1
> θ = b2e/b2ep1
> Cᵢ = 2α / (b2ep1^2)
> s.θ_pow[1] = 1
> s.τ[1] += α*nlb2ep1
> s.∇τ[1,i] = 2par[i]*nlb2ep1
> for k in 1:s.m # order does matter
> # These 4 instuctions are supposed to be executed in that order,
> # although the two in the middle could be extracted in theire own loop
> # (so 3 loops at all, the two first could be @tturbo'ed. )
> s.θ_pow[k+1] = s.θ_pow[k]*θ
> s.∇τ[k+1,i] = 2par[i]*s.θ_pow[k+1]
> s.τ[k+1] += α * s.θ_pow[k+1]
> s.θ_pow[k] *= k # order of instructions matter, be carrefull.
> end
> b2ep1 = -b2ep1 # additions are easier than substractions for the loop.
> @tturbo for j in 1:s.d # order does not matter.
> idx = i+j*s.n
> Cᵢⱼ = Cᵢ * par[idx] * s.e[j]
> s.∇τ[1,idx] += b2ep1 * Cᵢⱼ
> for k in 1:s.m # order does not matter.
> s.∇τ[k+1,idx] = s.θ_pow[k] * Cᵢⱼ
> end
> end
> end
> end
>
> # Second part: construct the final gradient from it.
> @tturbo s.tmp .= s.J*(s.τ-s.τₑ)
> s.loss[1] = sum(s.tmp.^2)
> # @tturbo s.∇ .= vec(2(s.tmp's.J)*s.∇τ)
> turbo_aBC!(s.∇,s.tmp,s.J,s.∇τ)
> return nothing
> end
> #########################################################################################################
> ####### 4) Adam structure to handle stochatic gradient descent.
> #########################################################################################################
> struct Adam{T}
> theta::Vector{T} # Parameter array 2
> m::Vector{T} # First moment
> v::Vector{T} # Second moment
> b1::T # Exp. decay first moment
> b2::T # Exp. decay second moment
> a::T # Step size
> epsilon::T # Epsilon for stability
> t::MVector{1,Int} # Time step (iteration)
>
> end
> function Adam(theta::AbstractArray{T}) where T
> m = zeros(T,size(theta))
> v = zeros(T,size(theta))
> b1 = T(0.9)
> b2 = T(0.999)
> a = T(0.001)
> epsilon = eps(T)
> t = MVector{1}([0])
> Adam(theta, m, v, b1, b2, a, epsilon, t)
> end
> function step!(opt::Adam{T},grad) where T
> opt.t[1] += 1
> mfac = 1 / (1 - opt.b1^opt.t[1])
> vfac = 1 / (1 - opt.b2^opt.t[1])
> @turbo for i in eachindex(opt.theta)
> opt.m[i] = opt.b1 * opt.m[i] + (1-opt.b1) * grad[i]
> opt.v[i] = opt.b2 * opt.v[i] + (1-opt.b2) * grad[i]^2
> opt.theta[i] -= opt.a * (opt.m[i] * mfac) / sqrt(opt.v[i]*vfac + opt.epsilon)
> end
> return nothing
> end
> #########################################################################################################
> ####### 5) Problem structure and final fitting functions
> #########################################################################################################
> function build_p0(n,d)
> return randn(n*(d+1))
> end
> struct Problem{T,M,d,MM,dd}
> s::StaticProjectionStorage{T,M,d,MM,dd}
> data::HybridMatrix{d, StaticArrays.Dynamic(), T, 2, Matrix{T}}
> opt::Adam{T}
> end
> function Problem(par,data,m)
> T = promote_type(eltype(par),eltype(data))
> par = T.(par)
> data = T.(data)
> d,N = size(data)
> l = length(par)
> @assert l % (d+1) == 0
> n = Int(l//(d+1))
> data = HybridArray{Tuple{d,StaticArrays.Dynamic()}}(data)
> opt = Adam(T.(par))
> Storage = StaticProjectionStorage(T,N,m,d,n,cov(data'))
> return Problem{T,m+1,d,(m+1)^2,d^2}(Storage,data,opt)
> end
> function fit!(P::Problem{T,M,d,MM,dd},epochs) where {T,M,d,MM,dd}
> for i_e = 1:epochs
> Random.rand!(P.s.e) # Choose a random direction of projection.
> build_model!(P.s,P.data) # Construct the projected model (precomputations for the gradient)
> loss_and_∇!(P.opt.theta,P.s) # Evaluate the gradient and store it in P.s.∇
> step!(P.opt,P.s.∇) # Perform the gradient decent step.
> end
> return nothing
> end
> function extract_rez(P::Problem{T,M,d,MM,dd}) where {T,M,d,MM,dd}
> n = Int(length(P.opt.theta)//(d+1))
> α = P.opt.theta[1:n].^2
> θ = reshape(P.opt.theta[(n+1):end].^2,(n,d))
> return TD.MultivariateGammaConvolution(α,θ)
> end
> #########################################################################################################
> ####### 6) Direct application.
> #########################################################################################################
> # Construct a mockup dataset and launch the algorithm :
> d,N = 15, 10000;
> Random.seed!(13);
> data= reshape(exp.(randn(d*N)),(d,N)); # Ofc, this is mockup data.
> p0 = randn(100*(d+1));
> P = Problem(p0,data,20);
> fit!(P,10000);
> @info "Loss: $(P.s.loss[1])"
> rez = extract_rez(P);
> #########################################################################################################
> ####### 7) Simple Profiling.
> #########################################################################################################
> using BenchmarkTools
> @btime fit!(P,1) # 60μs and 2 allocations per step.
>
> using ProfileView
> ProfileView.@profview fit!(P,100_000) # launch it twice and discard the first run.
>
> ```
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [January 7, 2022, 2:27pm UTC](https://discourse.julialang.org/t/review-wanted-adam-gradient-descent-with-complicated-gradient-to-fit-multivariate-gamma-convolutions/74183/2 "2022-01-07T14:27:37Z")
> [@lrnv](#):
>
> Most of the code is directed into the computation of the gradient, which is of combinatorial nature (and takes up 95% of the runtime)
That sounds suspicious. With reverse-mode / adjoint differentiation (either automatic or [manual](https://math.mit.edu/~stevenj/18.336/adjoint.pdf)), you should always be able to compute the gradient in with a computational cost comparable to that of the loss function. (Essentially, make sure you are evaluating the chain rule from left-to-right and not from right-to-left.)
> [@lrnv](#):
>
> I have written a code that fits a multivariate gamma convolution to a dataset using a stochastic gradient descent on a home-made loss
Not sure why you are using stochastic gradient descent, as opposed to a deterministic minimization algorithm, unless you are sampling a huge dataset in random batches during fitting?
---
### Author:  [@lrnv](https://discourse.julialang.org/u/lrnv)
#### Post date: [January 7, 2022, 2:33pm UTC](https://discourse.julialang.org/t/review-wanted-adam-gradient-descent-with-complicated-gradient-to-fit-multivariate-gamma-convolutions/74183/3 "2022-01-07T14:33:11Z")
> [@stevengj](#):
>
> > [@lrnv](#):
> >
> > Most of the code is directed into the computation of the gradient, which is of combinatorial nature (and takes up 95% of the runtime)
>
> That sounds suspicious. With reverse-mode / adjoint differentiation (either automatic or [manual](https://math.mit.edu/~stevenj/18.336/adjoint.pdf)), you should always be able to compute the gradient in with a computational cost comparable to that of the loss function. (Essentially, make sure you are evaluating the chain rule from left-to-right and not from right-to-left.)
In fact I compute both of them at the same time, this time includes both. I think I’m indeed doing it in the right order, but i’ll read up the ref, thanks.
> [@stevengj](#):
>
> > [@lrnv](#):
> >
> > I have written a code that fits a multivariate gamma convolution to a dataset using a stochastic gradient descent on a home-made loss
>
> Not sure why you are using stochastic gradient descent, as opposed to a deterministic minimization algorithm, unless you are sampling a huge dataset in random batches during fitting?
Well, I have a stochastic componant in my loss (the loss is written as an expectation on the random variable `s.e`, that the function `fit!()` randomises at each iteration). This is indeed a stochastic gradient descent.