PatchCore를 이용한 MIMII Dataset의 Anomaly Detection - Greedy Coreset Selection (260816)

Greedy Coreset Selection이란?

배경: 왜 coreset이 필요한가

PatchCore는 정상 데이터의 모든 패치 특징을 memory bank에 저장해두고, 테스트 시 새 패치가 이 memory bank 안의 가장 가까운 이웃과 얼마나 떨어져 있는지로 이상 여부를 판단합니다. 그런데 정상 학습 데이터의 모든 패치를 다 저장하면:

• memory bank가 너무 커져서 저장 공간과 최근접 이웃 탐색(nearest neighbor search) 비용이 감당이 안 됨
• 실제로는 인접한 패치들끼리 매우 유사한 정보를 담고 있어 중복이 많음(비효율)

그래서 전체 패치 집합 중 가장 "대표성" 있는 일부만 뽑아서 memory bank로 쓰자는 게 coreset selection의 목적입니다. 여기서 핵심은 "대표성"을 어떻게 정의하고 어떻게 뽑느냐입니다.

Greedy Coreset (Greedy k-center / Minimax Facility Location)

PatchCore 논문(Roth et al., 2022)은 이를 k-center problem으로 정식화합니다.

목표: 전체 패치 집합 $\mathcal{P}$에서 부분집합 $\mathcal{C} \subset \mathcal{P}$ (memory bank)를 골라, $\mathcal{P}$의 모든 점이 $\mathcal{C}$의 어떤 점과도 "너무 멀지 않도록" 만든다.

수식으로는 다음을 최소화하는 $\mathcal{C}$를 찾는 것:

$\max_{p \in \mathcal{P}} \min_{c \in \mathcal{C}} \| p - c \|_2$

즉, "전체 데이터 중 가장 멀리 떨어진(=가장 대표되지 못한) 점까지의 거리"를 최소화하는 것 — 이걸 minimax facility location이라고 부릅니다. 이 문제는 NP-hard라서 정확히 풀 수 없고, 대신 greedy 근사 알고리즘을 씁니다.

알고리즘 (제가 구현한 _greedy_sample과 동일한 절차):

1. 임의의 점 하나를 첫 center로 선택
2. 모든 점에 대해 "현재 선택된 center들 중 가장 가까운 것까지의 거리" (min_distances)를 계산
3. min_distances가 가장 큰 점(=현재 center들로부터 가장 소외된, 대표되지 못한 점)을 새 center로 추가
4. 새로 추가된 center를 기준으로 min_distances를 갱신 (기존 값과 새 거리 중 더 작은 값으로 업데이트)
5. 목표 개수(target)에 도달할 때까지 2~4 반복

이렇게 하면 매 스텝마다 "지금 memory bank가 가장 못 커버하고 있는 영역"을 채워나가는 방식이라, 결과적으로 데이터 공간 전체를 고르게 커버하는 점들이 뽑힙니다.

Random Projection을 쓰는 이유: 실제 패치 특징은 384차원 등으로 고차원인데, 이 차원 그대로 매 반복마다 거리 계산을 하면 비용이 큽니다. Johnson-Lindenstrauss 정리에 따르면, 랜덤 선형 사영으로 차원을 줄여도 점들 사이의 거리 관계가 근사적으로 보존되므로, 128차원 정도로 줄여서 거리 계산 비용을 낮추면서도 결과 품질은 거의 유지할 수 있습니다. (PatchCore 논문도 실제로 이 기법을 사용합니다.)

처음 방법(랜덤 샘플링)과의 차이점

왜 이게 실제로 중요한가 (지난번 실험으로 확인한 부분)

지난번 합성 데이터 실험에서, 전체의 2.5%만 차지하는 희소 영역(정상 데이터 중에서도 드문 패턴)에 대해:

• 랜덤 샘플링: 희소 영역에서 뽑힌 점이 0개
• Greedy coreset: 희소 영역에서 3개를 뽑아냄

이 차이가 실제 이상탐지 성능에 미치는 영향은 이렇습니다. PatchCore는 "memory bank에서 가장 가까운 점까지의 거리"로 이상 점수를 매기는데, 만약 정상 데이터 중 드문 패턴(예: 특정 운전 조건에서만 나오는 정상 소리)이 memory bank에서 완전히 빠지면, 테스트 시 그 패턴이 다시 나왔을 때 "memory bank와 멀다 = 이상하다"라고 잘못 판단(false positive) 할 위험이 커집니다. 즉 랜덤 샘플링은 운이 나쁘면 정상 패턴의 일부를 통째로 놓칠 수 있는 반면, greedy coreset은 그런 누락 위험을 구조적으로 줄여줍니다.

