Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebysheff Scalarization

이 논문은 선형 스칼라화의 한계를 극복하고 비볼록 파레토 프론트를 효과적으로 복원하기 위해 부드러운 체비셰프 스칼라화를 도입한 새로운 오프라인 강화학습 알고리즘인 STOMP 를 제안하며, 이를 통해 단백질 엔지니어링 등 다중 목적 최적화 작업에서 기존 최첨단 방법보다 우수한 성능을 입증했습니다.

핵심

1. 문제 상황: "완벽한 학생"을 찾는 딜레마

생각해 보세요. 여러분이 자녀를 교육한다고 가정해 봅시다.

• 목표 A: 수학 점수를 100 점 만점에 100 점 받아야 합니다.
• 목표 B: 동시에 예술적 감각도 최고 수준이어야 합니다.

문제는 이 두 가지가 서로 충돌할 수 있다는 점입니다. 수학 공부를 너무 많이 하면 예술 활동 시간이 줄어들고, 반대로 예술에 집중하면 수학 실력이 떨어질 수 있죠.

기존의 AI 기술들은 이 문제를 해결하기 위해 **"가중치"**라는 방식을 썼습니다.

"수학 점수에 70%, 예술 점수에 30% 를 줘서 합산한 점수가 가장 높은 학생을 찾자."

하지만 이 방법에는 치명적인 결함이 있습니다. 수학과 예술의 균형이 아주 미묘하게 다른 '최고의 학생'들 (예: 수학 90 점, 예술 95 점인 학생) 을 놓쳐버릴 수 있다는 것입니다. 마치 지형도가 울퉁불퉁한 산맥인데, 평평한 직선만 그어서 가장 높은 꼭대기를 찾으려다 보니, 사실은 그보다 더 멋진 구석진 곳에 있는 정상들을 놓치는 것과 같습니다.

2. 해결책: "스무스 체비셰프 (Smooth Tchebysheff)"라는 새로운 나침반

이 논문은 기존의 '가중치 합산' 방식 대신, **'스무스 체비셰프 (Smooth Tchebysheff)'**라는 새로운 나침반을 제안합니다.

• 기존 방식 (직선): "수학 점수가 조금 떨어지면 예술 점수가 많이 올라와야 전체 점수가 오른다"라고 계산합니다. 하지만 이 방식은 산맥의 구석진 부분 (비볼록 영역) 을 찾아내지 못합니다.
• 새로운 방식 (STOMP): "가장 낮은 점수 (약점) 가 가장 높은 학생을 찾아라"는 원리를 사용합니다.
  • 예를 들어, "수학 80 점, 예술 90 점인 학생"과 "수학 95 점, 예술 70 점인 학생"이 있다면, 전자의 '약점 (80 점)'이 후자의 '약점 (70 점)'보다 높으므로 전자를 더 선호합니다.
  • 이 방법은 어떤 구석진 산꼭대기에서도 가장 높은 지점을 찾아낼 수 있는 강력한 나침반입니다.

3. 핵심 기술: "STOMP" (스무스 체비셰프 최적화)

저자들은 이 나침반을 STOMP라는 이름의 새로운 알고리즘으로 만들었습니다. STOMP 는 다음과 같은 두 가지 똑똑한 작업을 합니다.

1. 척도 맞추기 (Standardization):

   • 예를 들어, '단백질 활성도'는 0100 점이고, '안정성'은 01000 점일 수 있습니다. 그냥 합치면 안정성 점수만 중요해집니다.
   • STOMP 는 각 점수들을 서로 다른 기준 (분포) 에 맞춰서公平하게 비교할 수 있도록 조정합니다. 마치 키가 150cm 인 사람과 200cm 인 사람을 비교할 때, 키의 절대값이 아니라 '자신 그룹 내에서 얼마나 큰지'로 비교하는 것과 같습니다.

2. 균형 잡기:

   • 이 나침반을 통해 AI 는 "수학 90 점, 예술 90 점" 같은 **모든 목표를 골고루 잘 수행하는 완벽한 학생 (파레토 최적 해)**들을 찾아냅니다.

4. 실험 결과: 단백질 설계에서의 대활약

이 논문에서는 이 STOMP 기술을 실제 단백질 설계에 적용해 보았습니다.

