Got 'nan' when using MetaAdam as a inner_opt

[hccz95]

I just made two modifications on examples/few-shot/maml-omniglot.py:

1. change the net structure in line 98 to my net:

net = nn.Sequential(
    nn.Conv2d(1, 64, (3, 3), 2, 0), nn.BatchNorm2d(64), nn.ReLU(),
    nn.Conv2d(64, 64, (3, 3), 2, 0), nn.BatchNorm2d(64), nn.ReLU(),
    nn.Conv2d(64, 64, (3, 3), 2, 0), nn.BatchNorm2d(64), nn.ReLU(),
    nn.Conv2d(64, 64, (2, 2), 1, 0), nn.BatchNorm2d(64), nn.ReLU(),
    nn.Flatten(), nn.Linear(64, args.n_way),
).to(device)

2. change the inner_opt from MetaSGD to MetaAdam in line 120.

Then I got 'nan' parameters after first time of meta_opt.step()

[JieRen98]

Hi,

The nan error is because of the numerical instability of the Adam in meta-learning. There exist an sqrt operation according to the Adam paper. However, the second derivation of the sqrt operation is not numerical stable, i.e. $\frac{d^2 sqrt(x)}{dx^2}|_{x=0}=\text{NaN}$. You can also refer to facebookresearch/higher#125.

For this problem, we have two recommended solutions.

• Put the sqrt operation into the whole equation, and compute the derivation of the output to the input manually. The second derivation of the sqrt operation will be eliminated. You can achieve this by setting the flag use_accelerated_op=True, you can follow the instructions in Functional Optimizer and Meta Optimizer.
• Register hook to the first-order gradients. During the back-propagation, the NaN gradients will be set to 0, which will have a similar effect to the first solution but much slower. Check the functions in https://github.com/metaopt/TorchOpt/blob/main/torchopt/_src/hook.py.

Best regards,
Jie

[hccz95]

Thank you for explaining this, it's very useful!