bhoov / energy-transformer-jax Goto Github PK

Hi, thank you for sharing very intreating works.
I have take a look the code, and I am confused about the potential inconsistent with code and equation described in the paper.

[1] Code seems to sum over all N by N attention matrix, where N is the number of token.

A = jax.nn.logsumexp(beta * jnp.einsum("qhz,khz->hqk", Q, K), -1)
return -1/beta * A.sum()

[2] Similarly, equation seems to calculate the energy via the attention matrix A, but does not sum the attention weights when the query and key come from the same token.

In that case, don't we need to remove the diagonal elements of the attention matrix?

Thanks in advance.