高斯分布的KLD总结

---

前置基础

矩阵内积

给定两个$m\times n$的矩阵$\mathbf{A}$和$\mathbf{B}$，其矩阵内积（也称为Frobenius inner product）定义为：
$$<\mathbf{A},\mathbf{B}>=\sum _{i=1}^m\sum _{j=1}^n{a}_{ij}{b}_{ij}=tr({\mathbf{A}}^T\mathbf{B})$$

分布与期望

给定随机变量$X\sim \mathcal{N}(\mathbf{\mu },\mathbf{\Sigma })$，则有如下等式成立：

1. $E[\mathbf{x}{\mathbf{x}}^T]=\mathbf{\Sigma }+\mathbf{\mu }{\mathbf{\mu }}^T$
2. $E[{\mathbf{x}}^T\mathbf{Ax}]=tr(\mathbf{A\Sigma })+{\mathbf{\mu }}^T\mathbf{A}\mathbf{\mu }$
   $proof:$ $E[{\mathbf{x}}^T\mathbf{Ax}]=E[tr({\mathbf{x}}^T\mathbf{Ax})]=E[tr(\mathbf{Ax}{\mathbf{x}}^T)]=tr(\mathbf{A}E[\mathbf{x}{\mathbf{x}}^T])=tr(\mathbf{A}(\mathbf{\Sigma }+\mathbf{\mu }{\mathbf{\mu }}^T))=tr(\mathbf{A\Sigma })+{\mathbf{\mu }}^T\mathbf{A}\mathbf{\mu }$
3. $E[(\mathbf{x}-{\mathbf{\mu }}_1{)}^T\mathbf{A}(\mathbf{x}-{\mathbf{\mu }}_1)]=tr(\mathbf{A\Sigma })+(\mathbf{\mu }-{\mathbf{\mu }}_1{)}^T\mathbf{A}(\mathbf{\mu }-{\mathbf{\mu }}_1)$

---

KLD定义

给定两个连续时间概率分布的概率密度函数分别为$p(x)$和$q(x)$，其KLD定义为：
$$D_{KL}(P||Q)=\int p(x)log(\frac{p(x)}{q(x)})dx$$
对于离散变量，给定两个概率分布$P(x)$和$Q(x)$，KLD定义为：
$$D_{KL}(P||Q)=\sum _{x}P(x)log(\frac{P(x)}{Q(x)})$$

一元高斯分布

假设连续时间的两个分布均为高斯分布，其中$P$分布均值${\mu }_1$，方差为${\sigma }_1$，$Q$分布均值${\mu }_2$，方差为${\sigma }_2$，则可以推导对应的KLD：
$$\begin{array}{c}\begin{array}{rl}{D}_{KL}(P||Q)& =\int p(x)log(\frac{p(x)}{q(x)})dx\\ & =\int \frac{1}{\sqrt{2\pi {\sigma }_1^2}}{e}^{-\frac{(x-{\mu }_1{)}^2}{2{\sigma }_1^2}}[log(\frac{{\sigma }_2}{{\sigma }_1})-\frac{(x-{\mu }_1{)}^2}{2{\sigma }_1^2}+\frac{(x-{\mu }_2{)}^2}{2{\sigma }_2^2}]dx\\ & =log(\frac{{\sigma }_2}{{\sigma }_1})-\frac{Var(x)}{2{\sigma }_1^2}+\frac{Var(x)+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}\\ & =log(\frac{{\sigma }_2}{{\sigma }_1})-\frac{1}{2}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}\end{array}\end{array}$$

多元高斯分布

对于$n$维随机变量$X$，假设$P$和$Q$分别满足$\mathcal{N}({\mathbf{\mu }}_1,{\mathbf{\Sigma }}_1)$和$\mathcal{N}({\mathbf{\mu }}_2,{\mathbf{\Sigma }}_2)$的分布，则其KLD推导如下：
$$\begin{array}{c}\begin{array}{rl}& D_{KL}(P||Q)\\ =& {\int }_{{\mathbb{R}}^n}p(\mathbf{x})log(\frac{p(\mathbf{x})}{q(\mathbf{x})})d\mathbf{x}\\ =& \int p(\mathbf{x})log(p(\mathbf{x}))d\mathbf{x}-\int p(\mathbf{x})log(q(\mathbf{x}))d\mathbf{x}\\ =& -\frac{1}{2}(log(2{\pi }^n|{\mathbf{\Sigma }}_1|)+E[(\mathbf{x}-{\mathbf{\mu }}_1{)}^T{\mathbf{\Sigma }}_1^{-1}(\mathbf{x}-{\mathbf{\mu }}_1)])+\frac{1}{2}(log(2{\pi }^n|{\mathbf{\Sigma }}_2|)+E[(\mathbf{x}-{\mathbf{\mu }}_2{)}^T{\mathbf{\Sigma }}_2^{-1}(\mathbf{x}-{\mathbf{\mu }}_2)])\\ =& \frac{1}{2}[log(\frac{|{\mathbf{\Sigma }}_2|}{|{\mathbf{\Sigma }}_1|})-n+tr({\mathbf{\Sigma }}_2^{-1}{\mathbf{\Sigma }}_1)+({\mathbf{\mu }}_1-{\mathbf{\mu }}_2{)}^T{\mathbf{\Sigma }}_2^{-1}({\mathbf{\mu }}_1-{\mathbf{\mu }}_2)]\\ =& \frac{1}{2}[<{\mathbf{\Sigma }}_2^{-1},{\mathbf{\Sigma }}_1>+||{\mathbf{\mu }}_1-{\mathbf{\mu }}_2|{|}_{{\mathbf{\Sigma }}_2^{-1}}-log(|{\mathbf{\Sigma }}_2{|}^{-1}|{\mathbf{\Sigma }}_1|)-n]\end{array}\end{array}$$

测试验证

% Generate sample data
% case 1
mu_p = [0.5, 1.0]';
sigma_p = diag([1.2, 0.8]);
mu_q = [0.5, 1.0]';
sigma_q = diag([1.2, 0.8]);

% case 2
% mu_p = [0.5, 1.0]';
% sigma_p = diag([1.2 0.8]);
% mu_q = [0.0, 1.5]';
% sigma_q = diag([0.9, 1.1]);

% Calculate KL divergence
kld = cal_KLD(mu_p, sigma_p, mu_q, sigma_q);

% Print the result
disp(['KL divergence: ', num2str(kl_loss)]);

% case 1 output: 0
% case 2 output: 0.44175

function kld = cal_KLD(mu_p, sigma_p, mu_q, sigma_q)
    eps = 1e-8;
    
    sigma_p = sigma_p .^ 2;
    sigma_q = sigma_q .^ 2;
    sigma_p_det = det(sigma_p);
    sigma_q_det = det(sigma_q);
    
    sigma_q_inv = inv(sigma_q);
    delta_u = (mu_q - mu_p);
    term1 = trace(sigma_q \ sigma_p);
    term2 = delta_u' * sigma_q_inv * delta_u;
    term3 = - length(mu_p);
    term4 = log(sigma_q_det + eps) - log(sigma_p_det + eps);
    kld = 0.5 * (term1 + term2 + term3 + term4);
    kld = max(kld, 0);
end

参考网址

[1] 两个高斯分布KL散度推导
[2] 多元高斯分布间的KL散度及其Pytorch实现
[3] 多变量高斯分布之间的KL散度
[4] 矩阵内积