본문 바로가기
Drawing (AI)/Paper review

[논문 리뷰] Confidence-aware Reward Optimization for Fine-tuning Text-to-Image Models

by 생각하는 이상훈 2024. 7. 15.
728x90

Intro

생성형 모델은 정확도로 평가하여 성능을 결정하는 것을 떠나서 인간이 보기에 자연스러운 것, 즉 alignment가 굉장히 중요하다. 인간의 의도와 생성형 모델의 결과물의 방향이 일치하는지를 따져야 한다는 것이다. 해당 논문은 Text-to-Image model의 alignment를 잘 평가할 수 있는 benchmark를 제시하고 잘 align된 모델을 만들 수 있는 reward model을 제시한다.


Background

Alignment

Text-Image ModelAlignment는 intro에서 말했듯 모델이 생성하는 이미지가 사용자의 의도와 정확하게 일치하는지를 보는 것이다. 예를 들어 사과와 사슴을 generate하라고 했을때 좌측은 비교적 잘 align된 반면 우측 image는 그렇지 않은 것을 볼 수 있다.

Alignment를 개선하면 모델의 신뢰도와 활용도를 크게 향상시킬 수 있을 것으로 기대된다.


Text-to-image model

Text-to-image model의 성능을 높이기 위해 인간 자원을 사용하는 방법은 여러 가지가 있다. 주요 방법 중 하나는 RLHF(Reinforcement Learning from Human Feedback, 인간 피드백을 통한 강화 학습)이다. 이 접근 방식은 모델의 출력에 대해 인간이 직접 피드백을 제공하여 모델을 개선하는 것을 포함한다.
우선, 모델이 학습할 수 있도록 대규모의 텍스트-이미지 쌍 데이터를 수집한다. 인간 레이블러가 각 텍스트 설명에 대해 정확하고 관련성 높은 이미지를 선택하거나 생성하며, 이미지의 품질과 관련성을 평가하고 부정확하거나 낮은 품질의 이미지를 제거한다. 이후에는 인간 피드백을 사용하여 모델의 출력을 평가하고, 이를 통해 모델을 지속적으로 업데이트한다. 
예를 들어 gpt3.5는 gpt3에 RLHF로 학습을 시켜서 성능을 굉장히 높히는 것에 성공하였다.


TIA2(Text-Image Alignment Assessment)

최근 몇 년 동안 대규모 텍스트-이미지 모델은 텍스트 프롬프트를 입력받아 고품질의 창의적인 이미지를 생성하는 데 성공했다. 그러나 이러한 모델은 구성적 생성 능력, 정확한 텍스트 렌더링, 공간 이해 등 여러 측면에서 여전히 한계가 있다. 또한, 모델을 훈련시키는 데 사용되는 대규모 데이터 세트는 악성 콘텐츠와 바람직하지 않은 편향을 포함할 수 있다. 이러한 한계를 극복하기 위해 인간 피드백 데이터를 활용한 보상 함수를 학습하여 모델의 동작을 인간의 의도에 맞추는 방법이 주목받고 있다. 그러나 이러한 보상 모델은 단순히 대리 목표일 뿐이므로 과도한 최적화는 모델의 성능을 저하시킬 수 있다. 

이 문제를 해결하기 위해 다양한 텍스트 프롬프트, 이미지 및 인간 주석으로 구성된 텍스트-이미지 정렬 평가(TIA2) 벤치마크를 도입했다. TIA2 벤치마크는 여러 출처에서 샘플링한 100개의 텍스트 프롬프트와 27,500개의 이미지를 포함하며, 각 텍스트-이미지 쌍에 대해 세 명의 인간 평가자로부터 바이너리 피드백(좋음 또는 나쁨)을 수집했다. 이를 통해 다양한 텍스트-이미지 정렬 시나리오에서 여러 최신 보상 모델을 평가할 수 있다. 과도 최적화 문제를 해결하기 위해, 우리는 모델의 보상 모델 신뢰도를 기반으로 정렬을 향상시키는 간단한 방법인 텍스트 정규화(TextNorm)를 제안했다. TextNorm은 입력 프롬프트와 의미적으로 대조적인 프롬프트 세트에 대한 보상을 비교하여 보상을 조정한다. 이를 통해 모델의 신뢰도를 측정하고 보상을 교정할 수 있다. 우리는 여러 보상 모델을 TIA2 벤치마크에서 평가했으며, TextNorm을 적용한 결과 과도 최적화를 효과적으로 줄일 수 있음을 입증했다. TextNorm은 인간 평가에서 기존의 보상 모델보다 두 배 더 많은 승리를 기록했다. TIA2 벤치마크와 TextNorm은 텍스트-이미지 모델의 보상 최적화 문제를 해결하는 데 중요한 도구다. 이를 통해 모델의 인간 의도 정렬을 개선하고 과도 최적화로 인한 성능 저하를 방지할 수 있다. 

