model-equivariance-auditor

模型等变性审计

即使使用了等变库，实现错误也可能破坏等变性。声称具有等变性但实际并不具备的模型，训练效果会很差，预测也会不一致。及早发现这些错误能节省大量调试时间。

工作流

复制此清单并跟踪进度：

Equivariance Audit Progress:
- [ ] Step 1: Gather model and symmetry specification
- [ ] Step 2: Run numerical equivariance tests
- [ ] Step 3: Test individual layers
- [ ] Step 4: Check gradient equivariance
- [ ] Step 5: Identify and diagnose failures
- [ ] Step 6: Document audit results

步骤 1：收集模型和对称性规范

收集：已实现的模型、预期的对称群、每个输出应具备不变性还是等变性，以及输入和输出空间的变换函数。查看设计阶段的架构规范。测试前与用户澄清所有歧义。

步骤 2：运行数值等变性测试

使用测试实现执行端到端等变性测试。对于不变性：验证 ||f(T(x)) - f(x)|| < ε。对于等变性：验证 ||f(T(x)) - T'(f(x))|| < ε。使用多个随机输入和变换。记录误差统计数据。阈值请参阅误差解读。可直接使用的测试代码请参阅测试代码模板。

步骤 3：测试各个层

如果端到端测试失败，则逐层测试以隔离问题。对于每一层：冻结其他层，单独测试该层的等变性。这样可以找出破坏等变性的层。使用逐层测试流程。尤其要仔细检查非线性函数、归一化和自定义操作。

步骤 4：检查梯度等变性

验证梯度也遵循等变性（这对训练很重要）。在 x 和 T(x) 处计算梯度。检查梯度是否按预期变换。梯度错误可能导致训练“丧失”等变性。请参阅梯度测试。

步骤 5：识别并诊断故障

如果测试失败，请使用常见故障模式进行诊断。检查：非等变非线性函数、批归一化问题、输出变换错误、数值精度问题、自定义层中的实现错误。提供具体的修复建议。分步排查请参阅调试指南。

步骤 6：记录审计结果

使用输出模板创建审计报告。包括：每项测试的通过/失败结果、误差大小、发现的问题和建议。区分：精确等变性（数值精度范围内）、近似等变性（误差可接受）和等变性遭到破坏（需要修复）。详细的审计方法请参阅方法详情。此输出的质量标准见质量评分标准。

测试实现

端到端等变性测试

import torch

def test_model_equivariance(model, x, input_transform, output_transform,
                            n_tests=100, tol=1e-5):
    """
    Test if model is equivariant: f(T(x)) ≈ T'(f(x))

    Args:
        model: The neural network to test
        x: Sample input tensor
        input_transform: Function that transforms input
        output_transform: Function that transforms output
        n_tests: Number of random transformations to test
        tol: Error tolerance

    Returns:
        dict with test results
    """
    model.eval()
    errors = []

    with torch.no_grad():
        for _ in range(n_tests):
            # Generate random transformation
            T = sample_random_transform()

            # Method 1: Transform input, then apply model
            x_transformed = input_transform(x, T)
            y1 = model(x_transformed)

            # Method 2: Apply model, then transform output
            y = model(x)
            y2 = output_transform(y, T)

            # Compute error
            error = torch.norm(y1 - y2).item()
            relative_error = error / (torch.norm(y2).item() + 1e-8)
            errors.append({
                'absolute': error,
                'relative': relative_error
            })

    return {
        'mean_absolute': np.mean([e['absolute'] for e in errors]),
        'max_absolute': np.max([e['absolute'] for e in errors]),
        'mean_relative': np.mean([e['relative'] for e in errors]),
        'max_relative': np.max([e['relative'] for e in errors]),
        'pass': all(e['relative'] < tol for e in errors)
    }

不变性测试（更简单的情况）

def test_model_invariance(model, x, transform, n_tests=100, tol=1e-5):
    """Test if model output is invariant to transformations."""
    model.eval()
    errors = []

    with torch.no_grad():
        y_original = model(x)

        for _ in range(n_tests):
            T = sample_random_transform()
            x_transformed = transform(x, T)
            y_transformed = model(x_transformed)

            error = torch.norm(y_transformed - y_original).item()
            errors.append(error)

    return {
        'mean_error': np.mean(errors),
        'max_error': np.max(errors),
        'pass': max(errors) < tol
    }

逐层测试

流程

