Diffusion-LM Improves Controllable Text Generation

• 목적 : 텍스트 생성 + 제어력 확보 모델

1. Introduction

Background

• LLM은 Controllability가 부족하여 특정 문장 구조나 문장 스타일을 강제하기 힘듬
  • 감정, 주제같은 단순 속성은 제어 가능하지만, 복잡한 문장 구조나 스타일의 일관성 유지 못함
  • 기존 Autoregressive base LLM은 왼쪽에서 오른쪽으로 next token prediction 하기에, 전체적인 문장 구조를 고려하기 어려움
  • 모델을 새로운 제어조건에 맞게 이용하려면 추가 학습을 해야하지만, 많은 비용이 듦

Related Works

Plug-and-Play

• LLM 을 그대로 두고, 외부 classifier를 이용하여 contrability 확보
  • PPLM : LLM의 hidden state를 조정하여 문장을 수정하는 방식
  • FUDGE : 문장의 앞부분을 보고 classifier를 통해 제어 가능성 판단 후, 다음 단어를 조정
• 하지만 해당 방법은 감정, 주제와 같은 간단한 제어만 가능하고, 문장 전체의 구문 구조를 통제하기는 힘듦.

Non-Autoregressive LM

• 기존 Autoregressive LLM 의 한계를 극복하기 위해 Non-Autoregressive LM이 연구되고 있음
• 주로 MT에서 사용되지만, 일반적인 자연어 생성에는 적합하지 않음

Diffusion Models for Text

• Diffusion은 이미지, 오디오 생성에서 성공적으로 많이 활용되고 있음
• 텍스트 생성에서도 활용하려는 시도가 있었으나, 대부분 Discrete Diffusion 방식을 채택하여 생성의 순서가 제한됨. 즉, 문장을 유연하게 생성하고 수정하기에는 한계가 있었음
  • Forward에서 모든 단어를 random하게 바꾸더라도, Reverse Process에서 토큰을 점진적으로 원래 토큰으로 변환해야 하므로, 바꾼 순서대로 역변환하기에 순서가 생김
  • 결국, Discrete Diffusion은 한번에 전체 문장을 생성하는 것이 아니라, 부분적으로 복원해 나가는 방식이기에, 생성 순서가 제한되어 구문 구조를 자유롭게 수정히기 힘듬
  • 토큰 단위 변경이기에, 문장 구조 유지가 힘듬
  • 토큰 단위 변경이기에, 문맥 고려 못함
  • discrete한 토큰이기에, 세밀한 미분 불가능
• Discrete Diffusion vs Continuous Diffusion
  • Discrete Diffusion : 단어 기반 Forward Process, 순차적으로 랜덤 단어로 변형
    • 문장 -> Tokenize -> Forward Process (랜덤 토큰으로 변경) -> Embedding -> Diffusion Transformer (Denoising Model) -> Reverse Process (Sampling) 하면서 문장 복원
  • Continuous Diffusion : 임베딩 벡터 기반 Forward Process, 임베딩 벡터를 Noise로 변형 (벡터 공간 입장에서 diffusion)
    • 문장 -> Tokenize -> Embedding -> Forward Proces (가우시안 노이즈로) -> Diffusion Transformer -> Reverse Process 하면서 문장 복원

Diffusion-LM

embedding space 에서 diffusion process를 통해 문장을 생성하기에, gradient based control이 세밀하게 가능하게 되어 특정 문체나 스타일을 강제할 수 있게됨

2. Method

Task-ReDefine

1. Text is Discrete Data !!

• (문제) diffusion은 연속적 데이터에서 잘 동작함 (image, audio), 하지만 text는 이산 데이터임
• (해결) text를 embedding vector로 바꾼 상태에서 diffusion 적용

2. Diffusion에서 문장의 의미를 유지해야함 !!

• (문제) diffusion은 noise를 추가하면서 원래 데이터 특성 잃어버림
• (해결) trainable한 단어 임베딩을 활용하고 diffusion 적용
  • 기존 LLM 처럼 고정된 단어 임베딩을 사용하면 noise add하는 과정에서 의미를 유지 못하게됨

3. Reverse 에서 단어를 정확히 복원해야함 !!

• (문제) diffusion은 gaussian noise 제거하는 방식인데, 단어를 정확히 예측해야함
• (해결) Softmax 기반 decoding과 추가적인 rounding 기법 도입

Continuous Embedding

기존 방법

• AutoRegressive : 문장 -> Tokenize -> Fixed Embedding -> Next Token Prediction
• Discrete Diffusion : 문장 -> Tokenize -> Forward Process -> Fixed Embedding -> Diffusion Model -> Reverse Process
  • 결국, 둘다 Token level에서 diffusion이 이뤄짐

Suggest

• Continuous Diffusion + Trainable Embedding
• 문장 -> Tokenize -> Trainable Embedding -> Forward Process -> Diffusion Model -> Reverse

Vector to Word

• vector를 word embedding으로 정확히 매핑시켜야함 -> 하단의 두 방법 사용

Rounding Algorithm

• Diffusion 은 노이즈를 부여하기 때문에 복원되어도 정확히 단어 벡터와 안맞을 수 있음
• 하지만 단어 임베딩과 정확히 일치해야 올바르게 복원됨
• 가장 가까운 단어 임베딩을 선택하는 방식 채용

Clamping Trick

• Rounding 과정에서 벡터가 단어 임베딩과 가까워지도록 학습 과정에서 추가 보정
• 특정 Threshold 이상 임베딩벡터 값 차이가 나면, 벡터를 강제로 가장 가까운 단어 임베딩으로 조정 (Clamping)

3. Generation

• 기존 방법들 (AutoRegressive, Discrete Diffusion) 은 순차적으로 생성
• 제안 방법은 노이즈 덩어리에서 한꺼번에 문장 생성하기에, 기존 LLM보다 유연한 수정이 가능해짐

Process

1. Random Noise Vector (XXX XXXX XXXX XXXX XXXXX)
2. Reverse Process - 학습된 모델을 이용하여 노이즈를 제거하면서 문장 복원
3. Rounding - 최종 생성된 Continuous Embedding을 Discrete한 Token으로 변환

• 문장 전체를 동시에 복원하면서 점진적으로 의미 있는 문장을 생성하는 방식

Controllable Generation

1. How to Control?

• Classifier-Guided Generation 수행
  • (기존) 스타일에 맞게 finetuning 해야함
  • (제안) 분류기를 사용하여 문장 생성 조절 가능 -> 추가학습 필요없음!

2. controllable generation

• 별도 목적을 갖는 Classifier 학습시켜놓고, Reverse Process에 이 classifier를 붙여서 Sampling 가능!!
• 해당 paper에서는 문장 유창성, 문장 스타일 classifier로 학습하여 제어기로 이용함
  • 기존 diffusion 모델처럼 노이즈만 제거하면 Fluency가 떨어질 수 있음
  • Fluency Classifier는 유창성 점수를 평가하며, gradient를 활용하지 않고, 유창성 점수에 따라 보정 수행하면서 Reverse process 수행하게 됨
  • 문장 스타일 classifier는 감정, 길이, 구문 구조 등 조정을 위해 활용되며, gradient를 직접 사용하여 문장을 수정하여 특정 스타일을 반영하도록 최적화하게 됨