x: text prompt, y: image, z: binary human label {0,1}
결과가 z와 일치하도록 하는 reward function r(x,y)이 목표이다. 이전의 평가방식을 이용하면 인간이 판단하는 결과와 조금은 상반된 결과가 나오는 것을 확인할 수 있다.


Confidence-Calibrated Reward Design

논문에서 보상 모델의 정렬을 향상시키기 위해 우리는 모델 신뢰도의 척도를 기반으로 한 보상 설계를 탐구했다. 인간 피드백 데이터로 미세 조정된 보상 모델이 인간 선호도를 충분히 반영하지 못하는 경우가 많다는 것을 발견한 저자들은 이를 해결하기 위해 텍스트 정규화(TextNorm)라는 간단한 방법을 제안했다. TextNorm은 입력 프롬프트와 의미적으로 대조적인 프롬프트 세트에 대한 보상을 비교하여 보상을 조정한다. 이는 모델이 입력 프롬프트에 대해 얼마나 신뢰할 수 있는지를 측정하고 그에 따라 보상을 교정하는 것이다. 예를 들어, 입력 프롬프트 "두 마리의 개 사진"에 대해 보상이 높은 경우와 대조 프롬프트 "세 마리의 개 사진"에 대해 보상이 낮은 경우, 이는 모델이 입력 프롬프트에 대한 신뢰도가 높음을 나타낸다. 이러한 방식으로 TextNorm은 보상 모델의 신뢰도를 반영하여 보상을 조정한다.
TextNorm을 구현하는 중요한 단계는 대조 프롬프트 세트를 생성하는 것이다. 간단한 입력 프롬프트의 경우 규칙 기반 접근 방식이 효과적일 수 있지만, 더 복잡한 경우에는 대형 언어 모델(LLM)을 활용하여 대조 프롬프트를 생성할 수 있다. LLM을 사용하면 입력 프롬프트의 실패 시나리오를 분석하고, 이에 대한 텍스트 설명을 포함한 대조 프롬프트를 생성할 수 있다. 이를 통해 TextNorm은 보상을 보다 정확하게 교정할 수 있다. 우리는 TextNorm의 효과를 TIA2 벤치마크에서 여러 보상 모델을 평가하여 입증했다. TextNorm은 기존 보상 모델보다 인간 판단과의 정렬을 현저히 향상시키며, 과도 최적화를 줄이는 데 성공했다. 이러한 결과는 TextNorm이 보상 모델의 신뢰도를 기반으로 보상을 교정하는 방법이 효과적임을 보여준다. 


Flow

전체 흐름을 보자면 input image와 그에 해당하는 prompt와 contrastive prompts를 입력하여 두가지의 scoring법을 이용하고 softmax에 태워서 확률을 ensemble하여 reward를 계산한다.

이때 ensemble 기법으로 Mean ensemble, Uncertainty-regularized ensemble이 존재한다. 이를 가볍게 살펴보면 아래와 같다.

Mean ensemble
Uncertainly-regularized ensemble

앙상블 기법은 mean 앙상블과 uncertainly-regularized 앙상블이 쓰였는데 mean은 단순한 평균을 이용하는 것이고 uncertainty-regularized는 분산을 기반으로 측정된 불확실성을 반영하여 불확실성이 크면 람다를 줄여 적게 반영하고 불확실성이 작으면 람다를 키워 많이 반영하는 기법이다.

