Nice code! I got it to work after a couple fixes:




instead of:

params = params.div(norm-EPS)

do:

params = params.div(norm) - EPS

θ/∥θ∥−ε  is what it says in the paper – EPS needs to be outside the denominator, otherwise you’re increasing the big embeddings instead of decreasing them.




instead of:

gamma = gamma.clamp(min=1)

do:

gamma = gamma.clamp(min=1+EPS)

This ensures that there is a minimal distance between equivalent embeddings (same word). Even though this seems counter-intuitive, doing this avoids getting a zero-divisor in the derivative (see the definition of γ in the partial derivative of Poincaré distance in the paper).