3D高斯前向传播及反向传播原理细节：从数学推导到代码实现

1次阅读

共计 2761 个字符，预计需要花费 7 分钟才能阅读完成。

3D 高斯分布是概率论和统计学中常用的一种连续概率分布，广泛应用于机器学习、计算机视觉等领域。其概率密度函数（PDF）定义为：

$$
\mathcal{N}(\mathbf{x} | \boldsymbol{\mu}, \boldsymbol{\Sigma}) = \frac{1}{(2\pi)^{3/2} |\boldsymbol{\Sigma}|^{1/2}} \exp\left(-\frac{1}{2} (\mathbf{x} – \boldsymbol{\mu})^T \boldsymbol{\Sigma}^{-1} (\mathbf{x} – \boldsymbol{\mu}) \right)
$$

其中，
– $\mathbf{x} \in \mathbb{R}^3$ 是输入向量
– $\boldsymbol{\mu} \in \mathbb{R}^3$ 是均值向量
– $\boldsymbol{\Sigma} \in \mathbb{R}^{3 \times 3}$ 是协方差矩阵，必须是对称正定的

在前向传播中，计算 3D 高斯值的性能瓶颈通常出现在以下几个方面：

协方差矩阵的逆运算 ：计算 $\boldsymbol{\Sigma}^{-1}$ 的时间复杂度为 $O(n^3)$，在 3D 情况下为 $O(27)$。
行列式计算 ：计算 $|\boldsymbol{\Sigma}|$ 的复杂度同样为 $O(n^3)$。
指数运算 ：exp 函数计算较为耗时。

为了提高计算效率，可以采用以下优化策略：

Cholesky 分解 ：将协方差矩阵分解为 $\boldsymbol{\Sigma} = \mathbf{L} \mathbf{L}^T$，其中 $\mathbf{L}$ 是下三角矩阵。这样可以高效计算行列式和逆矩阵。
对数空间计算 ：在计算 PDF 时，先计算对数 PDF，最后再取指数，避免数值下溢。

反向传播需要计算损失函数对高斯分布参数的梯度。假设损失函数为 $\mathcal{L}$，我们需要计算 $\frac{\partial \mathcal{L}}{\partial \boldsymbol{\mu}}$ 和 $\frac{\partial \mathcal{L}}{\partial \boldsymbol{\Sigma}}$。

均值梯度 ：

$$
\frac{\partial \mathcal{L}}{\partial \boldsymbol{\mu}} = \frac{\partial \mathcal{L}}{\partial \mathcal{N}} \cdot \frac{\partial \mathcal{N}}{\partial \boldsymbol{\mu}} = \frac{\partial \mathcal{L}}{\partial \mathcal{N}} \cdot \mathcal{N} \cdot \boldsymbol{\Sigma}^{-1} (\mathbf{x} – \boldsymbol{\mu})
$$

协方差矩阵梯度 ：

$$
\frac{\partial \mathcal{L}}{\partial \boldsymbol{\Sigma}} = \frac{\partial \mathcal{L}}{\partial \mathcal{N}} \cdot \frac{\partial \mathcal{N}}{\partial \boldsymbol{\Sigma}} = \frac{\partial \mathcal{L}}{\partial \mathcal{N}} \cdot \mathcal{N} \cdot \left[-\frac{1}{2} \boldsymbol{\Sigma}^{-1} + \frac{1}{2} \boldsymbol{\Sigma}^{-1} (\mathbf{x} – \boldsymbol{\mu}) (\mathbf{x} – \boldsymbol{\mu})^T \boldsymbol{\Sigma}^{-1} \right]
$$

import numpy as np

def gaussian_3d_pdf(x, mu, sigma):
    """
    计算 3D 高斯分布的概率密度函数

    参数:
    x: 输入向量, shape (3,)
    mu: 均值向量, shape (3,)
    sigma: 协方差矩阵, shape (3, 3)

    返回:
    概率密度值
    """
    # 确保输入形状正确
    assert x.shape == (3,)
    assert mu.shape == (3,)
    assert sigma.shape == (3, 3)

    # 计算差值向量
    diff = x - mu

    # Cholesky 分解计算行列式和逆矩阵
    L = np.linalg.cholesky(sigma)
    log_det = 2 * np.sum(np.log(np.diag(L)))
    inv_sigma = np.linalg.inv(sigma)

    # 计算指数部分
    exponent = -0.5 * diff.T @ inv_sigma @ diff

    # 计算归一化常数
    normalization = 1.0 / ((2 * np.pi) ** 1.5 * np.sqrt(np.exp(log_det)))

    # 计算最终概率密度
    pdf = normalization * np.exp(exponent)

    return pdf

# 示例用法
x = np.array([1.0, 2.0, 3.0])
mu = np.array([0.0, 0.0, 0.0])
sigma = np.array([[1.0, 0.5, 0.3],
                  [0.5, 1.0, 0.2],
                  [0.3, 0.2, 1.0]])

pdf = gaussian_3d_pdf(x, mu, sigma)
print(f"Probability density: {pdf}")