Remove the type ParamSpaceSGD

[Red-Portal]

This PR removes the use of the type ParamSpaceSGD, which provides a unifying implementation of VI algorithms that run SGD in parameter space. Instead, each parameter space SGD-based VI algorithm becomes its own AbstractVariationalAlgorithm, where the shared code implementing step is shared by dispatching over their Union.

This addresses #204

[Red-Portal]

Hi @yebai , could you check to make sure that this is what you asked for? Personally, I feel the ParamSpaceSGD-based interface is much cleaner and intuitive, especially in terms of project structure. So I still insist we keep it as an internal implementation detail.

[yebai]

Thanks @Red-Portal -- I left some comments above. In addition, let's simplify the folder structure a bit for clarity:

1. move all files in paramsspacesgd to algorithms, eg, "algorithms/paramspacesgd/constructors.jl" to "algorithms/constructors.jl"
2. keep each algorithm in its own file

Also, I'd suggest we consider renaming paramspacesgd.jl to interface.jl or something along the lines:

• "algorithms/paramspacesgd/paramspacesgd.jl" to "algorithms/interface.jl"

[Red-Portal]

Hi Hong, I planned to do the restructuring in a separate PR to keep things simple in this one. Though:

move all files in paramsspacesgd to algorithms, eg, "algorithms/paramspacesgd/constructors.jl" to "algorithms/constructors.jl"

"algorithms/paramspacesgd/paramspacesgd.jl" to "algorithms/interface.jl"

After the release of v0.5, we'll be adding algorithms that don't conform to the original ParamSpaceSGD formalism so I think these namings are not gonna withstand that change. In fact, part of the reason I kept everything under paramspacesgd/ was precisely for this reason.

[yebai]

It is okay to keep all algorithms under algorithms and remove the extra subfolder paramspacesgd:

1. These categorisations are nonstandard, so they are not helping clarity.
2. There aren't many algorithms, so it is okay to keep all of them under the same algorithms folder.

After the release of v0.5, we'll be adding algorithms that don't conform to the original ParamSpaceSGD formalism

You can add more algorithms to interface.jl, so long as these algorithms are clearly grouped in interface.jl.

[Red-Portal]

After the release of v0.5, we'll be adding algorithms that don't conform to the original ParamSpaceSGD formalism

You can add more algorithms to interface.jl, so long as these algorithms are clearly grouped in interface.jl.

I am saying that these new algorithms can't be grouped in interface.jl since they will need their custom implementation of step. The current algorithms all go through the same step, which is why they can be grouped.

[yebai]

"grouping" refers to grouping interface code together for similar VI algorithms in the proposed interface.jl. It doesn't require creating a new union type if an algorithm is distinct from others. In these cases, the algorithm could be a singleton group.

[sunxd3]

If I understand right, this PR flattens the existing AbstractVariationalAlgorithm → ParamSpaceSGD → (concrete algorithms) hierarchy. That intermediate layer exists today so that anything doing “SGD in parameter space” can share one abstraction.

Before we drop it, may I understand what concrete benefits the flattening delivers? In particular, are we planning to add other algorithm families alongside the current ParamSpaceSGD? If so, we then probably need to add the middle layer abstraction back.

At the moment, I haven't quite convince myself that the flattening of the type hierarchy is necessary.

[yebai]

I suggested the removal of the ParamSpaceSGD because it is not standard terminology in the variational inference literature. Adding this less well-known terminology adds mental overhead to understanding the code.

[Red-Portal]

@yebai @sunxd3 Thanks for chiming in. Actually, I have a new idea. So I believe the main complaint at the moment is that the term ParamSpaceSGD is not intuitive as an abstraction. What if we fix that directly: by changing the name to something more obvious.

In a nutshell, the nice thing about the current ParamSpaceSGD interface is that you only need to define a gradient estimator (estimate_gradient) of a corresponding objective to form an algorithm. So let me make that more explicit in the name. Here are a few candidates:

