ITADN

Add covariance_type option to GMM class

#49Closeddominik-strutz 创建于 2024-04-02
enhancement
D
dominik-strutzcommented
### Description Currently, the `zuko.flows.mixture.GMM` class only supports full covariance matrices. However, there are a number of use cases (especially high-dimensional) where a full covariance matrix is either not needed or infeasible to estimate. This issue proposes to add the option to choose between different covariance matrix types similar to [`sklearn.mixture.GaussianMixture`](https://scikit-learn.org/stable/modules/generated/sklearn.mixture.GaussianMixture.html) Here is an example of how different covariance types approximate a mixture of 3 Gaussians with varying covariance matrices. ![c6263e0a-5d78-45ef-8df4-414692a7871f](https://github.com/probabilists/zuko/assets/42805470/489d6f44-ebc2-4423-b09e-8537cbc22bc1) ### Implementation The current structure of the GMM `zuko.flows.mixture.GMM` class makes it very easy to add the above mentioned enhancements. I have implemented the changes in a [fork of the repository](https://github.com/dominik-strutz/zuko/commit/004db1e26cd393222efc258a6590b5749f6f4002) and could open a pull request if this change is wanted. I have only tested the code for the unconditional case, but I do not see any way I could break it when adding context features. ### Further improvements When generating the above figure, I (again) realised how easily mode collapse happens for GMMs. The `zuko.flows.mixture.GMM` class could, therefore, also benefit from some sort of initialisation procedure, again, similar to [`sklearn.mixture.GaussianMixture`](https://scikit-learn.org/stable/modules/generated/sklearn.mixture.GaussianMixture.html). I fully understand if that goes beyond the scope of what Zuko wants to achieve. The benefit is that Zuko is very convenient to use and ties in so well with Pytorch code that having such a procedure here could be nice. However, it might add another dependency (e.g., sklearn) if you want to use existing implementations of initialisation algorithms. I have some basic implementation of this (using sklearn) lying around and would be happy to polish it up and make another commit if this is wanted.
关闭于 2025-03-03 2 条评论