Rotation trick diverges when stochastic_sample_codes=True

[mahdip72]

@lucidrains Thanks for the brilliant repo. It’s been central to our new project: link

Recently, we are hitting divergence during training whenever rotation_trick and stochastic_sample_codes=True are both enabled, even with sample_codebook_temp=0.01, it does not fully solve the problem. However, disabling the rotation_trick restores stability. Is this combination expected to misbehave, or is there a workaround you’d recommend for us?

We are using these parameters for your reference:

dim: 64
decay: 0.99
codebook_size: 8196
commitment_weight: 0.25
orthogonal_reg_weight: 10
orthogonal_reg_max_codes: 512
orthogonal_reg_active_codes_only: True
rotation_trick: True
threshold_ema_dead_code: 3
kmeans_init: True
kmeans_iters: 10
stochastic_sample_codes: True
sample_codebook_temp: 0.1

[lucidrains]

@mahdip72 very cool application! so i think you don't need stochastic sampling of codes if you already have rotation trick. both address the need of increasing code utilization, but i think rotation trick achieves it more naturally

[lucidrains]

@mahdip72 are you seeing something different on your end where you want to turn both of them on?

[mahdip72]

@lucidrains I tried keeping both on, but in practice, we’re seeing the same pattern across two separate codebases (GCP-VQVAE and another VQ-VAE variant): with rotation_trick=True and stochastic sampling enabled, training runs fine for a few epochs and then blows up. Dropping sample_codebook_temp to 0.01 only delays the blow-up. I asked GPT-5 Codex for a sanity check and it suggested the gradient might be getting rotated toward essentially random directions when the sampled code jumps far from the input:

With stochastic sampling, the forward codeword can be far from the input vector. The rotation trick then forces the gradient to align with a largely random direction, amplifying updates instead of dampening them; this is consistent with the divergence you observe. When rotation_trick=False, the code falls back to STE (quantize = x + (quantize - x).detach()), so gradients stay identity-mapped and training remains stable. There is no guard to prevent the problematic combination.

Does that line up with how you expect the rotation trick to behave, or is there something else we should be looking at?

[lucidrains]

yea it may not be beneficial to rotate gradients when you draw some distant code

my question is why do you want the stochastic sampling turned on? is code utilization or some other metric even higher when you have both that and rotation trick on?

[mahdip72]

@lucidrains rotation trick does bump overall utilization a bit for us, but it doesn’t necessarily flatten the usage distribution. “Usage ratio” (PPL₁ / unique_codes) is really about the effective fraction of the codebook you pull from each step, while utilization just asks whether a code ever fires. I realized we can hit 100 % utilization yet keep a low usage ratio if a few codes dominate.

I’ve been digging into this because I’m working on a research project making VQ tokenizers more autoregressive-friendly. As you know, these papers use stochastic sampling methods:

1. Autoregressive Image Generation using Residual Quantization

2. Regularized Vector Quantization for Tokenized Image Synthesis

They show that stochastic sampling helps downstream AR modeling quite a bit (especially making this friendly with categorical sampling of AR models), so I’m trying to quantify multiple factors to build around it (usage ratio being an easy starting metric). That’s why I was experimenting with “rotation trick + sampling”. Ideally, we would keep the gradient benefits of the rotation while still getting the AR-friendly distribution that stochastic sampling brings, but the combo keeps diverging for us.

[lucidrains]

i see, i say you should try alternating the two settings every step and see if the divergence still happens

[mahdip72]

Here are the validation results I have got for different combinations in my new research project, which might be interesting for others:

Dark Blue plot:

rotation_trick: False
stochastic_sample_codes: False
sample_codebook_temp: 0.01

Light Blue plot:

rotation_trick: False
stochastic_sample_codes: True
sample_codebook_temp: 0.1

Green plot:

rotation_trick: False
stochastic_sample_codes: True
sample_codebook_temp: 0.5

Orange plot:

rotation_trick: False
stochastic_sample_codes: False

Red plot:

rotation_trick: True
stochastic_sample_codes: False

Other hyperparameters were identical in all trainings. sample_codebook_temp was constant throughout the training for those who had stochastic_sample_codes enabled. These are my other VQ layer parameters:

dim: 64
decay: 0.99
codebook_size: 8196
commitment_weight: 0.25
orthogonal_reg_weight: 10
orthogonal_reg_max_codes: 1024
orthogonal_reg_active_codes_only: True
threshold_ema_dead_code: 3
kmeans_init: True
kmeans_iters: 10

I also observed that when rotation_trick is enabled (Blue plot below), the gradient norm of VQ loss is higher:

@lucidrains Same as before, I could not make the combination of rotation_trick and stochastic_sample_codes work in any temperature; all blew up after some epochs.

As a conclusion, for this specific project, the default STE + sampling codebook with temperature 0.1 acts better than enabling rotation_trick.

Note:
The essence of this project was about using VQVAE for quantizing the continuous representation of a pretrained foundational model.

[lucidrains]

do you have a plot with both of them turned on before it diverges? your first plot, the x-axis is in thousands of steps? how about your final reconstruction loss?

[lucidrains]

@mahdip72 also, heard you are a friend of Alex (who was one of the most amazing collaborators i've ever met while implementing alphafold3)

small world

[mahdip72]

@mahdip72 also, heard you are a friend of Alex (who was one of the most amazing collaborators i've ever met while implementing alphafold3)

small world

Yeah, like you, he is brilliant.

do you have a plot with both of them turned on before it diverges? your first plot, the x-axis is in thousands of steps? how about your final reconstruction loss?

The utilization plot was based on epochs (each has 3906 steps). The metric I am monitoring on the validation set (same color as the utilization plot):

I am running some experiments for the rotation_trick + stochastic_sample_codes. Will drop them here soon.

[mahdip72]

I did the following two experiments in this setup:

dim: 64
decay: 0.99
codebook_size: 8196
commitment_weight: 0.25
orthogonal_reg_weight: 10
orthogonal_reg_max_codes: 1024
orthogonal_reg_active_codes_only: True
rotation_trick: True
threshold_ema_dead_code: 3
kmeans_init: True
kmeans_iters: 10
stochastic_sample_codes: True
sample_codebook_temp: 0.1

The only difference I made in them is enabling/disabling an experimental feature named adaptive_loss_coefficients to mitigate divergence, which is explained here.

Validation set
Red: adaptive_loss_coefficients was enabled.
Gray: adaptive_loss_coefficients was disabled.

Training set
Blue: adaptive_loss_coefficients was enabled.
Green: adaptive_loss_coefficients was disabled.

The reconstruction loss is computed as the sum of cosine and mse terms.
The unscaled term means that any scaling coefficients were removed during logging, allowing the visualization of raw, unweighted values.

Higher sample_codebook_temp values only make it worse.