This is great, would love to see an arxiv-style writeup. Include this in comparisons? http://proceedings.mlr.press/v130/shekelyan21a/shekelyan21a.pdf