[평범한 학부생이 하는 논문 리뷰] Detecting, Explaining, and Mitigating Memorization in Diffusion Models (ICLR 2024)

Paper : https://arxiv.org/abs/2407.21720

Detecting, Explaining, and Mitigating Memorization in Diffusion Models

Abstract

문제 : 생성모델의 몇 output은 training data를 단순히 복제함

해결 방안 : text-conditional predictions의 정도를 검사함으로써 memorized prompt를 detecting

1. Introduction

 첫번째 행처럼 training data가 그대로 생성되는 경우와 두번째 행처럼 training data의 구조를 가진 이미지를 생성하는 경우(memorization)가 발생한다.

 이러한 문제를 해결하기 위해서, 본 논문에서는 먼저 memorized prompt를 감지하는 새로운 방법을 제안한다. 본 논문은 이러한 prompt들에 대해, initialization이 무엇이든간에 text condition이 생성을 memorized solution을 향하도록 지속적으로 guide한다는 것을 발견했다. 이 현상은 denoising process동안 상당한 text guidance를 제안한다. 결국, 본 논문의 detection method는 text-conditional prediction의 정도를 기본적인 metric으로써 우선시 한다. 위 이미지에서처럼, memorized prompt에 대해서 더 높은 정도를 보인다. 본 논문의 method는 기존 framework에 추가적인 work를 추가하는 것 없이 그리고 심지어 많은 생성 없이 정확한 detection을 제공한다.

 이런 발견을 기초로, 본 논문은 특정한 trigger token을 찾는 것을 목표로 memorization을 이끄는 것에 대한 각 token의 영향을 강조하는 방식을 고안한다. 이러한 trigger token을 제거하는 것이 memorization을 상쇄한다는 직관에 따라, 그에 대응되는 text-conditional prediction의 정도 감소를 예상한다.

 마지막으로, memorization 문제를 해결하기위해 mitigation strategy를 도입한다. Inference와 training에 적용할 수 있는 두 가지 방식이 있다. 먼저 inference동안, text-conditional prediction의 정도를 최소화하는 perturbed prompt embedding을 사용하는 것을 제안한다. Training 동안, text-conditional predictions의 정도를 기초로 potential memorized image-text pair가 가려질 수 있다. 이 간단한 방식들은 prompt와 generation 사이의 더 일관된 alignment를 보장하고, memorization을 효과적으로 감소시킨다.

2. Detect Memorization Efficiently

2.1 Motivation

 Memorized prompt에 대해서, initialization은 관계가 없는 것으로 보인다. 생성된 이미지는 특정 memorized visualization으로 지속적으로 수렴한다. 이러한 행동은 model이 prompt와 특정 denoising trajectory에 overfitting하는 것을 암시하고 이는 memorized image를 이끈다. 결국, final image는 initial 상태로부터 상당히 멀어진다.

 이러한 발견은 간단한 detection strategy를 위한 foundation을 제공한다. text-conditional noise prediction의 작은 정도는 final image가 initialization과 가깝게 align된다는 것을 보여 memorized image가 아닐 것이다라는 것을 보인다. 반대로 정도가 큰 것은 potential memorization을 나타낸다. 이 상관관계는 Figure 2(a)에서 볼 수 있다.

2.2 An Effective Detecting Method

 위에서의 intuition을 기반으로, 본 논문은 간단하지만 효과적은 text-conditional noise prediction의 정도를 중심으로한 detection method를 제안한다.

$p$는 text embedding이고, $T$는 sampling step

Memorization은 tunable threshold $\gamma$아래로 떨어지면 식별된다.

 본 논문의 metric은 여러 이미지를 생성해야하는 기존의 방식과 다르게 한번의 generation에서도 강한 memorization 신호를 제공해 effective하고 reliable하다. 또한 기존의 방식과 다르게 training dataset에 대한 접근 없이 detection method를 사용할 수 있다. 또다른 이점은 detection metric을 계산하는데 있어서 적응력이다. First step으로부터 얻은 metric을 얻었을 때에도 reliable detection이 attainable하다. 그래서 즉시 memorized prompt를 식별할 수 있다.

2.3 Experiments

 Figure 2(b)를 보면 memorized prompt와 non-memorized prompt에서의 차이가 발생한다. 또한 memorized prompt에서는 bimodal인 것을 볼 수 있는데, 이는 matching verbatim(완전 memorization)과 template verbatim(동일한 형태지만 style은 다른 memorization)에서의 차이때문에 발생한다고 한다.

위 표에서, 본 논문의 방식의 precision과 efficiency에 대한 균형을 볼 수 있다. 많은 이미지를 생성하지 않아도, 많은 step을 밟지 않아도 결과가 좋게 나오는 것을 볼 수 있다.

3. Mitigate Memorization

3.1 A Straighforward Method to Detect Trigger Tokens

 기존 work에 따르면, memorized prompt에서 특정 word나 token이 generation process에서 중요한 영향을 끼친다고 한다. 앞서의 observation은 text-conditional noise prediction의 정도를 최소화하면서 각 토큰에 적용되는 변화의 정도를 확인함으로써 memorization과 관련된 각 토큰의 중요도를 구별하는 method를 제공한다. 상당한 변화를 가지는 토큰은 prediction을 지시하는데 있어서 중요한 역할을 제시하고 반대로 적은 변화를 가진 것은 덜 중요하다는 것을 보인다.

$N$개의 token을 가진 prompt $p$의 prompt embedding $e$가 주어졌을 때, objective는 다음과 같다.

각 토큰에 대한 significance score는 다음과 같이 정의한다.

Trigger token을 바꾸게되면 memorization을 완화하는데 도움을 준다.

3.2 An Effective Inference-Time Mitigation Method

어떠한 supervision없이 mitigation하는 방법은 (5)를 최소화함으로써 prompt embedding을 적용하는 것이다. 전체 step에 대해서 하면 computationally intensive하다. 본 논문에서는 초기 step에서 loss를 최소화하는 것이 이 후 step에서 더 작은 magnitude를 야기해서 효과적으로 memorization을 완화한다. Embedding은 loss 값 $l_{target}$를 통해 early stopping을 한다.

3.3 An Effective Training-Time Mitigation Method

Text-conditional noise prediction값이 $\tau$를 넘으면 mini-batch에서 sample을 배제함으로써 그 sample에 대해서 loss를 계산하지 않도록 한다. 이 방식은 10%의 학습시간 연장을 발생시킨다.

3.4 Experiments

4. Limitations and Future Work

 Detection의 경우 tunable threshold를 사용한다. 그래서 non-memorized prompt를 이용해서 적절한 threshold를 정해서 memorization을 detection을 하게된다. 이럴 경우 interpretability가 부족하다.