def test_layer_equivariance(layer, x, input_transform, output_transform):
    """Test a single layer for equivariance."""
    layer.eval()

    with torch.no_grad():
        T = sample_random_transform()

        # Transform then layer
        y1 = layer(input_transform(x, T))

        # Layer then transform
        y2 = output_transform(layer(x), T)

        error = torch.norm(y1 - y2).item()

    return {
        'layer': layer.__class__.__name__,
        'error': error,
        'pass': error < tolerance
    }

def audit_all_layers(model, x, transforms):
    """Test each layer individually."""
    results = []

    for name, layer in model.named_modules():
        if is_testable_layer(layer):
            result = test_layer_equivariance(layer, x, *transforms)
            result['name'] = name
            results.append(result)

    return results

各层需要测试的内容

层类型	检查内容
卷积	卷积核等变性
非线性函数	应保持等变性
归一化	通常会破坏等变性
池化	聚合是否正确
线性层	权重共享模式
注意力	置换等变性

梯度测试

为什么要测试梯度？

前向传播可能具备等变性，而反向传播却不具备。这会导致：

• 训练不稳定
• 模型“丧失”等变性
• 优化不一致

梯度等变性测试

def test_gradient_equivariance(model, x, loss_fn, transform, tol=1e-4):
    """Test if gradients respect equivariance."""
    model.train()

    # Gradients at original input
    x1 = x.clone().requires_grad_(True)
    y1 = model(x1)
    loss1 = loss_fn(y1)
    loss1.backward()
    grad1 = x1.grad.clone()

    # Gradients at transformed input
    model.zero_grad()
    T = sample_random_transform()
    x2 = transform(x.clone(), T).requires_grad_(True)
    y2 = model(x2)
    loss2 = loss_fn(y2)
    loss2.backward()
    grad2 = x2.grad.clone()

    # Transform grad1 and compare to grad2
    grad1_transformed = transform_gradient(grad1, T)
    error = torch.norm(grad2 - grad1_transformed).item()

    return {'error': error, 'pass': error < tol}

误差解读

误差阈值

误差级别	解读	操作
< 1e-6	完美（float32 精度）	通过
1e-6 到 1e-4	极佳（可接受）	通过
1e-4 到 1e-2	近似等变性	调查
> 1e-2	等变性遭到破坏	需要修复

相对误差与绝对误差

• 绝对误差：原始差值的大小
• 相对误差：按输出大小归一化后的误差

输出大小变化时使用相对误差。与数值精度比较时使用绝对误差。

常见失败模式

1. 非等变非线性

症状：经过非线性层后误差增大 原因：在等变特征上使用 ReLU、sigmoid 解决方法：使用门控非线性、基于范数的非线性，或仅对不变特征使用非线性

2. 批归一化破坏等变性

症状：误差随批次组成而变化 原因：BN 会针对不同方向计算不同的统计量 解决方法：使用 LayerNorm、GroupNorm 或等变批归一化

3. 输出变换不正确

症状：即使使用恒等变换，测试也会失败 原因：output_transform 与模型输出类型不匹配 解决方法：确认输出变换与层输出表示相匹配

4. 数值精度问题

症状：各处都有较小但非零的误差 原因：浮点数累积、插值 解决方法：测试时使用 float64，并接受较小的容差

5. 自定义层错误

症状：误差仅出现在特定层 原因：自定义等变层的实现有误 解决方法：根据等变约束检查层的实现

6. 填充/边界效应

症状：边缘附近的误差更大 原因：填充不符合对称性 解决方法：使用循环填充，或显式处理边界

输出模板

MODEL EQUIVARIANCE AUDIT REPORT
===============================

Model: [Model name/description]
Intended Symmetry: [Group]
Symmetry Type: [Invariant/Equivariant]

END-TO-END TESTS:
-----------------
Test samples: [N]
Transformations tested: [M]

Invariance/Equivariance Error:
- Mean absolute: [value]
- Max absolute: [value]
- Mean relative: [value]
- Max relative: [value]
- RESULT: [PASS/FAIL]

LAYER-WISE ANALYSIS:
--------------------
[For each layer]
- Layer: [name]
- Error: [value]
- Result: [PASS/FAIL]

GRADIENT TEST:
--------------
- Gradient equivariance error: [value]
- RESULT: [PASS/FAIL]

IDENTIFIED ISSUES:
------------------
1. [Issue description]
   - Location: [layer/component]
   - Severity: [High/Medium/Low]
   - Recommended fix: [description]

OVERALL VERDICT: [PASS/FAIL/NEEDS_ATTENTION]

Recommendations:
- [List of actions needed]