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 to 1e-4	優秀（許容範囲内）	合格
1e-4 to 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]