[论文解读] Technical report: Training Mixture Density Networks with full covariance matrices

Jakob Kruse|arXiv (Cornell University)|Mar 4, 2020

Bayesian Methods and Mixture Models参考文献 1被引用 6

一句话总结

本文提出一种通过参数化下三角矩阵的Cholesky分解来训练具有完整协方差矩阵的混合密度网络（MDNs）的方法，从而能够建模复杂且相关的条件分布。该方法支持通过可微对数似然优化实现端到端训练，并通过可逆变换实现高效采样。

ABSTRACT

Mixture Density Networks are a tried and tested tool for modelling conditional probability distributions. As such, they constitute a great baseline for novel approaches to this problem. In the standard formulation, an MDN takes some input and outputs parameters for a Gaussian mixture model with restrictions on the mixture components' covariance. Since covariance between random variables is a central issue in the conditional modeling problems we were investigating, I derived and implemented an MDN formulation with unrestricted covariances. It is likely that this has been done before, but I could not find any resources online. For this reason, I have documented my approach in the form of this technical report, in hopes that it may be useful to others facing a similar situation.

研究动机与目标

• 使混合密度网络能够使用完整协方差矩阵对条件概率分布进行建模，克服对角协方差假设的局限性。
• 提供一种可微且数值稳定的完整协方差矩阵参数化方法，适用于深度学习中的反向传播。
• 通过精度矩阵的Cholesky分解，支持具有无限制协方差的MDN的高效采样与训练。
• 记录并分享一个实用实现，供面临类似条件密度估计挑战的研究人员使用。

提出的方法

• 通过Cholesky分解使用下三角矩阵U对每个高斯分量的精度矩阵（逆协方差矩阵）进行参数化，确保其正定性。
• 利用Cholesky因子U计算对数密度贡献：−½‖U(x−μ)‖²₂ + Σⱼ log(Uⱼⱼ)，其中Uⱼⱼ为对角元素。
• 通过最小化数据在混合模型下的负对数似然来训练MDN，使用基于Jensen不等式的上界以提升训练稳定性。
• 通过反转Cholesky因子实现可微采样，生成样本x = μ + Lη，其中L = U⁻ᵀ，且η ∼ N(0, I)。
• 将完整协方差MDN作为可逆模块集成到归一化流中，使用神经网络输出的条件输入ω、μ和U。
• 使用PyTorch和FrEIA框架实现该模型，通过负对数似然损失实现GMM模块的反向传播。

实验结果

研究问题

• RQ1 混合密度网络能否通过使用完整协方差矩阵而非对角协方差矩阵，有效建模复杂且相关的条件分布？
• RQ2 如何在深度学习中以可微且数值稳定的方式参数化完整协方差矩阵？
• RQ3 在条件密度估计中，使用完整协方差矩阵对MDN的训练稳定性和性能有何影响？
• RQ4 完整协方差MDN能否在归一化流等可逆架构中集成，同时保持可微性与高效采样？

主要发现

• 所提出的精度矩阵Cholesky参数化确保了正定性，并实现了具有完整协方差矩阵的MDN的可微训练。
• 利用Cholesky因子U计算对数密度，可通过‖U(x−μ)‖²₂和迹项实现混合模型的稳定且高效的评估。
• 该方法通过使用基于Jensen不等式的可微负对数似然目标，支持端到端反向传播训练，并提升早期训练的稳定性。
• 通过反转精度矩阵的Cholesky因子，可高效生成多变量高斯分量的样本。
• 该实现已作为可逆GMM模块集成到FrEIA框架中，支持在归一化流和条件密度建模中的使用。
• 该方法为需要建模相关输出的研究人员提供了一种实用且可重用的完整协方差MDN解决方案。