Caching system

[Applics]

Hello,

I wanted to know if, as for the Manopt version in matlab, it was possible to save elements calculated in the objective function, gradient function or the hessian function. In order to use prevent us from redudant computations.

As for the example given on the Matlab version of manopt:

I find myself in the situation where I'm performing certain calculations in the cost function that also appear in the gradient and I'd like to avoid having to recalculate these same terms.

Is it possible to do this in Julia's version? I'm having trouble seeing the links between what I want to do and what I can find here https://manoptjl.org/stable/tutorials/CountAndCache/

Best,

Ben

[kellertuer]

Hi,
I am not 100% sure what you are referring to in your screenshot – do you mean to cache the computation of Ax?

The Caching system in Manopt directly caches arbitary calls in the objective, that is correct, but since we want to keep cost function signatures simple, they are of the form f(M,p) without further parameters. So compared to Matlab, we do pass the manifold into the cost but no further parameters like a storage. The challenge with that would be, that nearly all lines of code would be affected by that.

There is a major design difference here between Manopt in Matlab and Manopt.jl.
Manopt in Matlab works quite “monolithic” – you have this global – everywhere-passed storage.
Manopt.jl on the opposite adds most things in terms of decorators or actions. This makes for example the stopping criteria much moree fleexible, but with this “looser” coupling exchanging/sharing things becomes more complicated

You have two ways

1. You could try to write a functor that does the cache within the functors struct, something like

A = [... your matrix ]
mutable struct MyFunc{A,T}
    A::T # you could also keep this in global scope
    b::T
end
function (f::Myfunc)(:Cost,M,p)
    f.b .= f.A*p
# do your computations based on the common thing b and compute cost c
    return c
end
function (f::MyFunc)(:Gradient, M, p)
# use f.b here and compute your gradient (tangent vector) X
return X
end

now initialising mF = MyFunc(A,b) you can use f(M,p) = mF(:Cost,M,p) as your cost and grad_f(M,p) = mF(:Gradient, M, p) a your gradient. The problem might be (but as far as I see Matlab has the same problem) if you call f(M,p) and then grad_f(M,p2) it might use the wrong cached value.
The advantage of this is, that you could continue that to even share something to the Hessian (for exactly one point)

2. Use the CostGradObjective and a Cache

if you define a function

function costgrad(M,p)
# Compute both c and X from the last code here
return (c,X)

and https://manoptjl.org/stable/plans/objective/#Manopt.ManifoldCostGradientObjective
you can even use the variant where you compute X in place.

Together with a Cache (as you linked) this should work even better (faster) than the Matlab variant, since Matlab only has one storage for Ax if I read correctly and this would have a storage for each point p (in the cache). This is also the variant that for example Optim.jl provides

In very general one could think about extending the Cache to be able to cache anything and not just calls to the objective, but that would require quite some design thoughts (and I already spent the last 3 months designing the Cache in easy-tp-use, easy-to-add and performant), so that is maybe more like a long time feature still. The main challenge there would still be to keep everything modular while accessing a globally shared memory.

Hope this helps :)

[kellertuer]

To limit number of calls the cache is the main thing to help.

[mateuszbaran]

What I mean is that the solver has to call costgrad!! instead of separate calls to cost and gradient to effectively use ManifoldCostGradientObjective . Automatic caches are not meant to store intermediates. Hand-made caches like your mutable struct MyFunc{A,T} are fairly hard to get right, and ManifoldCostGradientObjective is meant to solve this exact problem.

[mateuszbaran]

So I would try to improve solvers to take advantage of costgrad!! if possible.

[mateuszbaran]

So, we've had a discussion and I learned that cached CostGrad objective calculated gradient also when only cost is requested. Depending on the specific solver, it may be a good or a bad thing, so it's hard to tell whether option 1 or option 2 would be better.

[kellertuer]

...but it does not only also compute the gradient, it also caches said gradient.

We also discussed that I can give a detailed description of these two variants – and one that gets even beyond the Matlab-interims-storage (to interims caching namely even), in an extension of the Cache tutorial. But I might need some time for that.
It would introduce 3 more examples to the tutorial above and might be a bit technical to explain ;)

[kellertuer]

I was very curious and hacked a few examples (but be careful it still has typos in there and some math does not render), just now after work. See the extension I wrote at

https://manoptjl.org/previews/PR258/tutorials/CountAndCache/#Advanced-Caching-Examples

and I would recommend the storage approach in
https://manoptjl.org/previews/PR258/tutorials/CountAndCache/#A-shared-storage-approach-using-a-functor

the cache approach after that (for the storage and in the solver) might be not worth it for gradient descent things; but it is also not much worse.

