Log Z term in loss

[sh0416]

I think it is intended for numerical stability, but I don't know how it works.

Could you explain it or provide a reference for that code?

jaxformer/jaxformer/models/decoder/inter/model.py

Line 71 in 9c41fd4

loss += z_loss * jax.lax.square(log_z)

[sh0416]

I think it prevents the logsumexp from being too large.. Is there any reference for doing this? or just a practical issue? I derived it to the following equation..

x_i - logsumexp(x)(1-0.0001*logsumexp(x))

[sh0416]

In the mentioned paper (https://arxiv.org/abs/2206.13517), regularization term is added to prevent the divergence.
Is it right? @enijkamp Can I get further details about the loss term?