Meancomposition data set에 대해, uncertainty-regularizedcounting, comprehensive data set에 대해 적용하였다. 정량적 평가를 위해 CLIP, BLIP, ImageReward, PickScore, TextNorm에 테스트를 해보니 TextNorm기법이 세가지 benchmark에 대해서 가장 좋은 성능을 내는 것을 볼 수 있었다. ablation study를 통해 textnorm과 앙상블을 제거하거나 변화를 줘보니 LLM을 이용한 contrastive prompt를 이용하는 것과, 앙상블 기법이 유의미한 성능향상을 불러 일으키는 것을 볼 수 있다.


Experiment

 

실험을 위해서 Stable Diffusion v2.1base model로 이용하였고 Best-of-n은 각 데이터셋에서 좋은 결과가 나온 30개의 prompt, Supervised Fine-tuningRL fine-tuning은 그 30개중 10개의 prompt로 실험했다.

prompt에 대해서 n개의 sample image를 생성하고 그중 score가 가장 높은 이미지를 선택하는 방식으로 n이 커질 수록 더 좋은 이미지가 나오는 경향이 있으나 BLIP-2의 경우 n이 증가하면 품질이 저하되고 왼쪽의 경우에서 n이 증가하면 ImageReward모델은 일부분이 짤리는 현상이 발생하고 오른쪽의경우 CLIP모델에서 n이 증가하면 비현실적으로 변하는 문제가 발생했다.

Supervised Fine-Tuning에서는 현재 policy를 기반으로 여러 state-action을 샘플링하고 각 샘플에 대한 reward 기반 가중치를 이용하여 regression 모델을 학습시켜 최적의 action을 예측한 뒤에 이를 기반으로 policyupdate하는 reward-weighted regression기법을 사용한다. 상위 10개의 이미지를 선택하여 데이터셋을 구성하고, reward-weighted likelihood를 최대로 만드는 방향으로 fine-tuning한. Fine-tuning 과정에서는 reward-weighted likelihood를 최대화하는 방향으로 학습한다. 수식을 보면 reward에 조건부 확률이 곱해져 확률에 따라 결과가 바뀌는 것을 알 수 있다. 또한, 정규화항에서 β는 원래 모델 p에서 너무 많이 벗어나지 않도록 방지하는 KL 정규화 계수이다.

RL fine-tuning은 다양한 데이터 분포를 수집하여 Multi-step Markov Decision ProcessDDPO를 사용해 최적의 정책을 학습한다. Denoising Diffusion Policy Optimization(DDPO)Diffusion과 강화학습의 policy 최적화를 결합한 기법이다. Diffusion같이 noise를 더하고 제거하는 과정에서 policy를 수정하며 학습시키는 방식이다. 또한, 여기서의 λ도 원래 모델에서 너무 멀리 벗어나는 것을 방지하는 정규화 항이다.

 

Human Evaluation에서 우리는 파인 튜닝된 모델을 사용하여 프롬프트당 8개의 이미지를 생성하고, 이를 4개의 이미지로 이루어진 세트로 나누었다. 60세트가 명의 독립적인 평가자에 의해 평가되었다. 결과를 보면 textnorm방식이 RL fine-tuning케이스 에서는 image quality 조금 밀리는 모습이 있으나 논문에서 target으로한 alignment 관점에서는 케이스 모두 압도적이었다.


Conclusion

결과적으로 해당 논문은 Human feedback을 이용한 강화학습에서 발생하는 Overoptimization으로 인한 문제를 신뢰도가 높고 일반화된 TIA2 benchmark를 제안하여 해당 benchmark에 대한 결과를 통해 보여주었고 이를 TextNorm기법을 이용한 reward model optimization을 통해 alignment를 개선하는 결과를 도출했다.


Evaluating text-tmage ge RL fine-tuning은 다양한 데이터 분포를 수집하여 Multi-step Markov Decision ProcessDDPO를 사용해 최적의 정책을 학습합니다. Denoising Diffusion Policy Optimization(DDPO)Diffusion과 강화학습의 policy 최적화를 결합한 기법입니다. Diffusion같이 noise를 더하고 제거하는 과정에서 policy를 수정하며 학습시키는 방식 입니다. 또한, 여기서의 λ도 원래 모델에서 너무 멀리 벗어나는 것을 방지하는 정규화 항입니다. nerat

728x90