트레이드오프

다만 공짜는 아닙니다:

• 속도: greedy는 반복적으로 거리를 계산해야 해서 랜덤보다 확실히 느립니다 (CORESET_RATIO를 올릴수록 이 차이는 더 벌어짐 — 지난번 논의한 test.py 속도 문제와 유사한 구조로, target이 커질수록 train.py의 coreset 단계도 느려집니다).
• 이상치(outlier) 민감성: greedy는 "가장 먼 점"을 계속 뽑기 때문에, 학습 데이터에 잘못 섞여 들어간 노이즈나 라벨링 오류가 있으면 그것부터 뽑힐 수 있습니다. (반면 랜덤은 이런 위험이 상대적으로 낮음 — 노이즈가 뽑힐 확률이 그냥 데이터 비율만큼임)

정리하면, 랜덤 샘플링은 빠르지만 운에 좌우되는 대표성을, greedy coreset은 느리지만 의도적으로 전체 공간을 커버하는 대표성을 제공합니다. PatchCore 논문이 greedy를 기본으로 채택한 이유도 이상탐지에서는 "흔치 않은 정상 패턴을 놓치지 않는 것"이 성능에 직결되기 때문입니다.

https://github.com/wontaeleeterry/PyTorch_Study/tree/main/mimii_patchcore_greedy_coreset_selection_260817

patchcore.py

import numpy as np
import torch
import torch.nn.functional as F

from tqdm import tqdm

import faiss


# ============================================================
# Feature map -> patch embedding
# ============================================================

def feature_map_to_patches(
    feature,
):
    """
    [B, C, H, W]
        ->
    [B, H*W, C]
    """

    b, c, h, w = feature.shape

    feature = feature.permute(
        0,
        2,
        3,
        1,
    )

    feature = feature.reshape(
        b,
        h * w,
        c,
    )

    return feature


def aggregate_features(
    features,
):
    """
    Combine layer2 and layer3.

    layer2:
        [B, C2, H2, W2]

    layer3:
        [B, C3, H3, W3]

    Returns:
        [B, H2*W2, C2+C3]
    """

    layer2 = features["layer2"]

    layer3 = features["layer3"]

    # --------------------------------------------------------
    # layer2
    # --------------------------------------------------------

    layer2_patches = (
        feature_map_to_patches(
            layer2
        )
    )

    b, n, c2 = (
        layer2_patches.shape
    )

    h = layer2.shape[2]
    w = layer2.shape[3]

    # --------------------------------------------------------
    # layer3 -> spatially resize
    # --------------------------------------------------------

    layer3 = F.interpolate(
        layer3,
        size=(h, w),
        mode="bilinear",
        align_corners=False,
    )

    layer3_patches = (
        feature_map_to_patches(
            layer3
        )
    )

    # --------------------------------------------------------
    # concatenate
    # --------------------------------------------------------

    features = torch.cat(
        [
            layer2_patches,
            layer3_patches,
        ],
        dim=-1,
    )

    return features


# ============================================================
# Coreset
# ============================================================

