mechanize + ema usage

[aaarrti]

Hi,

I want to use optax.contrib.mechanize together with the exponential moving average.
What would be the correct way to do it?
Option 1:

tx = optax.contrib.mechanize(optax.chain(optax.ema(...), optax.lion(...)))

Option 2:

tx = optax.chain(optax.ema(...), optax.contrib.mechanize(optax.lion(...)))

I would also appreciate it much, if you could share a bit of theory/background on why an option is correct, or just share any related resources. Thanks in advance!

[fabianp]

Perhaps @adefazio has some insights on this one?

[adefazio]

Mechanic is designed to work with any optimization algorithm as the base optimizer, so you should use the EMA within Mechanic (option 1) rather than wrapping it. The LR estimation may not work correctly if you apply the EMA via option 2, or at least the theory doesn't tell us what will happen.