从零推导MMD：用NumPy和PyTorch构建图像质量评估器

在计算机视觉领域，判断两组图像是否来自同一分布是个基础但关键的问题。想象你训练了一个生成模型，如何量化它生成的图片与真实照片的差距？传统方法如PSNR、SSIM只能捕捉像素级差异，而**最大均值差异（MMD）**提供了一种更本质的分布相似性度量方式。今天我们将抛开框架封装，从数学原理出发，用纯手工实现揭开MMD的神秘面纱。

1. MMD的核心思想与数学原理

1.1 分布差异的直观理解

假设你面前有两堆沙子，如何判断它们是否来自同一个沙坑？最直接的方法是抓取样本进行对比：

• 传统方法：比较沙粒的平均大小、颜色（类似PSNR/SSIM）
• MMD方法：将沙粒放入特殊显微镜（核函数）观察微观结构，比较所有可能视角下的平均特征

数学上，MMD通过将数据映射到再生核希尔伯特空间（RKHS），在该空间中计算两个分布样本均值的距离。关键公式如下：

MMD^2(P,Q) = \mathbb{E}_{x,x'}[k(x,x')] + \mathbb{E}_{y,y'}[k(y,y')] - 2\mathbb{E}_{x,y}[k(x,y)]

提示：当且仅当P=Q时，MMD=0。高斯核是最常用的核函数，其带宽参数σ控制着特征空间的尺度敏感性。

1.2 核函数的选择艺术

不同的核函数就像不同的显微镜镜头：