• ObjectiveSGD
• ObjectiveDescentInducedAlgorithm
• MinObjectiveSGD
• MinObjectiveWithSGD
• MinObjectiveBySGD
• ObjectiveGradientEstimateDescent

or something along these lines? Would that resolve your concern?

[yebai]

I think I see your point. But, I am not sure that helps.

define a gradient estimator (estimate_gradient) of a corresponding objective to form an algorithm.

That probably includes every learning algorithm in ML.

[Red-Portal]

define a gradient estimator (estimate_gradient) of a corresponding objective to form an algorithm.

That probably includes every learning algorithm in ML.

Yes, that is indeed almost true! But the point is that there are a couple of important algorithms that don't quite conform to this formalism, as they result in a custom update rule. They don't fall out of a gradient estimator, but modify the parameter update step too. So this is the reason I wish to allow for two different abstraction levels. But as you said, most algorithms only require defining a gradient estimator. So the lower-level interface helps unify the code for all those algorithms.

[yebai]

a couple of important algorithms that don't quite conform to this formalism, as they result in a custom update rule. They don't fall out of a gradient estimator, but modify the parameter update step too.

We are at risk of premature abstraction and introducing heuristic terminology here. It is better to work with concrete algorithms, and define a union type if sharing code is needed (eg, step) across algorithms.

There might be some insights we can learn by taking a unifying view of parameter space gradient descent VI, but that is a discussion we should have offline for a review paper.

[Red-Portal]

We are at risk of premature abstraction and introducing heuristic terminology here. It is better to work with concrete algorithms, and define a union type if sharing code is needed (eg, step) across algorithms.

My main beef with using Unions here is the following:

• This results in duplicating code, as can already be seen in the PR (the structs KLMinRepGradDescent, KLMinRepGradProxDescent, KLMinScoreGradDescent all have the same fields.)
• To use the shared step function, a certain <:AbstractVariationalAlgorithm struct has to contain a certain list of fields, which is now implicit since we're not using a proper interface.
• I am not sure whether I should document the use of the shared step function. But if I do document it, it's not going to be pretty since it assumes a whole lot of implicit things (as stated in the item above).
• The structuring of the project will need a bit of discussion since things are a bit more complicated. (What should we call the file containing the shared step function? How should we structure the directories?) As mentioned in a previous comment, we can't just use generic names since I will be adding algorithms that don't use the shared step function in the following PRs. Also, people first looking at the code base will probably have to do some mental gymnastics since lots of things are made implicit but not explicit through some interface.

With that said, do you find the solution below still unsatisfying? At least I hope that this resolves your concern that the terminology is non-standard.

So let me make that more explicit in the name. Here are a few candidates:

• ObjectiveSGD
• ObjectiveDescentInducedAlgorithm
• MinObjectiveSGD
• MinObjectiveWithSGD
• MinObjectiveBySGD
• ObjectiveGradientEstimateDescent
  or something along these lines? Would that resolve your concern?

If you think we should still go with an implicit interface, then I'll follow for the sake of moving forward.

[Red-Portal]

Hi @yebai , what is your final verdict on what we should do here given my last comment above?

[yebai]

Hi @ Red-Portal, I don't think the proposed interfaces are better. For clarity, I'd suggest we remove paramspacesgd types. It's okay to keep the step etc functions in interface.jl instead.

Happy to continue the discussion offline and explore whether parameter space can be useful for novel algorithmic insights.

[Red-Portal]

Hi @yebai , I applied all the changes you requested. Let me know if you're satisfied with the code now. Updating the docs will be separate since it will be a lot. (And painful on my end 😅 )

[yebai]

Many thanks, @Red-Portal, for the contributions. I'm happy with the changes; @penelopeysm / @mhauru might want to review this too.

[github-actions]

Benchmark Results

