CAVI-CMN trains conditional mixture networks via coordinate-ascent variational inference (CAVI), giving closed-form updates and avoiding gradient-based optimization entirely. The result: fast convergence, well-calibrated uncertainty, and competitive accuracy on classification tasks — without ever computing a gradient through the network.

Built with custom JAX conjugate variational inference code. Read the paper or check out the code on GitHub.