class CoresetSampler:
    """
    PatchCore의 approximate greedy coreset selection.

    Greedy k-center (minimax facility location):
    현재까지 선택된 center 집합과의 거리(최근접 center까지의 거리)가
    가장 먼 점을 반복적으로 새 center로 추가한다.

    연산량을 줄이기 위해:
      1) Johnson-Lindenstrauss random projection으로 차원을 축소해서
         거리 계산 비용을 낮추고,
      2) 매 반복마다 전체 pairwise distance를 다시 계산하지 않고
         min_distance(각 점 -> 가장 가까운 center)만 갱신한다.
         (O(target * n) matrix 연산, target 회의 O(n) 갱신)

    method="random"으로 두면 기존의 단순 랜덤 샘플링도 그대로 사용할 수 있다
    (속도 비교/디버깅용).
    """

    def __init__(
        self,
        ratio=0.01,
        random_seed=42,
        method="greedy",
        projection_dim=128,
        device=None,
    ):

        self.ratio = ratio

        self.random_seed = (
            random_seed
        )

        self.method = method

        self.projection_dim = (
            projection_dim
        )

        self.device = (
            device
            if device is not None
            else torch.device("cpu")
        )

    # ------------------------------------------------------------
    # Random projection (Johnson-Lindenstrauss)
    # ------------------------------------------------------------

    def _random_projection(
        self,
        features,
    ):
        """
        [N, D] -> [N, projection_dim]

        projection_dim이 None이거나 원래 차원보다 크면
        projection을 적용하지 않는다.
        """

        dim = features.shape[1]

        if (
            self.projection_dim is None
            or self.projection_dim >= dim
        ):

            return features

        generator = (
            torch.Generator()
            .manual_seed(self.random_seed)
        )

        projection_matrix = (
            torch.randn(
                dim,
                self.projection_dim,
                generator=generator,
            )
            / (self.projection_dim ** 0.5)
        )

        projection_matrix = (
            projection_matrix
            .to(features.device)
        )

        return features @ projection_matrix

    # ------------------------------------------------------------
    # Greedy k-center coreset
    # ------------------------------------------------------------

    def _greedy_sample(
        self,
        features_np,
        target,
    ):

        n = len(features_np)

        features_t = (
            torch.from_numpy(features_np)
            .to(self.device)
        )

        proj = self._random_projection(
            features_t
        )

        rng = np.random.default_rng(
            self.random_seed
        )

        # ----------------------------------------------------
        # 첫 center는 무작위로 선택
        # ----------------------------------------------------

        first_idx = int(
            rng.integers(0, n)
        )

        selected_indices = [
            first_idx
        ]

        min_distances = (
            torch.cdist(
                proj,
                proj[first_idx : first_idx + 1],
            )
            .squeeze(1)
        )

        # 이미 선택된 점이 다시 뽑히지 않도록 마스킹
        min_distances[first_idx] = -1.0

        # ----------------------------------------------------
        # 반복적으로 min_distance가 최대인 점을 center로 추가
        # ----------------------------------------------------

        for _ in tqdm(
            range(1, target),

            desc="Greedy coreset sampling",
        ):

            next_idx = int(
                torch.argmax(
                    min_distances
                ).item()
            )

            selected_indices.append(
                next_idx
            )

            new_distances = (
                torch.cdist(
                    proj,
                    proj[next_idx : next_idx + 1],
                )
                .squeeze(1)
            )

            min_distances = (
                torch.minimum(
                    min_distances,
                    new_distances,
                )
            )

            min_distances[next_idx] = -1.0

        return np.array(
            selected_indices
        )

    # ------------------------------------------------------------
    # Random sampling (fallback / 비교용)
    # ------------------------------------------------------------

    def _random_sample(
        self,
        n,
        target,
    ):

        rng = np.random.default_rng(
            self.random_seed
        )

        return rng.choice(
            n,
            size=target,
            replace=False,
        )

    # ------------------------------------------------------------
    # Entry point
    # ------------------------------------------------------------

    def sample(
        self,
        features,
    ):

        features = np.asarray(
            features,
            dtype=np.float32,
        )

        n = len(features)

        target = max(
            1,
            int(n * self.ratio)
        )

        if target >= n:

            return features

        if self.method == "greedy":

            indices = self._greedy_sample(
                features,
                target,
            )

        elif self.method == "random":

            indices = self._random_sample(
                n,
                target,
            )

        else:

            raise ValueError(
                f"Unknown coreset method: "
                f"{self.method}"
            )

        return features[indices]



class PatchCoreMemory:

    def __init__(self, k=1):

        self.k = k
        self.index = None


    def fit(self, memory):

        memory = np.asarray(
            memory,
            dtype=np.float32
        )

        dimension = memory.shape[1]

        self.index = faiss.IndexFlatL2(
            dimension
        )

        self.index.add(
            memory
        )


    def predict(self, features):

        features = np.asarray(
            features,
            dtype=np.float32
        )

        distances, indices = (
            self.index.search(
                features,
                self.k
            )
        )

        # ----------------------------------------------------
        # k nearest neighbor distance
        # ----------------------------------------------------

        patch_scores = distances[:, 0]

        # ----------------------------------------------------
        # Image-level anomaly score
        # ----------------------------------------------------

        score = float(
            np.max(
                patch_scores
            )
        )

        return (
            score,
            patch_scores
        )