MARBLE: AI 보상 최적화의 새로운 해법
- •MARBLE은 복잡한 AI 훈련 보상을 수동 조정 대신 경사 공간 최적화로 자동 균형을 맞춘다.
- •5가지 서로 다른 이미지 품질 보상을 동시에 최적화하며 97%의 높은 훈련 효율을 달성했다.
- •연구진은 가중치 합산 방식의 한계를 극복하여 모델 미세 조정 과정에서 발생하던 보상 간 충돌 문제를 해결했다.
최근 확산 모델의 급격한 발전은 인간의 선호도를 모델에 반영하는 능력에 힘입은 바가 크다. 흔히 RLHF라 불리는 이 과정은 단순히 이미지를 생성하는 것을 넘어, 사용자가 보기에 미학적으로 뛰어나고 프롬프트에 정확하며 안전한 결과를 만들도록 AI를 가르치는 핵심 미세 조정 메커니즘이다. 그러나 우리가 시스템에 더 높은 수준의 성능을 요구함에 따라, 종종 상충하는 여러 목표를 동시에 최적화해야 하는 복잡한 공학적 과제에 직면하게 되었다.
지금까지 개발자들은 이를 해결하기 위해 여러 보상을 가중치 합산하는 방식을 사용해 왔다. 예를 들어 사진과 같은 사실성, 구도의 균형, 텍스트 정확성이라는 세 가지 목표를 한꺼번에 학습시키기 위해 각각 0.4와 0.6 같은 가중치를 부여해 하나의 점수로 합치는 방식이다. 하지만 이 방식은 보상 간의 경쟁이라는 결정적 결함이 존재한다. 사실성에 완벽한 결과물이 텍스트 정확성에는 무용지물일 수 있으며, 결국 가중치 합산은 양쪽 신호를 모두 희석시켜 어느 기준도 충족하지 못하는 결과를 초래하곤 했다.
새롭게 제안된 MARBLE 프레임워크는 이러한 보상 균형 문제를 근본적으로 변화시켰다. 연구진은 서로 경쟁하는 목표들을 하나의 합으로 억지로 묶는 대신, 경사 공간 최적화라는 방법을 도입했다. 이 기술은 각 보상에 대해 독립적인 이점 추정치를 유지함으로써, 모델이 자신의 어떤 행동이 특정 지표에 의해 평가받는지 명확히 이해하도록 돕는다. 정렬 과정을 가중 평균이 아닌 수학적 최적화 문제로 접근함으로써, 시스템은 모든 목표의 미묘한 차이를 존중하면서 경사를 조화롭게 조정할 수 있다.
이번 돌파구의 핵심은 Quadratic Programming을 활용해 모델 업데이트 방향을 결정한 점이다. 이를 통해 시스템은 복잡한 가중치 계획을 수동으로 조정할 필요 없이, 모든 보상 차원을 동시에 최적화하는 통합된 업데이트 방향을 찾아낸다. 실제로 이 방식은 한 단계씩 작업을 수행해야 했던 다단계 커리큘럼 학습의 필요성을 제거하며, 훨씬 더 효율적이고 간소화된 통합 훈련 프로세스를 가능하게 한다.
성능 향상도 매우 고무적이다. 연구진은 SD3.5 Medium 아키텍처 실험을 통해 MARBLE이 5가지 보상 차원을 모두 동시에 개선할 수 있음을 확인했다. 무엇보다 기존 훈련 배치에서 80% 이상 발생하던 보상 차원 간의 충돌, 즉 부정적 경사 문제를 효과적으로 해결했다. 수학적 복잡성이 더해졌음에도 불구하고 이 방식은 표준 단일 보상 기준 대비 약 97%의 속도로 작동하며, 고품질 정렬이 반드시 연산 효율의 희생을 수반하는 것은 아님을 증명했다.