Numerical stability of attention layers

Hello!

Thank’s for a cool implementation! I’ve been playing around with the code on CIFAR-10 trying to replicate SOTA fid values. However, I get some nasty loss explosions all of a sudden in the middle of training in fp16 with sensible configs, e.g.,

class BaseUnet32(Unet):
    def __init__(self, *args, **kwargs):
        kwargs.update(
            dict(
                dim=192,
                dim_mults=(1, 2, 4),
                num_resnet_blocks=2,
                layer_attns=(False, True, True),
                layer_cross_attns=False,
                attn_heads=1,
                ff_mult=2.0,
            )
        )
        super().__init__(*args, **kwargs)

I’ve looked into attention implementation in guided-diffusion and it seems that they do scaling by 1/sqrt(dim) a bit differently by redistributing it as 1/sqrt(sqrt(dim)) in q and k (link), as well as doing softmax on float32 inputs. Testing it at the moment, will let you know if it helps

[KhrulkovV]

@lucidrains Hey, I set attention heads to 4 and lowered learning rate, going good so far (both fp32 and fp16 runs are roughly identical)

[KhrulkovV]

But the boost was really nice - 1200 img/sec to 1800 img/sec until it diverged 😹