Paper accepted at NeuRIPS 2021.