Benchmark suite	Current: 7d3ed86	Previous: 9b2eabb	Ratio
normal/RepGradELBO + STL/meanfield/Zygote	3689036472.5 ns	3990679915 ns	0.92
normal/RepGradELBO + STL/meanfield/ReverseDiff	970375750 ns	1124955886 ns	0.86
normal/RepGradELBO + STL/meanfield/Mooncake	913048608 ns	1261092892.5 ns	0.72
normal/RepGradELBO + STL/fullrank/Zygote	3670088355.5 ns	3971814683.5 ns	0.92
normal/RepGradELBO + STL/fullrank/ReverseDiff	1489955135.5 ns	1621498433.5 ns	0.92
normal/RepGradELBO + STL/fullrank/Mooncake	972304504 ns	1285149814.5 ns	0.76
normal/RepGradELBO/meanfield/Zygote	2365990011 ns	2859724373.5 ns	0.83
normal/RepGradELBO/meanfield/ReverseDiff	660842017.5 ns	777886417 ns	0.85
normal/RepGradELBO/meanfield/Mooncake	799140056 ns	1137178738 ns	0.70
normal/RepGradELBO/fullrank/Zygote	2401497236 ns	2864869952 ns	0.84
normal/RepGradELBO/fullrank/ReverseDiff	831429492.5 ns	954006564 ns	0.87
normal/RepGradELBO/fullrank/Mooncake	842196960.5 ns	1144634418 ns	0.74
normal + bijector/RepGradELBO + STL/meanfield/Zygote	5641882447 ns	5743469386 ns	0.98
normal + bijector/RepGradELBO + STL/meanfield/ReverseDiff	1944678817 ns	2419710523 ns	0.80
normal + bijector/RepGradELBO + STL/meanfield/Mooncake	3867940233 ns	4227822021.5 ns	0.91
normal + bijector/RepGradELBO + STL/fullrank/Zygote	5657272286 ns	5787110379 ns	0.98
normal + bijector/RepGradELBO + STL/fullrank/ReverseDiff	2646661599.5 ns	3116964174 ns	0.85
normal + bijector/RepGradELBO + STL/fullrank/Mooncake	3997669962.5 ns	4366946807.5 ns	0.92
normal + bijector/RepGradELBO/meanfield/Zygote	4058043520 ns	4464663738.5 ns	0.91
normal + bijector/RepGradELBO/meanfield/ReverseDiff	1570570758 ns	2102610294 ns	0.75
normal + bijector/RepGradELBO/meanfield/Mooncake	3674711961 ns	4074397467 ns	0.90
normal + bijector/RepGradELBO/fullrank/Zygote	4150263319.5 ns	4560198939 ns	0.91
normal + bijector/RepGradELBO/fullrank/ReverseDiff	1831914975 ns	2331323017 ns	0.79
normal + bijector/RepGradELBO/fullrank/Mooncake	3821471618.5 ns	4175468965.5 ns	0.92

This comment was automatically generated by workflow using github-action-benchmark.

[penelopeysm]

hi @yebai @Red-Portal I'm a bit ill now but if you give me a few days I'll be happy to take a look, if it's not a rush

[Red-Portal]

@penelopeysm Sure thing, please take care!

[penelopeysm]

I only have one comment, which is that it's probably better to keep this as a struct (although I recognise you may have some difficulty in choosing a name for it). NamedTuples are too flexible, if you return one you don't know what fields are in it.

[Red-Portal]

I propose we stick with this for now and improve this later once we can make things more organized. (The new structure in the PR is making a whole lot of things implicit anyway, so I think the fields being implicit doesn't make things much worse)

[penelopeysm]

Sure, I'm not that fussed.

[Red-Portal]

@yebai @penelopeysm Are we good to go?

edit: Penny told me that this is a yes.

[github-actions]

AdvancedVI.jl documentation for PR #205 is available at:
https://TuringLang.github.io/AdvancedVI.jl/previews/PR205/