0
$\begingroup$

I am currently trying to make an autoencoder that compresses a 3D volume where each value represents the density of said volume. The architecture is a UNet without skip connections. The optimizer is AdamW with a learning rate of 0.0002 and beta1 of 0.5. I use three losses: MSE between the input and reconstructed, the 3d frobenius norm of the reconstructed density and a segmentation loss where I want to areas with non-zeros density to match.

My issue is that the autoencoder does not seem to converge. After a few epochs, it will get results that are closer to the initial input, but if I train it for too long, the loss shoots up and the output looks nothing like the input.

From your experience, what could be causing this? What are the tools/tips that I can use to mitigate non-convergence? What questions should I be asking myself?

Here's a plot of the MSE error for my validation set.

enter image description here

Improve this question
$\endgroup$
2
  • $\begingroup$ Usually using a lower learning rate and a higher value of beta1 can help the optimizer converge more quickly and also UNet architectures are typically used for image segmentation tasks, and they may not be well-suited for compressing 3D density volumes. $\endgroup$
    – Vic
    Commented Dec 12, 2022 at 12:52
  • $\begingroup$ I'm giving it a shot with a smaller learning rate. Regarding UNet, I saw that the Stable Diffusion paper used a similar architecture, they used a UNet to compress 2d images. $\endgroup$
    – truvaking
    Commented Dec 12, 2022 at 13:14

1 Answer 1

Reset to default
0
$\begingroup$

Before compressing 3D volumes, I recommend compress 2D surfaces as it is simpler for identifying divergence causes and it would be also easy to scale up to 3D.

Then, you can try SGD or RMSProp instead of AdamW. Even if Adam is a great algorithm, it is less robust than SGD: Adam uses a different update rule that incorporates moving averages of the model's gradients and second moments. This can cause Adam to oscillate more, especially when the data is noisy or the optimization landscape is heavily non-convex. In contrast, SGD only uses the gradient of the loss function, which can make it more stable and easier to tune.

Finally, you can use Velo, which is a learned optimizer. It should be able to solve complex convergence issues as you have.

https://github.com/google/learned_optimization/tree/main/learned_optimization/research/general_lopt

Improve this answer
$\endgroup$
1
  • $\begingroup$ Does it answer your question? If not, please let me know to provide additional information. $\endgroup$
    – Nicolas Martin
    Commented Dec 13, 2022 at 13:12

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.