For all models and datasets we run the same grid search and optimize the distortion loss, applying Rsgd. The implementation of all baselines are taken from Geoopt [1]}. We train for 3000 epochs, reducing the learning rate by a factor of 5 if the model does not improve the performance after 50 epochs, and early stopping based on the average distortion if the model does not improve after 150 epochs. We use the burn-in strategy [2]}, [3]} training with a 10 times smaller learning rate for the first 10 epochs.
We experiment with learning rates from \(\lbrace 0.02, 0.01, 0.005\rbrace \) , batch sizes from \(\lbrace 128, 512\rbrace \)  and max gradient norm from \(\lbrace 100, 300\rbrace \) .
