Hi,
A possible solution is to scale-down the values one of the two matrices, or both of them before doing matmul operation, you can try something similar to what in Softmax-Attention.
for example you can try doing it this way:
f = torch.matmul(theta_x / math.sqrt(your_hidden_size), phi_x)