MARS gamma parameter

[tom-jod]

In the implementation, I can't see a way to dynamically change the gamma parameter in the MARS variance reduction feature. It seems to have been hard-coded as the value from the paper (0.025) and is not being optimized.

• If I wanted to optimize gamma could I just add it into the list of optimized params in the trial function from benchmarks/utils.py
• If I additionally wanted to set up a schedule for gamma how would you recommend doing that?

Thanks!

[ClashLuke]

I'm not sure; these are excellent questions!

2. You can manually modify the gamma parameter before/after every optimizer pass using the opt.parameter_groups. It already uses compiler guards, so it won't recompile or cause other overheads. For example, the following code should work:

opt = heavyball.AdamW(params, lr=1, mars=True)
for i in range(100):
  opt.step(closure)
  for group in opt.param_groups:
    group["mars_gamma"] = i / 100

1. All heavyball optimizers are wrapped in no_grad(), so this won't work. However, the approach is very interesting! Adapting _compilable_mars_correction_ to manually compute the gradient for gamma (and update it) would likely be easiest. The closest existing functions you could reference are _compilable_d_adapt_ and _compilable_lr_adapt_.

Let me know if you need anything else :)

[tom-jod]

Thanks Luke! I'm currently looking at expanding on the MARS implementation in HeavyBall and just wanted to clarify some things:

1. Am I correct in thinking that the implementation actually uses MARS-approx where we simply use the previous gradient update rather than doing backprop again with the old parameters on the new batch?
2. The MARS algorithm uses gradient clipping after applying the variance reduction update in c (algorithm 1 in https://arxiv.org/html/2411.10438v1#S3.E13), is this carried out in the current mars implementation? If not could it be done by using the gradient clipping flag or would that apply it before the variance reduction step?
3. How can I run standard SGD in the repo for a baseline comparison?

Thanks,
Tom