核类型	公式	适用场景
高斯核	exp(-
拉普拉斯核	exp(-
线性核	xᵀy	计算简单，表达能力弱

在图像领域，我们通常采用多尺度高斯核组合来捕捉不同层次的特征差异。

2. NumPy基础实现

2.1 核矩阵计算

我们先实现核心的高斯核函数：

import numpy as np

def gaussian_kernel(X, Y, sigma=1.0):
    """
    计算样本集X和Y之间的高斯核矩阵
    X: (m,d) numpy数组
    Y: (n,d) numpy数组
    sigma: 高斯核带宽
    返回: (m,n)核矩阵
    """
    XX = np.sum(X**2, axis=1)[:, np.newaxis]
    YY = np.sum(Y**2, axis=1)[np.newaxis, :]
    distances = XX + YY - 2 * np.dot(X, Y.T)
    return np.exp(-distances / (2 * sigma**2))

2.2 完整MMD计算流程

结合核函数实现MMD：

def mmd_naive(X, Y, kernel='gaussian', sigma=1.0):
    """基础MMD实现"""
    if kernel == 'gaussian':
        K_XX = gaussian_kernel(X, X, sigma)
        K_YY = gaussian_kernel(Y, Y, sigma)
        K_XY = gaussian_kernel(X, Y, sigma)
    else:
        raise ValueError("Unsupported kernel")
    
    m = X.shape[0]
    n = Y.shape[0]
    
    # 无偏估计版本
    term1 = (K_XX.sum() - np.trace(K_XX)) / (m*(m-1))
    term2 = (K_YY.sum() - np.trace(K_YY)) / (n*(n-1))
    term3 = K_XY.sum() * 2 / (m*n)
    
    return np.sqrt(term1 + term2 - term3)

注意：这里使用了无偏估计版本，避免对角线元素（自相似度）对结果的影响。实际应用中可能需要考虑计算效率的优化。

3. PyTorch优化实现

3.1 批处理与GPU加速

用PyTorch重构实现，支持自动微分和GPU加速：

import torch

def mmd_rbf(X, Y, kernel_mul=2.0, kernel_num=5):
    """多尺度高斯核MMD实现"""
    batch_size = X.size(0)
    total = torch.cat([X, Y], dim=0)
    
    # 计算成对距离矩阵
    XX = torch.sum(X**2, dim=1, keepdim=True)
    YY = torch.sum(Y**2, dim=1, keepdim=True)
    distances = XX + YY.T - 2 * torch.mm(X, Y.T)
    
    # 自适应带宽选择
    median_distance = torch.median(distances.detach())
    sigma_list = [median_distance * (kernel_mul**i) for i in range(-kernel_num//2, kernel_num//2+1)]
    
    # 多核组合
    kernel_val = sum(torch.exp(-distances / (2 * sigma**2)) for sigma in sigma_list)
    
    # MMD计算
    K_XX = kernel_val[:batch_size, :batch_size]
    K_YY = kernel_val[batch_size:, batch_size:]
    K_XY = kernel_val[:batch_size, batch_size:]
    
    return torch.mean(K_XX) + torch.mean(K_YY) - 2 * torch.mean(K_XY)

3.2 实际图像评估案例

让我们在CIFAR-10数据上测试：

from torchvision import datasets, transforms

# 数据准备
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
])
real_data = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
fake_data = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 提取特征
def get_features(dataloader, sample_size=1000):
    features = []
    for i, (images, _) in enumerate(dataloader):
        features.append(images.flatten(start_dim=1))
        if len(features) * images.size(0) >= sample_size:
            break
    return torch.cat(features)[:sample_size]

real_features = get_features(torch.utils.data.DataLoader(real_data, batch_size=100))
fake_features = get_features(torch.utils.data.DataLoader(fake_data, batch_size=100))

# 计算MMD
distance = mmd_rbf(real_features, fake_features)
print(f"MMD距离: {distance.item():.4f}")

4. 工业级实现对比

4.1 Inception特征空间的重要性

原始像素空间的MMD可能无法捕捉语义级差异。工业界常用做法：

1. 使用Inception-v3的中间层作为特征提取器
2. 在特征空间计算MMD（即著名的FID指标基础）
3. 采用多尺度金字塔结构增强鲁棒性

from torchvision.models import inception_v3

class InceptionMMD:
    def __init__(self, device='cuda'):
        self.model = inception_v3(pretrained=True)
        self.model.fc = torch.nn.Identity()  # 移除全连接层
        self.model.to(device)
        self.model.eval()
        
    def get_features(self, images):
        with torch.no_grad():
            return self.model(images)
            
    def __call__(self, X, Y):
        feat_X = self.get_features(X)
        feat_Y = self.get_features(Y)
        return mmd_rbf(feat_X, feat_Y)

4.2 实际应用建议

在真实项目中：

• 对于256x256以上图像，建议使用多层级特征提取
• 批量大小至少64以上以获得稳定估计
• 可结合自适应带宽选择策略
• 考虑使用线性时间估计版本处理大数据集

def linear_mmd(feat_X, feat_Y):
    """线性复杂度MMD估计"""
    kernel = lambda x,y: torch.exp(-torch.norm(x-y, p=2)**2 / feat_X.size(1))
    phi_X = torch.mean(torch.stack([kernel(x, x) for x in feat_X]), dim=0)
    phi_Y = torch.mean(torch.stack([kernel(y, y) for y in feat_Y]), dim=0)
    return torch.norm(phi_X - phi_Y, p=2)

5. 常见问题与调优技巧

5.1 超参数选择策略

• 带宽σ：通常取数据 pairwise 距离的中位数
• 核数量：3-5个不同尺度通常足够
• 样本量：每类至少1000个样本可获得稳定估计

5.2 数值稳定性处理

当数据尺度差异大时：

def stable_mmd(X, Y):
    # 数据标准化
    X = (X - X.mean()) / X.std()
    Y = (Y - Y.mean()) / Y.std()
    
    # 添加小常数防止数值溢出
    epsilon = 1e-6
    distances = torch.cdist(X, Y) + epsilon
    
    # 使用log-sum-exp技巧
    max_val = torch.max(distances)
    kernel_val = torch.exp(-(distances - max_val)**2 / 2)
    return kernel_val.mean()

5.3 与其他指标对比

在图像生成评估中：

指标	计算复杂度	感知相关性	是否需要参考图像
MMD	O(n²)	中	是
FID	O(n²)	高	是
IS	O(n)	中	否
PSNR	O(n)	低	是

实际项目中，我通常会同时计算MMD和FID作为互补指标。当发现MMD降低但FID升高时，往往意味着模型出现了模式崩溃。