• 상황: 연구원들은 단백질이 "효소 활성 (일 잘함)"과 "안정성 (오래 견딤)"을 동시에 갖기를 원했습니다.
• 결과: 기존 방법들 (DPO 등) 보다 STOMP 를 사용한 AI 가 훨씬 더 다양한, 그리고 더 훌륭한 단백질들을 만들어냈습니다.
  • 마치 다양한 맛 (신맛, 단맛, 쓴맛) 을 모두 만족시키는 요리를 만드는 셰프처럼, STOMP 는 서로 충돌하는 목표 사이에서 최고의 균형을 찾아냈습니다.

5. 요약: 왜 이 논문이 중요한가요?

이 논문은 **"하나의 정답만 찾는 것이 아니라, 다양한 상황에서의 '최고의 균형'을 찾아내는 AI"**를 만들었습니다.

• 비유: 기존 AI 가 "가장 높은 산 하나만 찾는 등산가"였다면, STOMP 는 **"산맥 전체의 모든 정상과 아름다운 골짜기를 모두 찾아내는 탐험가"**입니다.
• 의의: 이 기술은 단백질 설계뿐만 아니라, 챗봇 (유용함과 안전함의 균형), 이미지 생성 (화질과 지시사항의 균형) 등 우리가 살면서 겪는 모든 복잡한 의사결정 문제에 적용될 수 있는 강력한 도구입니다.

결론적으로, 이 논문은 **"서로 충돌하는 여러 목표를 동시에 만족시키는, 더 똑똑하고 균형 잡힌 AI"**를 만드는 길을 열었습니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: 대규모 언어 모델 (LLM) 은 인간 선호도에 맞춰 정렬 (Alignment) 되기 위해 종종 오프라인 강화 학습 (Offline RL) 을 사용합니다. 기존 연구는 주로 단일 목적 함수 (예: 유용성) 를 최적화하는 데 집중했습니다.
• 핵심 문제: 실제 응용 분야 (예: 단백질 공학, 챗봇) 에서는 상충되는 여러 목표를 동시에 최적화해야 하는 경우가 많습니다.
  • 예시: 단백질 공학에서는 '촉매 활성'과 '특이성'을 동시에 높여야 하거나, 챗봇에서는 '유용성'과 '안전성'을 모두 만족해야 합니다.
• 기존 방법의 한계: 다목적 최적화 (Multi-Objective Optimization) 를 해결하기 위해 기존에는 선형 보상 스칼라화 (Linear Reward Scalarization) 를 주로 사용했습니다. 이는 각 보상에 가중치를 부여하여 합산하는 방식입니다.
  • 결함: 선형 스칼라화는 수학적으로 증명된 바와 같이 비볼록 (non-convex) 영역의 파레토 프론트 (Pareto front) 를 복원할 수 없습니다. 즉, 상충되는 목표 간의 최적 절충안 (compromise) 인 중요한 해들을 놓치게 됩니다.

2. 방법론 (Methodology)

이 논문은 선형 스칼라화의 한계를 극복하기 위해 Smooth Tchebysheff Scalarization (STS) 을 오프라인 다목적 RL 에 적용한 새로운 알고리즘 STOMP (Smooth Tchebysheff Optimization of Multi-Objective Preferences) 를 제안합니다.

핵심 아이디어

1. 문제 재정의: 단순히 보상을 스칼라화하는 것이 아니라, 다목적 RL 문제 자체를 스칼라화할 최적화 문제로 간주합니다.
2. Smooth Tchebysheff Scalarization (STS) 적용:
   • 기존 TS 는 Min-Max 형태라 미분이 불가능하여 최적화가 어렵습니다. Lin et al. [42] 이 제안한 STS 는 로그 - 합 - 지수 (logsumexp) 를 사용하여 미분 가능하게 만듭니다.
   • STS 는 파레토 프론트의 모든 해 (비볼록 영역 포함) 를 찾을 수 있습니다.
3. 보상 표준화 및 분산 기반 정규화:
   • STS 는 개별 보상의 스케일에 매우 민감합니다. 이를 해결하기 위해 논문은 관측된 분포에 기반한 보상 표준화를 도입합니다.
   • 각 보상 $r_i$를 해당 보상의 분산 $\sigma_i$로 나누고, 분할 함수 (partition function) $Z_i(x)$를 추정하여 분산 상대적 보상 (distribution-relative reward) $\rho_i$를 정의합니다.
   • 이를 통해 보상의 스케일 차이를 자동으로 보정하고, 이상치 (outliers) 에 덜 민감하게 만듭니다.
