Robust Representation Consistency Model via Contrastive Denoising
핵심 요약
이 논문은 딥러닝 모델을 적대적 공격(adversarial perturbations)으로부터 보호하기 위한 새로운 방법을 제안합니다.
기존의 확산 모델(diffusion model) 기반 방어 기법들은 높은 정확도를 보였지만, 추론 시간이 매우 길다는 큰 단점이 있었습니다.
이 논문은 확산 모델의 노이즈 제거 과정을 생성(generative) 작업이 아닌 판별(discriminative) 작업으로 재해석하여, 성능은 뛰어나면서도 추론 속도를 획기적으로 개선한 모델(rRCM)을 개발했습니다.
그 결과, 기존 확산 모델 기반 방법보다 평균 85배 빠른 속도로 더 높은 인증 정확도(certified accuracy)를 달성했습니다.
그림 2는 이 논문의 핵심적인 두 단계를 시각적으로 보여주는 매우 중요한 그림입니다.
(a)와 (b)로 나누어 자세히 설명해 드리겠습니다.
그림 2(a): 사전 학습 (Pre-training) - 어떻게 모델을 똑똑하게 만드는가?
이 부분은 모델의 '기초 체력'을 기르는 단계입니다. 최종 목표인 분류를 하기 전에, 모델이 노이즈가 낀 이미지로부터 강건하고 일관된 특징(representation)을 추출하는 법을 배우는 과정입니다. 그림에서 볼 수 있듯이, 두 가지 손실 함수(Loss)를 동시에 사용합니다.
- 일관성 손실 (Consistency Loss):
- 입력: 깨끗한 이미지(x_0)에 노이즈(ε)를 더해, 시간적으로 아주 약간 차이 나는 두 개의 노이즈 이미지(x_t_n, x_t_n-1)를 만듭니다. 이 둘은 같은 복원 경로 상에 있는 '이웃'과 같습니다.
- 목표: 모델(Encoder)이 이 두 '이웃' 이미지를 보고 거의 동일한 특징을 추출하도록 훈련시킵니다.
- 의미: 이 과정을 통해 모델은 확산 모델의 노이즈 구조를 학습하고, 노이즈의 미세한 차이를 극복하여 본질적인 이미지를 파악하는 능력을 기릅니다. 이것이 이 논문의 핵심적인 아이디어입니다.
- 대조 손실 (Contrastive Loss):
- 입력: 깨끗한 이미지(x_0)에 회전, 자르기 등 일반적인 데이터 증강(augmentation)을 적용해 두 개의 약간 다른 버전(z_1, z_2)을 만듭니다.
- 목표: 모델(Encoder + Projector)이 이 두 버전의 이미지는 '같은 것'으로, 다른 이미지로부터 온 것들은 '다른 것'으로 인식하도록 학습합니다. 즉, 같은 이미지의 다른 뷰는 가깝게, 다른 이미지는 멀게 만듭니다.
- 의미: 이 과정은 모델이 이미지의 의미적, 시각적 특징 자체를 더 잘 이해하도록 돕습니다. 일관성 손실이 '노이즈'에 대한 이해를 돕는다면, 대조 손실은 '이미지 콘텐츠' 자체에 대한 이해를 돕습니다.
사전 학습 후: 이 복잡한 학습이 끝나면, 대조 손실 계산에만 사용되던 보조 장치인 '프로젝터(Projector)'는 버려집니다. 우리는 오직 잘 훈련된 **'인코더(Encoder)'**만 다음 단계로 가져갑니다.
그림 2(b): 순전파 (Forward Pass) - 실제로 어떻게 사용하는가?
이 부분은 사전 학습으로 똑똑해진 인코더를 실제 분류 문제에 사용하는 방법을 보여줍니다. 이 단계는 매우 간단하고 빠릅니다.
- 입력: 노이즈가 낀 이미지 (x + σε) 하나가 주어집니다.
- 과정:
- 이 이미지가 사전 학습된 **인코더(Encoder)**를 통과하여 강건한 특징 벡터를 추출합니다.
- 이 특징 벡터가 간단한 선형 헤드(Linear Head), 즉 분류기를 통과하여 최종 클래스 레이블(예: '고양이')을 예측합니다.
- 의미: 이것이 바로 이 모델의 효율성이 극대화되는 지점입니다. 기존 방법들처럼 여러 단계에 걸쳐 노이즈를 제거할 필요 없이, 단 한 번의 순전파(single forward pass)만으로 노이즈가 낀 이미지로부터 바로 정답을 예측합니다. 모델이 사전 학습 단계에서 노이즈를 '암묵적으로' 제거하고 핵심 특징을 파악하는 법을 이미 배웠기 때문입니다.
요약하자면, (a)는 다양한 종류의 노이즈와 데이터 변형 속에서도 핵심을 꿰뚫어 보는 강력한 '눈'(Encoder)을 만들기 위한 복잡하고 정교한 훈련 과정이며, (b)는 그렇게 잘 훈련된 '눈'을 사용하여 실제 문제를 빠르고 정확하게 해결하는 실전 활용법을 보여줍니다.
문제점: 기존 방법의 한계
딥러닝 모델은 약간의 노이즈(적대적 공격)에도 쉽게 성능이 저하되는 취약점을 가집니다.
이를 해결하기 위해 randomized smoothing라는 인증된 방어 기법이 사용됩니다.
최근에는 확산 모델을 이용해 입력 이미지의 노이즈를 먼저 제거하고 분류하는 방식이 높은 성능을 보여주었지만, 두 가지 큰 문제가 있었습니다.
- 높은 계산 비용: 노이즈 제거 과정에서 여러 단계의 샘플링을 거치기 때문에 추론 시간이 매우 깁니다. 예를 들어 ImageNet 데이터셋의 이미지 한 장을 인증하는 데 수십 분이 걸리기도 합니다.
- 번거로운 구조: 노이즈를 제거하는 확산 모델과 분류를 수행하는 분류기, 이렇게 두 개의 독립된 모델이 필요하여 관리 오버헤드가 컸습니다.
제안된 해결책: rRCM의 핵심 아이디어
저자들은 이 문제를 해결하기 위해 확산 모델의 노이즈 제거 과정을 새로운 관점에서 접근했습니다.
확산 모델은 깨끗한 이미지에 점진적으로 노이즈를 추가하고(forward process), 반대로 노이즈로부터 깨끗한 이미지를 복원(reverse process)합니다.
이때 노이즈가 낀 이미지와 조금 덜 낀 이미지는 같은 복원 경로(trajectory) 상에 놓여있다고 볼 수 있습니다.
rRCM의 핵심 아이디어는 **"같은 복원 경로 상에 있는 이미지들은 잠재 공간(latent space)에서 서로 유사한 표현(representation)을 가져야 한다"**는 것입니다.
이는 이미지를 생성하는 대신, 주어진 두 이미지가 같은 경로에서 왔는지를 판별하는 문제로 바꾼 것입니다.
이 아이디어를 구현하기 위해 2단계 학습 방식을 사용합니다.
1단계: 사전 학습 (Pre-training)
- 대조적 노이즈 제거 모델이 노이즈 낀 이미지들의 관계를 배우도록 사전 학습을 진행합니다. 여기서는 '대조 학습(contrastive learning)' 기법을 활용합니다.
- 긍정 쌍 (Positive Pair):
같은 깨끗한 이미지로부터 생성된, 시간적으로 인접한 두 노이즈 이미지를 하나의 쌍으로 묶습니다. 이 둘은 잠재 공간에서 가깝도록 학습됩니다. - 부정 쌍 (Negative Pair):
서로 다른 이미지로부터 생성된 노이즈 이미지들은 부정 쌍으로 간주되어, 잠재 공간에서 멀어지도록 학습됩니다. 이 과정을 통해 모델은 노이즈의 구조를 이해하고, 노이즈가 낀 이미지로부터 일관된 특징을 추출하는 능력을 학습하게 됩니다.
2단계: 미세 조정 (Fine-tuning)
사전 학습된 모델에 분류를 위한 선형 레이어(linear head)를 추가한 뒤, 노이즈가 낀 이미지와 실제 레이블을 이용해 전체 모델을 미세 조정합니다.
이 단계에서 모델은 노이즈 제거와 분류를 하나의 과정으로 통합하여 한 번의 순전파(single forward pass)만으로 빠르고 강건한 예측을 수행할 수 있게 됩니다.
주요 실험 결과 및 기여
- 성능과 효율성 동시 달성:
- ImageNet 데이터셋에서 기존 확산 모델 기반 방법들보다 평균 5.3%, 특히 더 큰 노이즈(perturbation radii)에서는 최대 11.6% 더 높은 인증 정확도를 달성했습니다.
- 동시에 추론 비용은 평균 85배나 절감하여 기존의 고전적인 방법들과 비슷한 수준의 속도를 보였습니다.
- 뛰어난 확장성:
- 모델의 크기나 학습 배치 사이즈를 늘릴수록 성능이 지속적으로 향상되는 것을 보여주었습니다.
- 이는 더 많은 계산 자원을 투입하면 더 높은 성능을 기대할 수 있다는 것을 의미합니다.
- 새로운 방향 제시:
- 확산 모델의 구조적인 노이즈 스케줄을 분류 모델의 강건성(robustness) 향상에 직접 활용한 첫 번째 연구입니다.
- 이는 노이즈가 낀 샘플과 깨끗한 샘플 간의 관계를 활용하여 모델의 강건성을 높이는 새로운 연구 방향을 제시합니다.
- 결론적으로, 이 논문은 기존의 느리고 비효율적이던 확산 모델 기반 방어 기법의 패러다임을 바꾸어, 빠르고 정확하며 확장 가능한 새로운 해결책인 rRCM을 성공적으로 제시했습니다.
이 모델은 생성(generative) 모델의 원리를 판별(discriminative) 작업에 매우 효과적으로 적용한, 일종의 하이브리드적인 접근 방식을 사용합니다.
둘 다 된다고 말할 수 있지만, 역할과 목적이 다릅니다.
- 생성(Generative) 원리 활용:
- 이 모델의 근간은 확산 모델(Diffusion Model)입니다.
- 확산 모델은 본질적으로 노이즈로부터 데이터를 '생성'하는 모델입니다.
- rRCM은 이 생성 과정, 즉 노이즈가 점차 제거되며 깨끗한 이미지로 복원되는 구조적인 경로(trajectory)의 원리를 그대로 가져와 활용합니다. 깨끗한 이미지와 노이즈 낀 이미지가 어떻게 연결되는지에 대한 물리적(수학적) 원리를 제공하는 셈이죠.
- 판별(Discriminative) 작업으로의 변환:
- 하지만 rRCM은 실제로 이미지를 생성하지는 않습니다. 대신, 생성 모델의 원리를 판별적인 학습 목표를 위해 사용합니다.
- 사전 학습 단계: 모델은 "두 노이즈 이미지가 같은 원본 이미지로부터 왔는가?"를 **구분(판별)**하도록 학습됩니다. 이것이 바로 논문에서 "생성 작업을 판별 작업으로 재구성했다"고 말하는 핵심 부분입니다.
- 최종 목표: 모델의 최종적인 임무는 이미지를 특정 클래스(예: 고양이, 개)로 **분류(판별)**하는 것입니다.
정리
이미지를 '만드는' 생성 모델의 설계도(원리)를 가져와서, 이미지를 '구분하는' 판별 모델을 훈련시키는 데 사용한 것입니다.
그렇기 때문에 rRCM은 생성 모델의 강력한 이론적 기반 위에서, 실제로는 매우 효율적인 판별 모델(분류기)로 작동할 수 있게 된 것입니다.
이 접근 방식 덕분에 노이즈 제거와 분류를 한 번에 처리하면서도 높은 성능을 달성할 수 있었습니다.
Projector는 무슨 역할을 하나요?
결론부터 말하면, '사전 학습'을 위한 임시 보조 장치입니다.
- 역할: 인코더(Encoder)가 추출한 특징(representation)을 또 다른 공간으로 한 번 더 변환(projection)하는 역할을 합니다. 이 변환된 공간에서 대조 손실(Contrastive Loss)을 계산합니다.
- 왜 필요한가?: 만약 프로젝터 없이 인코더의 출력에서 바로 대조 손실을 계산하면, 인코더는 분류와 같은 다른 유용한 정보는 다 버리고 오직 '대조 학습' 작업에만 과도하게 최적화될 위험이 있습니다. 프로젝터는 이런 정보 손실을 막아주는 완충 장치 역할을 합니다. 인코더는 풍부하고 일반적인 특징을 배우는 데 집중하고, 프로젝터는 대조 학습이라는 특정 임무만 처리하는 것이죠.
- 가장 중요한 점: 사전 학습이 끝나면 이 프로젝터는 과감히 버려집니다. 우리는 잘 훈련된 인코더만 가져가서 최종 분류 모델을 만드는 데 사용합니다.
2. stop-grad는 무엇인가요?
'기울기(Gradient) 흐름을 멈춰라'라는 신호입니다.
- 역할: 딥러닝 모델은 '역전파(backpropagation)' 과정에서 오차(loss)로부터 발생한 기울기를 거꾸로 전파시키며 가중치를 업데이트(학습)합니다. stop-grad는 특정 경로로 기울기가 흐르는 것을 인위적으로 차단하는 장치입니다.
- 왜 사용하는가?: 그림 2(a)의 '일관성 손실' 부분을 보면, x_t_n-1이 들어가는 아래쪽 경로에 stop-grad가 있습니다. 이는 모델이 꼼수를 부리는 것을 방지하기 위함입니다. 만약 양쪽 경로 모두에서 기울기가 흐른다면, 모델은 두 입력에 대해 그냥 똑같은 출력만 내보내도록 학습하여 손실을 0으로 만드는 쉬운 길을 택할 수 있습니다. 하지만 이렇게 되면 아무것도 배우지 못하죠.
- '선생님과 학생' 비유: stop-grad가 적용된 아래쪽 경로를 '선생님(정답)'으로, 위쪽 경로를 '학생'으로 생각할 수 있습니다. 학생은 선생님의 출력을 따라 하도록 배우지만, 학생의 오답이 선생님의 지식에 영향을 주지는 않습니다. 이렇게 한쪽을 고정된 목표(target)로 만들어 모델이 의미 있는 방향으로 학습하도록 유도하는 것입니다.
3. ViT(Vision Transformer)를 쓰는 건가요?
네, 맞습니다. 이 논문에서는 Vision Transformer(ViT)를 모델의 기본 구조(backbone)로 사용합니다. ViT는 이미지를 여러 패치로 나누어 처리하는 트랜스포머 구조로, 최근 컴퓨터 비전 분야에서 매우 높은 성능을 보이고 있어 SOTA(State-of-the-art) 연구에서 널리 사용됩니다.
4. 노이즈 낀 이미지에서 바로 클래스 레이블 예측? 이게 Forward? 왜 Backward가 아닌가요?
이 질문은 딥러닝의 두 가지 핵심 단계를 정확히 짚어주셨습니다.
- "노이즈 낀 이미지에서 바로 클래스 레이블 예측인가요?" 네, 정확합니다! 이것이 바로 이 모델의 가장 큰 장점입니다. 사전 학습을 통해 인코더는 노이즈를 보고도 그 안에 숨겨진 진짜 이미지의 특징을 뽑아내는 능력을 이미 갖추었습니다. 그래서 여러 번의 노이즈 제거 과정 없이, 단 한 번의 계산만으로 바로 최종 예측을 할 수 있습니다.
- "이게 Forward Pass인가요? 왜 Backward Pass가 아닌가요?" 네, 이것은 Forward Pass가 맞습니다.
- Forward Pass (순전파, 추론): 입력 데이터(이미지)가 신경망의 첫 계층부터 마지막 계층까지 순서대로 통과하며 결과를 출력하는 과정입니다. 즉, 모델을 **사용(Inference)**하여 예측값을 얻는 과정입니다. 그림 2(b)는 바로 이 '사용' 단계를 보여줍니다.
- Backward Pass (역전파, 학습): 이 과정은 학습(Training) 단계에서만 일어납니다. Forward Pass로 얻은 예측값이 실제 정답과 얼마나 다른지(오차, loss) 계산한 후, 이 오차를 마지막 계층부터 첫 계층까지 거꾸로 전파시키면서 각 계층의 가중치를 업데이트(학습)하는 과정입니다.
결론적으로, 모델을 훈련시킬 때는 Forward Pass와 Backward Pass가 모두 필요하고, 잘 훈련된 모델을 실제로 사용할 때는 Forward Pass만 수행하여 결과를 얻습니다. 그림 2(b)는 '실제 사용' 상황을 묘사하고 있기 때문에 Forward Pass라고 부르는 것이 맞습니다.
reference
https://velog.io/@yelim421/The-information-bottleneck-method