본 포스팅에서는 이전 포스팅에서 설명한 적대적 공격과 방어 기법의 구현 코드를 분석합니다.

공격 기법 코드는 torchattacks 라이브러리(Harry24k/adversarial-attacks-pytorch)를 기반으로 합니다. torchattacks는 PyTorch 기반의 적대적 공격 구현체를 모아놓은 라이브러리로, 다양한 공격 기법을 간결한 인터페이스로 제공합니다. 방어 기법 코드는 MAIR 라이브러리(Harry24k/MAIR)를 기반으로 합니다. MAIR는 적대적 학습 기반의 방어 기법 구현체를 제공합니다.

분석할 대상 공격/방어 기법들은 다음 표와 같습니다.

이름 유형 논문 Distance
FGSM 공격 Explaining and Harnessing Adversarial Examples \(L_{\infty}\)
C&W 공격 Towards Evaluating the Robustness of Neural Networks \(L_2\)
PGD 공격 Towards Deep Learning Models Resistant to Adversarial Attacks \(L_{\infty}\)
PGD-L2 공격 Towards Deep Learning Models Resistant to Adversarial Attacks \(L_2\)
TPGD 공격 Theoretically Principled Trade-off between Robustness and Accuracy \(L_{\infty}\)
AT 방어 Towards Deep Learning Models Resistant to Adversarial Attacks -
TRADES 방어 Theoretically Principled Trade-off between Robustness and Accuracy \(L_{\infty}\)

공격 코드 리뷰

FGSM

FGSM(Fast Gradient Sign Method)는 가장 기초가 되는 적대적 예제 생성 방법입니다. FGSM 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[x_{adv} = x + \epsilon \cdot \text{sign}(\nabla_x J(\theta, x, y))\]

위 수식은 \(L_{\infty}\) 제약 조건 하에, 손실 함수의 gradient에 sign 함수를 취해 부호만 추출한 뒤, 각 픽셀을 \(\epsilon\)만큼 이동시키는 수식입니다. 구현 코드를 살펴보겠습니다.

def __init__(self, model, eps=8 / 255):
    super().__init__("FGSM", model)
    self.eps = eps
    self.supported_mode = ["default", "targeted"]

eps=8/255는 이미지가 [0, 1]로 정규화된 상태에서 사람의 눈으로 거의 식별 불가능한 수준의 perturbation 크기입니다. supported_mode에 targeted가 포함되어 있어 특정 클래스로의 오분류를 유도하는 targeted attack도 지원합니다.

def forward(self, images, labels):
    images = images.clone().detach().to(self.device)
    labels = labels.clone().detach().to(self.device)

    if self.targeted:
        target_labels = self.get_target_label(images, labels)

clone().detach()로 원본 텐서와 computation graph를 분리합니다. 이후 requires_grad = True를 설정할 것이기 때문에 원본에 영향을 주지 않도록 하기 위함입니다.

    loss = nn.CrossEntropyLoss()
    images.requires_grad = True
    outputs = self.get_logits(images)

    if self.targeted:
        cost = -loss(outputs, target_labels)
    else:
        cost = loss(outputs, labels)

일반적인 학습에서는 model.parameters()에 대한 gradient를 구하지만, 여기서는 이미지에 대한 gradient를 구합니다. 모델 파라미터는 고정한 채 입력 이미지를 어떻게 바꾸면 loss가 올라가는지를 계산하는 것입니다. targeted 모드에서 -loss를 사용하는 이유는 target label에 대한 loss를 최소화하는 방향으로 이미지를 조작하기 위해서입니다.

    grad = torch.autograd.grad(
        cost, images, retain_graph=False, create_graph=False
    )[0]

    adv_images = images + self.eps * grad.sign()
    adv_images = torch.clamp(adv_images, min=0, max=1).detach()

    return adv_images

