Comment by elpalek

have you thought about SPALDE models? ex: https://arxiv.org/abs/2109.10086