[kellertuer]

We tested the updated docs carefully and they are included in the newest version, see

https://manoptjl.org/stable/tutorials/CountAndCache/#Advanced-Caching-Examples

if you have further questions – feel free to ask (Caching is sometimes a bit tricky for me for example).

[IvanBioli]

Hi,

your updated docs are amazing in explaining how to use the LRU cache in A local cache approach. However, I’m struggling to figure out how to cache multiple related values efficiently.

To be clearer, say that I want to cache both A*x and A*(A*x). In the MATLAB version of Manopt, I would do something like

% SOME CODE
if ~isfield(store, 'Ax')
    store.Ax = A * x
end
if ~isfield(store, 'AAx')
    store.AAx = A * store.Ax
end
% SOME OTHER CODE

Is there a way to achieve a similar result in Julia?

Best,
Ivan

[kellertuer]

That is one thing not yet fully generically available in Julia.

First you have to distinguish 2 things:

1. Caching within the objective – that is for example the cache= keyword or the local cache you link to. That should generically work for most components of an objective

2. Caching things the algorithm computes along the way – that is within the state. For that we have a “local, maybe shared” cache, see https://manoptjl.org/stable/plans/state/#Manopt.StoreStateAction.
   While for the objective an “easy cache wrapper” is doable (see 1) that is not so easy for the state, since the step would be inside and the cache “outside of that” in a wrapper and they would not communicate.
   The linked StoreAction is for example used to store previous gradients/iterates within CG coefficients

   Manopt.jl/src/plans/conjugate_gradient_plan.jl

   Lines 324 to 337 in 2aa8f62

   	function (u::DirectionUpdateRuleStorage{FletcherReevesCoefficientRule})(
   	amp::AbstractManoptProblem, cgs::ConjugateGradientDescentState, i
   	)
   	M = get_manifold(amp)
   	if !has_storage(u.storage, PointStorageKey(:Iterate)) ||
   	!has_storage(u.storage, VectorStorageKey(:Gradient))
   	update_storage!(u.storage, amp, cgs) # if not given store current as old
   	end
   	p_old = get_storage(u.storage, PointStorageKey(:Iterate))
   	X_old = get_storage(u.storage, VectorStorageKey(:Gradient))
   	coef = inner(M, cgs.p, cgs.X, cgs.X) / inner(M, p_old, X_old, X_old)
   	update_storage!(u.storage, amp, cgs)
   	return coef
   	end

but to do that for any interims computation I have not yet had a good idea how to approach that – and I have thought about that several times. So in the super generic way you want to do this, this is simply not available.
The simple reason is that I have no clue how that could be achieved.
If you really need it – maybe use Matlab.

[IvanBioli]

I thought a bit about it, and this could work for my usecase:

using LRUCache
@kwdef mutable struct myStore
    Ax::Union{AbstractVector, Nothing} = nothing
    A2x::Union{AbstractVector, Nothing} = nothing
end
struct CacheG{C,M}
    A::M
    cache::C
end


function (g::CacheG)(::Val{:Cost}, x)
    store = get!(myStore, g.cache, x)
    if isnothing(store.Ax)
        println("Computing Ax")
        store.Ax = g.A * x
    end 
    Ax = store.Ax
    return -x'*Ax
end
function (g::CacheG)(::Val{:Gradient}, x)
    store = get!(myStore, g.cache, x)
    if isnothing(store.Ax)
        println("Computing Ax")
        store.Ax = g.A * x
    end
    if isnothing(store.A2x)
        println("Computing A2x")
        store.A2x = g.A * store.Ax
    end 
    Ax = store.Ax
    return Ax
end

A = [1 2; 3 4]
mycache = CacheG(A, LRU{AbstractVector, myStore}(maxsize=10))

cost(p) = mycache(Val(:Cost), p)
grad(p) = mycache(Val(:Gradient), p)

x = [1, 2]
cost(x)
cost(x)
grad(x)
grad(x)

Does it cover also what you are mentioning? I did not fully understand your second point.

[kellertuer]

ah yes, that looks good.

Well you were not so precise what do cache.

For the first part (both the Caching system following https://manoptjl.org/stable/tutorials/CountAndCache/#Caching as well as the manual one you mentioned are about caching stuff while evalutating the objective (i.e. during function calls, gradient calls,...)

The second part was about caching stuff during the solver iterations (but outside the calls from the first case). That is what would be super tricky to impossible

Bu sure your approach would work, I would maybe not use AbstractVector but a concrete T<:AbstractVector for efficiency, any abstract field in a struct would be quite inefficient.