retain_graph=False, create_graph=False는 gradient 계산 후 computation graph를 즉시 해제하는 옵션입니다. FGSM은 gradient를 한 번만 사용하기 때문에 메모리 절약을 위해 이렇게 설정합니다. grad.sign()으로 gradient의 부호만 취해 수식을 그대로 구현하고, torch.clamp로 결과 이미지가 유효한 픽셀 범위 [0, 1]을 벗어나지 않도록 합니다.

C&W

C&W(Carlini & Wagner)는 변수 변환을 통해 [0, 1] box 제약 조건을 우회하여 적대적 예제를 생성하는 공격 기법입니다. C&W 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[\min_{w} \left\| \frac{1}{2}(\tanh(w)+1) - x \right\|^2_2 + c \cdot f\left(\frac{1}{2}(\tanh(w)+1)\right)\] \[f(x') = \max\left(\max_{i \neq t} Z(x')_i - Z(x')_t, -\kappa\right)\]

위 수식은 \(L_2\) 제약 조건 하에서, 손실 함수 \(f\)를 통해 오분류를 유도하면서 동시에 perturbation의 \(L_2\) 크기를 최소화하는 방향으로 최적화하는 수식입니다. 구체적으로 코드를 분석해보겠습니다.

def __init__(self, model, c=1, kappa=0, steps=50, lr=0.01):
    super().__init__("CW", model)
    self.c = c
    self.kappa = kappa
    self.steps = steps
    self.lr = lr
    self.supported_mode = ["default", "targeted"]

c는 \(L_2\) 거리와 오분류 손실 사이의 균형을 조절하는 하이퍼파라미터입니다. c가 클수록 오분류에 더 집중하고, 작을수록 perturbation 크기를 줄이는 데 집중합니다. kappa는 confidence로, 오분류의 확신 정도를 조절합니다. 이 값이 클수록 적대적 예제에 대해서 강한 확신을 가지고 오분류를 하게 됩니다.

def forward(self, images, labels):

    images = images.clone().detach().to(self.device)
    labels = labels.clone().detach().to(self.device)

    if self.targeted:
        target_labels = self.get_target_label(images, labels)

    w = self.inverse_tanh_space(images).detach()
    w.requires_grad = True

    best_adv_images = images.clone().detach()
    best_L2 = 1e10 * torch.ones((len(images))).to(self.device)
    prev_cost = 1e10
    dim = len(images.shape)

    MSELoss = nn.MSELoss(reduction="none")
    Flatten = nn.Flatten()

    optimizer = optim.Adam([w], lr=self.lr)

C&W는 이미지를 직접 최적화하지 않고 w라는 변수를 도입합니다. 이미지가 [0, 1] 범위를 벗어나지 않는 동시에 제약 없는 최적화가 가능하도록 tanh 변환을 활용하는 것입니다. inverse_tanh_space로 원본 이미지를 w 공간으로 변환한 뒤 이를 최적화하고, tanh_space로 다시 이미지 공간으로 복원합니다. best_L2를 1e10으로 초기화해 이후 최솟값 갱신을 위한 기준값으로 사용합니다.

    for step in range(self.steps):
        adv_images = self.tanh_space(w)

        current_L2 = MSELoss(Flatten(adv_images), Flatten(images)).sum(dim=1)
        L2_loss = current_L2.sum()

        outputs = self.get_logits(adv_images)
        if self.targeted:
            f_loss = self.f(outputs, target_labels).sum()
        else:
            f_loss = self.f(outputs, labels).sum()

        cost = L2_loss + self.c * f_loss

        optimizer.zero_grad()
        cost.backward()
        optimizer.step()

매 스텝마다 w를 tanh_space로 변환해 적대적 이미지를 생성하고, \(L_2\) 거리 손실과 오분류 손실을 합산해 cost(총 손실)를 계산합니다. FGSM과 달리 torch.autograd.grad가 아닌 Adam optimizer로 w를 반복적으로 업데이트합니다.

        mask = condition * (best_L2 > current_L2.detach())
        best_L2 = mask * current_L2.detach() + (1 - mask) * best_L2

        mask = mask.view([-1] + [1] * (dim - 1))
        best_adv_images = mask * adv_images.detach() + (1 - mask) * best_adv_images

오분류에 성공하면서 동시에 현재 \(L_2\) 거리가 기존 최솟값보다 작을 때만 best_adv_images를 갱신합니다. mask를 이용해 조건을 만족하는 이미지만 선택적으로 업데이트하는 방식입니다.

        if step % max(self.steps // 10, 1) == 0:
            if cost.item() > prev_cost:
                return best_adv_images
            prev_cost = cost.item()

    return best_adv_images

전체 스텝의 10% 간격으로 loss 수렴 여부를 확인합니다. 이전 cost보다 현재 cost가 크면 조기 종료합니다.

def tanh_space(self, x):
    return 1 / 2 * (torch.tanh(x) + 1)

def inverse_tanh_space(self, x):
    return self.atanh(torch.clamp(x * 2 - 1, min=-1, max=1))

def f(self, outputs, labels):
    one_hot_labels = torch.eye(outputs.shape[1]).to(self.device)[labels]
    other = torch.max((1 - one_hot_labels) * outputs, dim=1)[0]
    real = torch.max(one_hot_labels * outputs, dim=1)[0]

    if self.targeted:
        return torch.clamp((other - real), min=-self.kappa)
    else:
        return torch.clamp((real - other), min=-self.kappa)

tanh_space는 w를 [0, 1] 범위로 변환하고, inverse_tanh_space는 그 역변환입니다. f 함수는 논문의 목적함수를 구현한 것으로, non-targeted의 경우 정답 클래스의 logit과 나머지 클래스 중 최댓값의 차이를 계산합니다. 이 값이 \(-\kappa\) 이하로 내려가지 않도록 clamp로 제한합니다.

PGD

PGD(Projected Gradient Descent)는 FGSM을 반복 최적화 형태로 확장한 공격 기법입니다. PGD 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[\max_{\delta \in \mathcal{S}} \mathcal{L}(\theta, x+\delta, y), \quad \mathcal{S}=\{\delta:\|\delta\|_\infty \le \epsilon\}\] \[x^{t+1} = \Pi_{x+\mathcal{S}}\Big(x^t + \alpha \cdot \mathrm{sign}\big(\nabla_x \mathcal{L}(\theta, x^t, y)\big)\Big)\]

여기서 \(\Pi_{x+\mathcal{S}}\)는 원본 \(x\) 주변 \(L_\infty\) 허용 집합으로의 projection입니다.

코드에서 핵심 과정은 다음 2단계입니다.

  1. 손실을 키우는 방향으로 한 스텝 이동
  2. 원본 기준 \(L_\infty\) 반경 \(\epsilon\) 내부로 projection

이를 구현한 코드는 다음과 같습니다.

def __init__(self, model, eps=8 / 255, alpha=2 / 255, steps=10, random_start=True):
    super().__init__("PGD", model)
    self.eps = eps
    self.alpha = alpha
    self.steps = steps
    self.random_start = random_start
    self.supported_mode = ["default", "targeted"]

def forward(self, images, labels):

    images = images.clone().detach().to(self.device)
    labels = labels.clone().detach().to(self.device)

    if self.targeted:
        target_labels = self.get_target_label(images, labels)

    if self.random_start:
        adv_images = adv_images + torch.empty_like(adv_images).uniform_(-self.eps, self.eps)
        adv_images = torch.clamp(adv_images, min=0, max=1).detach()

random start는 clean 이미지 주변 [-eps, eps] 구간에서 좌표별 균등 샘플링을 수행합니다. 이는 탐색 시작점을 다양화해 약한 local optimum에 머무는 문제를 줄이는 역할을 합니다.

    for _ in range(self.steps):
        adv_images.requires_grad = True
        outputs = self.get_logits(adv_images)

        if self.targeted:
            cost = -loss(outputs, target_labels)
        else:
            cost = loss(outputs, labels)

비표적 공격은 CE(outputs, labels)를 키워 정답 클래스에서 멀어지게 하고, 표적 공격은 -CE(outputs, target)를 키워 타겟 클래스 쪽으로 이동시킵니다.

        grad = torch.autograd.grad(cost, adv_images, retain_graph=False, create_graph=False)[0]
        adv_images = adv_images.detach() + self.alpha * grad.sign()
        delta = torch.clamp(adv_images - images, min=-self.eps, max=self.eps)
        adv_images = torch.clamp(images + delta, min=0, max=1).detach()

grad.sign()은 \(L_\infty\) 제약에서의 steepest-ascent 방향입니다. 업데이트 후 delta를 [-eps, eps]로 clip하여 “각 픽셀별 변화량 상한”이라는 \(L_\infty\) 제약을 보장하고, 마지막에 [0,1]로 clamp하여 유효 입력 범위를 유지합니다.

PGDL2

PGDL2는 PGD 구조를 유지하면서 제약 집합을 \(L_2\)로 바꾼 버전입니다. PGDL2 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[\max_{\delta \in \mathcal{S}} \mathcal{L}(\theta, x+\delta, y), \quad \mathcal{S}=\{\delta:\|\delta\|_2 \le \epsilon\}\] \[x^{t+1} = \Pi_{x+\mathcal{S}}\Big(x^t + \alpha \cdot \frac{\nabla_x \mathcal{L}(\theta, x^t, y)}{\|\nabla_x \mathcal{L}(\theta, x^t, y)\|_2+\eta}\Big)\]

즉, sign(grad) 대신 정규화된 gradient 방향을 사용하고, 좌표별 clip 대신 벡터 길이를 축소하는 방식으로 projection합니다.

def __init__(self, model, eps=1.0, alpha=0.2, steps=10, random_start=True, eps_for_division=1e-10):
    super().__init__("PGDL2", model)
    self.eps = eps
    self.alpha = alpha
    self.steps = steps
    self.random_start = random_start
    self.eps_for_division = eps_for_division
    self.supported_mode = ["default", "targeted"]

eps_for_division 수식의 \(\eta\)는 gradient norm이 매우 작을 때 분모 안정성을 확보하기 위한 상수입니다.

if self.random_start:
    delta = torch.empty_like(adv_images).normal_()
    d_flat = delta.view(adv_images.size(0), -1)
    n = d_flat.norm(p=2, dim=1).view(adv_images.size(0), 1, 1, 1)
    r = torch.zeros_like(n).uniform_(0, 1)
    delta *= r / n * self.eps
    adv_images = torch.clamp(adv_images + delta, min=0, max=1).detach()

random start에서는 각 샘플마다:

  • delta/n으로 방향벡터를 L2 norm 1로 정규화
  • r*eps를 곱해 최종 길이를 [0, eps] 범위로 설정

구(정확히는 고차원 L2 볼) 관점으로 보면, 중심은 원본 샘플 \(x\)이고 delta/n은 구의 중심에서 껍데기 방향으로 향하는 단위벡터(방향)입니다. 여기에 r*eps를 곱하면 “그 방향으로 얼마나 멀리 갈지”를 정하는 길이가 되어, 최종 시작점은 \(x + r\epsilon u\) 형태(\(u\): 단위벡터)가 됩니다. 즉 random start는 “방향은 무작위 + 길이는 0~\(\epsilon\) 사이”라는 방식으로 구 내부 점을 고릅니다.

    grad = torch.autograd.grad(cost, adv_images, retain_graph=False, create_graph=False)[0]
    grad_norms = torch.norm(grad.view(batch_size, -1), p=2, dim=1) + self.eps_for_division
    grad = grad / grad_norms.view(batch_size, 1, 1, 1)
    adv_images = adv_images.detach() + self.alpha * grad

업데이트 단계에서는 gradient를 샘플별로 L2 정규화한 뒤 alpha만큼 이동합니다. 즉 크기 정보보다 방향 정보를 중심으로 이동합니다.

    delta = adv_images - images
    delta_norms = torch.norm(delta.view(batch_size, -1), p=2, dim=1)
    factor = self.eps / delta_norms
    factor = torch.min(factor, torch.ones_like(delta_norms))
    delta = delta * factor.view(-1, 1, 1, 1)
    adv_images = torch.clamp(images + delta, min=0, max=1).detach()

이 부분이 L2 projection의 핵심입니다. “현재 교란 \(\delta\)의 길이 \((\|\delta\|_2)\)가 \(\epsilon\)을 넘는지”를 보고, 넘으면 같은 방향으로 길이만 줄입니다.

  • factor = eps / ||delta||_2: 현재 길이를 허용 반경에 맞추기 위한 스케일 비율 계산
  • factor = min(factor, 1): 이미 내부(||delta||_2 <= eps)면 factor=1로 유지, 외부면 factor<1로 축소
  • delta *= factor: 방향은 유지한 채 크기만 조정되어, 외부 점은 경계(||delta||_2 = eps)로 정확히 투영

이를 식으로 쓰면 다음과 같습니다.

\[\Pi_{\|\delta\|_2 \le \epsilon}(\delta)=\delta\cdot \min(1,\epsilon/\|\delta\|_2)\]

그리고 최종 단계에서 adv_images = torch.clamp(images + delta, min=0, max=1).detach()를 적용해, L2 반경 제약과는 별도로 입력 자체가 유효 픽셀 범위 [0,1]를 벗어나지 않도록 보정합니다.

TPGD

TPGD(TRADES PGD)는 TRADES 방어 기법 논문에서 제안된 내부 최적화(inner maximization)를 위한 적대적 예제 생성 기법입니다. TPGD 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[x_{adv} = x + \alpha \cdot \text{sign}(\nabla_x KL(f(x) \| f(x_{adv})))\]

위 수식은 정답 레이블을 사용하는 대신, 원본 이미지에 대한 예측 확률 분포와 적대적 이미지에 대한 예측 확률 분포 사이의 쿨백-라이블러 발산(KL Divergence)을 최대화하는 방향으로 픽셀을 이동시키는 수식입니다. 구현 코드를 살펴보겠습니다.

def __init__(self, model, eps=8 / 255, alpha=2 / 255, steps=10):
    super().__init__("TPGD", model)
    self.eps = eps
    self.alpha = alpha
    self.steps = steps
    self.supported_mode = ["default"]

TPGD는 레이블 자체를 활용하여 오분류를 유도하는 것이 아니라, 원본 이미지의 예측 분포와의 차이를 벌리는 데 목적이 있으므로 ‘targeted’ 모드를 지원하지 않고 ‘default’ 모드만 지원합니다.

def forward(self, images, labels=None):
    images = images.clone().detach().to(self.device)
    logit_ori = self.get_logits(images).detach()

    adv_images = images + 0.001 * torch.randn_like(images)
    adv_images = torch.clamp(adv_images, min=0, max=1).detach()

    loss = nn.KLDivLoss(reduction="sum")

원본 이미지의 예측 값(logit_ori)을 미리 계산하고 detach()로 고정하여 기준점으로 삼습니다. 이후 최적화를 시작하기 전, 초기 위치를 무작위로 약간 이동(random start)시켜 지역 최적해(local optima)에 빠지는 것을 방지합니다. 이때 일반적인 PGD가 균등 분포(uniform)를 쓰는 것과 달리, 계수 0.001과 함께 가우시안 노이즈(torch.randn_like)를 더해주는 디테일이 들어있습니다. 손실 함수로는 reduction="sum" 옵션이 적용된 KLDivLoss를 준비합니다.

for _ in range(self.steps):
    adv_images.requires_grad = True
    logit_adv = self.get_logits(adv_images)

    # Calculate loss
    cost = loss(F.log_softmax(logit_adv, dim=1), F.softmax(logit_ori, dim=1))

매 스텝마다 적대적 이미지에 대한 모델의 출력(logit_adv)을 계산합니다. PyTorch의 KLDivLoss는 첫 번째 인자로 log-probability를, 두 번째 인자로 probability를 요구하므로 적대적 출력에는 log_softmax를, 원본 출력에는 softmax를 취해줍니다. 이 KL 손실값을 계산하여 최종 cost로 정의합니다.

    # Update adversarial images
    grad = torch.autograd.grad(
        cost, adv_images, retain_graph=False, create_graph=False
    )[0]

    adv_images = adv_images.detach() + self.alpha * grad.sign()
    delta = torch.clamp(adv_images - images, min=-self.eps, max=self.eps)
    adv_images = torch.clamp(images + delta, min=0, max=1).detach()

return adv_images

계산된 cost에 대해 adv_images의 gradient를 계산한 후, 부호(grad.sign())를 추출하여 alpha만큼 이미지를 업데이트합니다. 이후 원본 이미지와의 차이(delta)가 eps 범위를 넘지 않게 제약하고, 최종 픽셀 값도 [0, 1]의 유효 범위를 유지하도록 잘라냅니다(clipping).

방어 코드 리뷰

AT

AT(Adversarial Training)는 가장 직관적이고 널리 쓰이는 적대적 방어 기법으로, 모델 학습 과정 자체에 적대적 예제를 포함시키는 방법입니다. AT 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[\min_{\theta} \mathbb{E}_{(x, y)} \left[ \max_{\|x'-x\|_{\infty} \leq \epsilon} L(f_\theta(x'), y) \right]\]

위 수식은 내부 최대화(inner maximization)를 통해 오차를 가장 크게 만드는 적대적 예제 \(x'\)를 찾고, 외부 최소화(outer minimization)를 통해 해당 적대적 예제에 대한 손실 함수를 최소화하는 방향으로 모델 파라미터 \(\theta\)를 번갈아 학습시키는 Min-Max 최적화 수식입니다. 구체적으로 코드를 분석해보겠습니다.

def __init__(self, rmodel, eps, alpha, steps, random_start=True):
    super().__init__(rmodel)
    self.atk = PGD(rmodel, eps, alpha, steps, random_start)

AT는 모델을 강건하게 학습시키기 위해 내부적으로 강력한 1차 미분 기반 공격인 PGD(Projected Gradient Descent)를 사용합니다. 생성자에서 rmodel(방어 모델)과 공격 하이퍼파라미터(eps, alpha, steps), 그리고 초기 노이즈 추가 여부를 결정하는 random_start를 받아 PGD 공격 모듈을 초기화합니다.

def calculate_cost(self, train_data, reduction="mean"):
    # ... (데이터 디바이스 할당 생략) ...
    adv_images = self.atk(images, labels)
    logits_adv = self.rmodel(adv_images)

미니배치로 들어온 원본 images와 labels를 앞서 정의한 PGD 공격 모듈에 통과시켜 즉석에서 적대적 예제(adv_images)를 생성합니다. 그리고 방어 모델(rmodel)에 이 적대적 예제를 입력으로 넣어 예측값(logits_adv)을 산출합니다. 학습 시 원본 이미지가 아닌 적대적 예제로 피드포워드(feed-forward)를 수행하는 것이 핵심입니다.

    cost = nn.CrossEntropyLoss(reduction="none")(logits_adv, labels)
    self.add_record_item("CALoss", cost.mean().item())

    return cost.mean() if reduction == "mean" else cost

생성된 적대적 예제들의 예측값과 실제 정답 레이블 사이의 Cross-Entropy 손실(cost)을 계산합니다. 이때 reduction="none"을 사용하여 배치 내 각 샘플의 손실값을 개별적으로 구합니다. 이후 MAIR 프레임워크의 로깅 시스템을 통해 CALoss(Cross-Entropy Adversarial Loss)라는 이름으로 평균 손실값을 기록하여 에포크별 강건성 손실 추이를 추적합니다. 마지막으로 함수 호출부에서 요구하는 reduction 방식(일반적으로 “mean”)에 맞춰 최종 손실값을 반환하고 역전파에 사용하게 됩니다.

TRADES

TRADES(TRadeoff-inspired Adversarial DEfense via Surrogate-loss minimization)는 자연 정확도와 적대적 강건성 사이의 trade-off를 이론적으로 다루는 방어 기법입니다. TRADES 수식을 간단하게 리뷰해보면 다음과 같습니다.

\[\min_f \mathbb{E}\left[ L(f(x), y) + \beta \cdot \max_{\|x'-x\|_{\infty} \leq \epsilon} KL(f(x) \| f(x')) \right]\]

위 수식은 clean 이미지에 대한 Cross-Entropy 손실과, clean 예측과 적대적 예측 사이의 KL divergence를 \(\beta\)로 조절하며 동시에 최소화하는 수식입니다. \(\beta\)가 클수록 강건성에, 작을수록 자연 정확도에 집중합니다.

def __init__(self, rmodel, eps, alpha, steps, beta):
    super().__init__(rmodel)
    self.atk = TPGD(rmodel, eps, alpha, steps)
    self.beta = beta

TRADES는 inner maximization에 TPGD를 사용합니다. TPGD는 정답 레이블 기준이 아닌 clean 예측과 적대적 예측 사이의 KL divergence를 최대화하는 방향으로 적대적 이미지를 생성합니다. beta는 수식의 \(\beta\)에 해당하는 trade-off 파라미터입니다.

def calculate_cost(self, train_data, reduction="mean"):

    images, labels = train_data
    images = images.to(self.device)
    labels = labels.to(self.device)

    logits_clean = self.rmodel(images)
    loss_ce = nn.CrossEntropyLoss(reduction=reduction)(logits_clean, labels)

    adv_images = self.atk(images)
    logits_adv = self.rmodel(adv_images)

clean 이미지에 대한 logit을 먼저 계산해 Cross-Entropy 손실을 구하고, TPGD로 적대적 이미지를 생성한 뒤 해당 logit을 따로 계산합니다. 두 logit을 분리해서 계산하는 이유는 이후 KL divergence 계산에서 각각 다른 역할을 하기 때문입니다.

    probs_clean = F.softmax(logits_clean, dim=1)
    log_probs_adv = F.log_softmax(logits_adv, dim=1)
    loss_kl = nn.KLDivLoss(reduction="none")(log_probs_adv, probs_clean).sum(dim=1)

    cost = loss_ce + self.beta * loss_kl

KL divergence를 계산할 때 clean 쪽은 softmax, 적대적 쪽은 log_softmax를 적용합니다. PyTorch의 KLDivLoss는 입력으로 log-probability를 요구하기 때문입니다. 최종 cost는 Cross-Entropy 손실과 KL divergence 손실을 beta로 가중합산한 값으로, 수식을 그대로 구현한 것입니다.

    self.add_record_item("Loss", cost.mean().item())
    self.add_record_item("CELoss", loss_ce.mean().item())
    self.add_record_item("KLLoss", loss_kl.mean().item())

학습 과정에서 전체 손실, Cross-Entropy 손실, KL divergence 손실을 각각 기록합니다. 세 값을 분리해서 추적하면 beta조절에 따라 두 손실이 어떻게 변하는지 모니터링할 수 있습니다.

이상으로 적대적 예제 생성의 기본이 되는 공격 기법들과, 방어 기법들에 대한 코드 분석 포스팅을 마칩니다.