[논문 리뷰] R.A.C.E : Robust Adversarial Concept Erasure for Secure Text-to-Image Diffusion Model

R.A.C.E.: Robust Adversarial Concept Erasure for Secure Text-to-Image Diffusion Model

이번에는 컨택드렸던 자대 교수님의 논문을 가져와봤다. 내가 주로 관심가지고 연구하는 분야인 Robustness / Security of AI에 부합하는 논문이라 평소에 컨택할 때 읽었던 논문들보다 조금은 더 자세히 읽었던 것 같다. 지난번 랩실 컨택 때에는 그냥 아무 생각도 안하고 지원했던 것 같은데....ㅎㅎ

Abstract

In the evolving landscape of text-to-image (T2I) diffusion models, the remarkable capability to generate high-quality images from textual descriptions faces challenges with the potential misuse of reproducing sensitive content. To address this critical issue, we introduce Robust Adversarial Concept Erase (RACE), a novel approach designed to mitigate these risks by enhancing the robustness of concept erasure method for T2I models. RACE utilizes a sophisticated adversarial training framework to identify and mitigate adversarial text embeddings, significantly reducing the Attack Success Rate (ASR). Impressively, RACE achieves a 30 percentage point reduction in ASR for the “nudity” concept against the leading white-box attack method. Our extensive evaluations demonstrate RACE’s effectiveness in defending against both white-box and black-box attacks, marking a significant advancement in protecting T2I diffusion models from generating inappropriate or misleading imagery. This work underlines the essential need for proactive defense measures in adapting to the rapidly advancing field of adversarial challenges. Our code is publicly available: https://github.com/chkimmmmm/ R.A.C.E.

논문의 주 Contribution은 아래와 같다.

• Adversarial Training을 통한 T2I Diffusion Model에서의 Concept Erasure
• 또한, Adversarial Training을 통한 Adversarial Robustness에 대한 보장

Introduction

이 논문에서 진행하고 싶은 'Concept Erasure'라는 개념의 한장 요약이라고 볼 수 있을 것이다. 이미지에서 특정 '스타일', 혹은 'Concept'라고 불릴 수 있는 것을 제거하는 것이라고 보면 되는데, 이는 논문에서 나와 있듯 저작권이 있는 콘텐츠나 부적절한 콘텐츠가 제작되는 것ㅇ르 막기 위해 연구된다고 보면 될 것이다. 기존 Concept Erasure 기법들의 경우에는 Adversarial Attack에 상당히 취약한 단점을 보였는데, 본 연구에서는 TTI Diffusion Model에서의 Adversarial Attack에 Robust한 Concept Erasure Method를 제안한다. 한 줄로 정리하자면, '기존 Concept Erasure에 Adversarial Training을 접목한' 방식이라고 볼 수 있을 것이다.

Algorithm

여기서 보면 |> Perform targeted Attack이라고 되어 있는 곳에 PGD(Projected Gradient Descent) 공격이 가해지는 것을 알 수 있다. 이렇게 학습 과정 내에서 Adversarial Example을 사용하여서 학습을 시키는 것을 Adversarial Training이라고 하고, 이는 Adversarial Attack에 대한 Robustness를 강화하기 위해 주로 사용되는 방식이라고 볼 수 있을 것이다. 

Erasing Stable Diffusion 모델 구조.

이게 R.A.C.E의 기저 논문으로 사용되는 Erasing Stable Diffusion의 모델 구조이다. 여기서 Concept Erase라는 개념이 처음 도입되었는데, 이 모델에 추가적으로 Adversarial Training을 수행한다고 생각하면 될 것이다.

단일 시간 단계 공격(Single-Timestep Adversarial Attack)

Diffusion Model의 경우에는 여러 단계에 거쳐 noising / denoising을 거치며 원하는 concept / image를 만들어나가는 과정이라고 보아도 무방하다(이 말이 이해되지 않는다면, diffusion 모델을 다시 보고 오길 바란다. 그렇기 때문에 adversarial training을 진행할 때에 모든 timestep에 거쳐서 adversarial example을 생성하고 training을 진행한다면 computational cost가 상당히 증가하게 될 것이다. 그렇기 때문에, 본 논문에서는 Stable Diffusion의 Loss Function을 이용하여 Adversarial Attack을 한 timestep에서만 진행하고, 추후 이를 이용하여 모든 timestep을 학습시키는 방식으로 computational cost를 낮출 수 있었다.

Experiments

일단 ASR(Attack Success Rate) 관점에서 분석한 RACE 방식의 성능을 먼저 알아보도록 하자. 기존 방식들은 Adversarial Attack에 대해서 ASR이 매우 높은 모습을 보이는 경향이 있었는데, 이를 특정 부분에서 확실히 감소시키는 것을 볼 수 있었다. 하지만 어쩔 수 없듯이, 'Illegal'과 같은 다소 vague한 concept들에 대해서는 제대로 잡아내지 못하는 모습을 보였다. 이는 뭐 사람의 영역이지 아직까지는 모델의 영역이 아닌 것 아닐까... 싶다.

추가적으로 ASR 감소 성능 이외에도, 생성되는 이미지의 퀄리티 또한 중요한 평가 지표가 되어야 할 것이다. 이를 측정하기 위해 위의 Table 우측을 보게 되면 CLIP-Score, FID(Frechet Inception Distance) 2가지를 추가적으로 측정하였다. FID의 경우에는 생성된 이미지 - 실제 이미지 사이의 유사도를 측정하는 지표이고, CLIP Score는 생성된 이미지와 그 이미지에 대한 텍스트 설명 간의 일치도를 측정하는 지표라고 보면 된다. 이 2가지 지표에서 모두 기존 방식들에 비해 ASR은 유의미하게 감소시켰지만, 이미지의 퀄리티는 크게 떨어지지 않은 것을 알 수 있다. Adversarial Training은 주로 이미지 생성 품질을 약화시킨다는 우려가 존재하고 실제로 기존 모델들은 이에 부합하는 행보를 보였는데, 이러한 기존 한계점을 충분히 보완하는 모델 구조였다고 생각한다. 단일 단계에서만 공격을 진행하여 이와 같은 결과를 얻을 수 있었다는 생각이 든다.

Conclusion

결론적으로 얘기해 보자면, 이 논문은 T2I Diffusion Model들의 Robustness, Security를 향상시키는 데에 크게 기여했다고 할 수 있다. Adversarial Training을 통한 강건한 Concept Erasure를 제안함과 동시에 단일 시간 단계 공격을 이용하여 비용의 감소와 성능 유지 또한 챙길 수 있었다. 이와 같은 연구들이 계속 진행되어야 할 것이라고 생각하며, 나 또한 이러한 모델과 방법론들에 대한 연구를 꾸준히 이어나갈 생각이다.