3D高斯散射中的反向传播链式求梯度：原理详解与新手避坑指南

1次阅读

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

3D 高斯散射（3DGS）是近年来神经渲染领域的重要技术，它通过大量可学习的高斯分布来表示 3D 场景，能够实现高质量的实时渲染。但在实际应用中，梯度计算往往成为性能瓶颈：

内存消耗大：传统实现中每个高斯参数都需要独立存储中间梯度，导致显存占用随高斯数量线性增长
数值不稳定：透明度乘积和指数计算容易引发梯度爆炸 / 消失

3DGS 的渲染结果可表示为：
$$I = \sum_{i=1}^N \alpha_i \cdot c_i \cdot \prod_{j=1}^{i-1}(1-\alpha_j)$$
其中：
– $\alpha_i$ = $\text{sigmoid}(\text{MLP}(\mu_i, \Sigma_i))$
– $c_i$ = $\exp(-\frac{1}{2}(x-\mu_i)^T\Sigma_i^{-1}(x-\mu_i))$

对均值 $\mu$ 的梯度传播：
$$\frac{\partial I}{\partial \mu} = \sum_{i=1}^N \frac{\partial I}{\partial c_i} \cdot \frac{\partial c_i}{\partial \mu_i}$$
其中 $\frac{\partial c_i}{\partial \mu_i} = c_i \cdot \Sigma_i^{-1}(x-\mu_i)$

class GaussianRenderer(nn.Module):
    def __init__(self, N):
        super().__init__()
        self.mu = nn.Parameter(torch.randn(N, 3))  # 均值
        self.L = nn.Parameter(torch.randn(N, 3, 3))  # 协方差矩阵的 Cholesky 分解
        self.alpha = nn.Parameter(torch.rand(N))  # 透明度

    def forward(self, rays):
        # 计算各高斯对光线的贡献
        diff = rays[:, None] - self.mu[None]  # [R, N, 3]
        cov = self.L @ self.L.transpose(1,2)  # [N, 3, 3]
        exp_term = -0.5 * (diff @ cov.inverse() * diff).sum(-1)  # [R, N]
        c = exp(exp_term)  # [R, N]

        # 累积透明度
        alpha = torch.sigmoid(self.alpha)  # [N]
        weights = alpha[None] * c * torch.cumprod(1 - alpha[None] + 1e-10, dim=1)  # [R, N]
        return weights.sum(-1)  # [R]

使用 torch.no_grad() 包装不需要梯度的中间计算
对大型张量操作启用 inplace=True 选项

# 对数空间计算替代原始乘积
def safe_alpha_prod(alpha):
    log_alpha = torch.log(alpha + 1e-10)
    log_1malpha = torch.log1p(-alpha + 1e-10)
    return torch.exp(log_alpha + torch.cumsum(log_1malpha, dim=0))

梯度维度不匹配 ：使用assert grad.shape == param.shape 进行校验
数值异常检测：添加torch.autograd.set_detect_anomaly(True)

def check_grad(func, inputs, eps=1e-3):
    # 数值梯度验证
    analytic_grad = torch.autograd.grad(func, inputs)[0]
    num_grad = torch.zeros_like(inputs)
    for i in range(inputs.numel()):
        inputs_plus = inputs.clone()
        inputs_plus.view(-1)[i] += eps
        inputs_minus = inputs.clone()
        inputs_minus.view(-1)[i] -= eps
        num_grad.view(-1)[i] = (func(inputs_plus) - func(inputs_minus))/(2*eps)
    return torch.allclose(analytic_grad, num_grad, rtol=1e-2)