How to update EMA only at desired steps

[dohyun1411]

Hi,

I am training my model (VQ + other models) with gradient accumulation.
Since the VQ module uses EMA updates at every forward pass,
it is difficult to align them with the gradient updates.

What I want is to accumulate the necessary data during forward passes,
but defer the actual EMA update until a specific step (e.g., when gradients are applied).

I also noticed a parameter, manual_ema_update.
Can this be used to manually trigger EMA updates at the right time?

[lucidrains]

@dohyun1411 yea it can! it was originally done for another feature, but can be repurposed

let me think about it

[lucidrains]

@dohyun1411 hey Dohyun, do you want to see if this works for you on the latest version?

[dohyun1411]

Thanks for the quick improvement. I think there are a few related details worth considering.
For example, "k-means init" now runs on the very first mini-batch. But with gradient accumulation, the actual “first” batch is the effective batch formed after several mini-batches are combined into one optimizer step. It would be better to delay k-means init until that effective batch is ready.
The same applies to "expire codes," which should ideally also be evaluated over aggregated data rather than a single mini-batch.
Thanks again with your great work!

[lucidrains]

@dohyun1411 thanks for the suggestions! will get around to those later this month!