[Paper] 텍스트-이미지 Diffusion에서 Critic 기반 Reinforcement Unlearning

Source: arXiv - 2601.03213v1

개요

이 논문은 텍스트‑투‑이미지 확산 모델(예: Stable Diffusion, DALL·E 2 등)이 실제 환경에 점점 더 많이 적용됨에 따라 중요성이 커지고 있는 실용적인 문제를 다룹니다: 특정 개념—예를 들어 저작권이 있는 스타일이나 유해한 시각적 모티프—을 모델 전체 성능을 손상시키지 않으면서 “잊히게” 하는 방법. 저자들은 확산 디노이징 과정을 순차적 의사결정 문제로 보는 강화학습(RL) 프레임워크를 제안하고, 기존 RL 시도보다 더 안정적으로 언러닝을 유도하기 위해 시간‑스텝 인식 비평가(timestep‑aware critic) 를 도입합니다.

주요 기여

• RL‑기반 언러닝 프레임워크로 각 디노이징 단계를 행동으로 모델링하여 세밀한 크레딧 할당을 가능하게 함.
• 시간‑인식 비평가는 CLIP‑학습 보상 예측기를 기반으로 하며, 모든 확산 단계에서 노이즈가 섞인 잠재 표현을 평가해 단계별 밀도 높은 피드백을 제공함.
• 정책‑그래디언트 업데이트는 역확산 커널에 적용되며, 오프‑폴리시 데이터를 재사용할 수 있어 기존 확산 파이프라인과 호환 가능함.
• 실증 검증을 통해 여러 목표 개념에 대해 강력한 감독 기반 베이스라인과 동등하거나 더 나은 망각 효과를 보이며 이미지 품질 및 프롬프트 준수를 유지함을 입증함.
• 오픈‑소스 공개를 통해 코드, 평가 스크립트, 사전 학습된 비평가를 제공하여 재현성 및 향후 연구를 가속화함.

방법론

1. Sequential View of Diffusion – 역확산 과정(노이즈를 이미지로 변환하는 과정)을 마르코프 의사결정 과정으로 정의합니다:

   • State: 시간 단계 t에서의 현재 노이즈가 섞인 잠재 변수.
   • Action: 모델이 예측한 디노이징 방향(확산 커널의 출력).
   • Transition: 확산 단계를 적용하여 다음 시간 단계로 이동.

2. Critic Design – CLIP 기반 네트워크를 미세조정하여 노이즈가 섞인 잠재 변수와 목표 개념 텍스트로부터 스칼라 “언러닝 보상”을 예측하도록 합니다. 중요한 점은 비평가가 노이즈가 섞인 잠재 변수를 받는다는 것으로(깨끗한 이미지는 받지 않음), 따라서 각 확산 단계마다 학습 신호를 제공할 수 있습니다.

3. Reward Signal – 보상은 잠재 변수가 원하지 않는 개념으로부터 멀리 떨어져 있을 때(클립 유사도 기준) 높게, 그렇지 않을 때 낮게 설정됩니다. 비평가가 노이즈가 섞인 잠재 변수에서 작동하기 때문에 보상 자체가 자연스럽게 노이즈가 섞여 있으며 시간 단계에 따라 변동합니다. 이는 정책이 확산 궤적 중 개념이 가장 취약한 지점을 학습하는 데 도움이 됩니다.

4. Policy Update – 단계별 보상을 이용해 저자들은 어드밴티지 추정치를 계산하고, 표준 REINFORCE 스타일 정책 그래디언트를 적용하여 확산 커널의 파라미터를 조정합니다. 오프‑폴리시 샘플(예: 원본 모델이 생성한 잠재 변수)도 재사용할 수 있어 샘플 효율성이 향상됩니다.

5. Training Loop – 과정은 다음을 번갈아 수행합니다:

   • 목표 개념을 포함한 프롬프트 배치를 샘플링합니다.
   • 확산 과정을 실행하면서 상태, 행동, 비평가 보상을 수집합니다.
   • 비평가를 (주기적으로) 업데이트하고, 계산된 어드밴티지를 통해 확산 정책을 업데이트합니다.

결과 및 발견

• 시간 단계 인식 비평자는 그래디언트 업데이트의 분산을 크게 감소시켜 더 안정적인 학습과 더 빠른 수렴을 가능하게 합니다 (목표 망각 수준에 도달하기 위해 약 30 % 적은 diffusion 단계).
• 소거 실험 결과 (i) 단계별 비평자를 제거하고 (ii) 클린 이미지 전용 보상을 사용할 경우 모두 성능이 저하되어 언러닝이 부족하거나 눈에 띄는 아티팩트가 발생함을 확인했습니다.
• 정성적 예시들은 모델이 특정 아티스트의 스타일을 지우면서도 관련 없는 프롬프트에 대해 고품질 이미지를 생성할 수 있음을 보여줍니다.

실용적인 함의

• 컴플라이언스 및 IP 관리 – 기업은 배포된 확산 모델에서 처음부터 재학습 없이도 저작권이 있거나 상표가 등록된 시각 요소를 사후에 제거할 수 있습니다.
• 안전 및 중재 – 해롭거나 허용되지 않은 시각 개념(예: 극단주의 상징)을 실시간으로 제거할 수 있어 우발적인 생성 위험을 줄입니다.
• 모듈식 업데이트 – 이 방법이 기존 확산 백본 위에 플러그인 형태의 정책‑그라디언트 레이어로 작동하기 때문에, 개발자는 이를 CI 파이프라인에 통합하여 지속적인 “개념 위생”을 구현할 수 있습니다.
• 샘플 효율성 – 오프‑폴리시 재사용을 통해 이전에 생성된 이미지 로그를 활용할 수 있어 전체 감독식 파인튜닝에 비해 계산 비용을 낮출 수 있습니다.

제한 사항 및 향후 연구

• CLIP에 대한 보상 의존성 – 평가는 CLIP의 편향을 물려받는다; CLIP이 개념을 오분류하면 학습 제거 신호가 잡음이 많거나 잘못된 방향일 수 있다.
• 다수 개념에 대한 확장성 – 현재 설정은 목표 개념당 별도의 평론가를 학습한다; 동시에 여러 개념을 제거하도록 확장하는 것은 아직 해결되지 않은 과제이다.
• 이론적 보장 – 경험적 망각 효과는 강하지만, 개념이 얼마나 제거되었는지에 대한 공식적인 경계는 제공되지 않는다.
• 향후 방향은 저자들이 제시한 바와 같이 다중 작업 평론가 탐색, 보다 강인한 보상 모델 통합(예: 확산 기반 분류기), 망각 속도와 다운스트림 작업 성능 사이의 트레이드오프 연구 등을 포함한다.

저자

• Mykola Vysotskyi
• Zahar Kohut
• Mariia Shpir
• Taras Rumezhak
• Volodymyr Karpiv

논문 정보

• arXiv ID: 2601.03213v1
• 카테고리: cs.LG
• 출판일: 2026년 1월 6일
• PDF: Download PDF