4. STOMP 알고리즘 (Loss Function):
   • Direct Preference Optimization (DPO) 프레임워크를 다목적 환경으로 확장합니다.
   • 이론적 유도: KL 발산 제약 하에서 최적 정책 $\pi^*$ 들 간의 거리를 최소화하는 Upper Bound 를 유도하여 스칼라화된 보상 $R_{ST}$를 도출합니다.
   • 손실 함수: OffsetDPO 를 기반으로 하되, 선형 가중치 대신 STS 기반의 보상 차이를 사용합니다.
     • 선호도 쌍 (Winner/Loser) 결정 시: 정책 독립적인 근사 보상 $R_{ST}$ 사용 (안정성 확보).
     • 개별 Loss 항 계산 시: 정책 의존적인 보상 $R_{ST}^\pi$ 사용 (더 엄격한 최적화 문제 반영).
     • 정규화: 승자 (Winner) 시퀀스의 로그 가능도 (NLL) 를 패널티로 추가하여 오프라인 RL 에서 발생하는 분포 이탈 (distribution shift) 문제를 완화합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 알고리즘 STOMP: 다목적 오프라인 RL 을 위해 설계된 최초의 체계적인 알고리즘으로, DPO 를 다목적 설정으로 자연스럽게 확장합니다.
2. 비볼록 파레토 프론트 복원: 선형 스칼라화가 실패하는 비볼록 영역의 해를 성공적으로 찾아내는 것을 이론적으로 및 실험적으로 증명합니다.
3. 자동 보상 표준화: 개별 보상의 스케일 조정을 위한 하이퍼파라미터 튜닝 없이, 데이터 분포를 기반으로 자동으로 보상을 표준화하는 메커니즘을 제시합니다.
4. 단일 목적 RL 과의 호환성: 기존 DPO 및 OffsetDPO 와 동일한 프레임워크 내에서 구현 가능하며, 기존 RLHF 파이프라인에 통합하기 용이합니다.

4. 실험 결과 (Results)

논문은 단백질 공학 (Protein Engineering) 작업을 통해 STOMP 를 검증했습니다.

• 데이터셋:
  • DHFR: TMP 존재 유무에 따른 DHFR 활성 (부정 상관).
  • PbrR: Pb2+ 결합 증가 및 Zn2+ 결합 감소 (강한 부정 상관).
  • $\alpha$-Amylase: 활성, 발현, 열안정성 (양정 상관).
• 모델: ProGen3-3B, ProGen-RA-3B, ProGen-RA-10B 등 3 개의 단백질 언어 모델 사용.
• 평가 지표: 초부피 (Hypervolume) - 파레토 프론트에서 비우세 해들이 차지하는 영역의 크기로, 값이 클수록 다목적 최적화 성능이 우수함을 의미합니다.
• 성과:
  • 오프라인 오프-폴리시 평가: 9 가지 설정 중 8 가지에서 STOMP 가 가장 높은 초부피를 기록했습니다.
  • 생성 평가 (Generative Evaluation): 학습된 모델이 생성한 단백질 시퀀스를 평가한 결과에서도 8 가지 설정 중 8 가지에서 STOMP 가 최상의 성능을 보였습니다.
  • 비교: 기존 선형 스칼라화 (DPO-Lin, ODPO-Lin) 및 기존 STS 기반 방법 (ODPO-STZ) 보다 일관되게 우수한 성능을 보였으며, 특히 PbrR 과 같은 강한 상충 관계를 가진 데이터셋에서 격차가 두드러졌습니다.

5. 의의 및 결론 (Significance)

• 다목적 정렬의 새로운 표준: 단일 목표 최적화를 넘어, 상충되는 여러 목표를 동시에 고려해야 하는 현실적인 AI 응용 (챗봇, 의료, 과학 발견 등) 에 필수적인 파레토 최적 해를 효율적으로 찾을 수 있는 방법을 제공합니다.
• 과학적 발견 가속화: 단백질 공학 분야에서 활성, 안정성, 발현량 등 여러 특성을 동시에 만족하는 새로운 단백질을 설계하는 데 있어 기존 방법보다 훨씬 강력한 도구가 됩니다.
• 범용성: 단백질 언어 모델에 국한되지 않고, 텍스트 생성, 이미지 생성 등 모든 다목적 정렬이 필요한 LLM 응용 분야에 적용 가능한 범용 알고리즘입니다.

이 논문은 Smooth Tchebysheff Scalarization을 RL 프레임워크에 통합함으로써, 기존 RLHF 의 한계를 넘어 더 복잡하고 실용적인 다목적 의사결정 문제를 해결할 수 있는 강력한 기반을 마련했습니다.