I’m developing an extended/augmented version of the “Deep InfoMax” model described in this paper: https://openreview.net/forum?id=Bklr3j0cKX 4
While looking for the root cause of differences between 0.4.1 and 1.0.0, I stripped the model down until it was approximately the same as the “local DIM” model in the ICLR paper, modulo some data augmentation tricks.
I’m measuring model “performance” during training in terms of two tasks – (i) a log-likelihood loss from a noise-contrastive estimation objective that drives self-supervised representation learning, and (ii) a classifier which is trained at the same time using the features learned by the self-supervised objective. The classifier is a simple MLP with one ReLU hidden layer and does not backprop its loss into the features provided by the encoder. The classifier is for monitoring the quality of features learned via self-supervised learning.
The NLL for the classifier and the encoder’s NCE cost are both significantly worse when trained in 1.0.0 than in 0.4.1, whether I compute the losses in fp32 or fp64. If I cast the relevant tensors to fp64 prior to computing the NCE and classifier losses, and then cast back to fp32 after computing the losses, the performance in 0.4.1 is unaffected but the performance in 1.0.0 improves relative to the case where I use fp32 inside the loss computations. I believe everything is computed on the GPU, but it’s possible that the pytorch backend silently shifts some things off GPU.
The main step in computing the NCE cost is to do F.log_softmax on a large tensor of pair-wise “scores” between “global” and “local” features taken from all the inputs in a minibatch. The softmax normalization is over a set of 1k-10k (potentially) high-variance logits. The classifier cost is pretty standard, and just passes the MLP’s output logits to an nn.CrossEntropy instance.
Training in 0.4.1 has always been very stable, including while developing the models described in the ICLR paper. I did not encounter any instability until I started working with 1.0.0. One weird guess for what could cause the problem is that “improvements” in numerical stability may be causing some grads/losses which underflow to 0 in 0.4.1 to no longer underflow, and these grads are the ones that explode in backprop and cause the instability. But, idk.