Token-level Direct Preference Optimization

Fine-tuning pre-trained Large Language Models (LLMs) is essential to align them with human values and intentions. This process often utilizes methods like pairwise comparisons and KL divergence against a reference LLM, focusing on the evaluation of full answers generated by the models. However, the generation of these responses occurs in a token level, following a sequential, auto-regressive fashion. In this paper, we introduce Token-level Direct Preference Optimization (TDPO), a novel approach to align LLMs with human preferences by optimizing policy at the token level. Unlike previous methods, which face challenges in divergence efficiency, TDPO incorporates forward KL divergence constraints for each token, improving alignment and diversity. Utilizing the Bradley-Terry model for a token-based reward system, TDPO enhances the regulation of KL divergence, while preserving simplicity without the need for explicit reward modeling. Experimental results across various text tasks demonstrate TDPO's superior performance in balancing alignment with generation diversity. Notably, fine-tuning with TDPO strikes a better balance than DPO in the controlled sentiment generation and single-turn dialogue datasets, and significantly improves the quality of generated responses compared to both DPO and PPO-based RLHF methods. Our code is open-sourced at https://github.com/Vance0124/Token-level-Direct-Preference-Optimization.

한 줄 요약

TDPO는 토큰 수준에서 KL 발산을 제어하여 LLM의 정렬과 생성 다양성을 균형 있게 향상시키는 새로운 정렬 방법이다.

핵심 기여도

TDPO는 토큰 수준에서 KL 발산 제약을 적용하여 DPO 대비 더 높은 정렬과 생성 다양성을 달성함.
Bradley-Terry 모델을 토큰 수준으로 변환하여 별도의 보상 모델 없이도 정렬을 수행함.
IMDb, Anthropic HH, MT-Bench 데이터셋에서 DPO와 PPO 기반 RLHF 대비 생성 품질을 80% 이상의 인간 선호도로 개선함.
TDPO는 α=0.5로 설정 시 KL 발산 조절 능력이 가장 우수함.

핵심 아이디어

기존 DPO는 문장 수준에서 KL 발산을 제어하지만, 생성 과정은 토큰 단위로 이루어지므로 토큰 수준에서 발산을 제어하는 것이 더 효과적일 수 있다. TDPO는 이 점을 고려하여, 각 토큰에 대해 forward KL divergence 제약을 적용함으로써 생성 다양성을 유지하면서도 정렬 성능을 향상시킨다. Bradley-Terry 모델을 토큰 수준으로 확장하여 보상 함수를 정의하고, Bellman 방정식을 통해 문장 수준 보상과 토큰 수준 생성 간의 관계를 연결함. 이는 별도의 보상 모델 학습 없이도 정렬을 가능하게 한다.

기술적 접근법

**모델**: GPT-2 Large, Pythia-2.8B를 기반 모델로 사용.
**알고리즘**: TDPO는 DPO의 목적 함수를 토큰 수준으로 재구성하고, forward KL divergence 제약을 추가함.
**하이퍼파라미터**: α=0.5로 설정하여 KL 발산 조절과 보상 최적화 간 균형을 유지.
**보상 모델**: TDPO는 별도의 보상 모델 없이 Bradley-Terry 모델을 토큰 수준으로 변환하여 사용.
**KL 발산 제약**: 각 토큰에 대해 forward KL divergence를 적용하여 생성 다양성을 유지.

주요 결과

**IMDb 데이터셋**: TDPO는 DPO 대비 더 높은 보상과 낮은 KL 발산을 달성함.
**Anthropic HH 데이터셋**: TDPO는 DPO 대비 생성 다양성과 정렬 성능을 균형 있게 향상시킴.
**MT-Bench**: TDPO는 GPT-4 기준으로 80% 이상의 인간 선호도를 달성함.
**KL 발산 제어**: TDPO는 DPO 대비 dispreferred 응답의 KL 발산 증가 속도를 30% 이상 억제함.
**α 파라미터**: α=0.5일 때 TDPO는 KL 발산 조절 능력이 가장 우수함.

의의 및 한계

TDPO는 기존 DPO의 KL 발산 제어 한계를 토큰 수준에서 해결함으로써, 정렬과 생성 다양성을 동시에 향상시키는 새로운 접근법을 제시한다. Bradley-Terry 모델을 토큰 수준으로 확장하여 별도의 보상 모델 없이도 정렬을 수행할 수 있다는 점에서 실용적 가치가 높다. 그러나 TDPO는 토큰 수준의 복잡한 제약을 적용하기 때문에, 대규모 모델에서의 계산 효율성은 추가 연구가 필요하다. 또한, α 값의 선택이 성능에 큰 영향을 미치므로, 다양한 도메인에서의 최적화가 필요하다.

실용적 활용

TDPO는 감성 생성, 대화 시스템, 코드 생성 등 토큰 수준의 생성 다양성과 정렬이 중요한 분야에 적용 가능하다. 특히, 인간 선호도를 반영한 맞춤형 콘텐츠 생성, 안전한 대화 시스템 개발, 그리고 생성 모델의 윤리적 정렬에 유용하게 활용될 수 있다.