주요 컨텐츠로 이동

TAO: 레이블이 없는 데이터로 효율적인 LLM을 훈련시키기 위해 테스트 시간 계산 사용

Test-time Adaptive Optimization (TAO)

Published: March 25, 2025

Mosaic 리서치2분 소요

Summary

  • TAO (Test-time Adaptive Optimization) 소개, 라벨이 없는 데이터를 필요로 하지 않고 작업에서 LLM 성능을 향상시키는 새로운 접근법, 테스트 시간 계산을 사용하여 모델 튜닝 과정을 보강하고 빠르고, 비용 효율적이며, 고품질의 LLM을 생성합니다.
  • 이미 가지고 있는 데이터, 기존의 예산 내에서의 추론 – TAO는 LLM 사용 데이터만 필요하지만, 수천 개의 레이블이 지정된 예제에서 전통적인 미세 조정을 뛰어넘을 수 있으며, 저렴한 오픈 소스 모델을 비싼 독점적인 모델보다 더 뛰어나게 만들 수 있습니다. 모델 품질은 튜닝 중에 더 많은 계산 예산을 투자함으로써 향상시킬 수 있습니다. 최종 모델의 추론 비용은 변하지 않습니다.
  • 테스트 시간 계산과 강화 학습에 의해 구동 – TAO는 이러한 기법을 확장하여 새로운 알고리즘과 새로운 기업 중심의 보상 모델, DBRM에 의해 구동되는 작업에서 모델을 자동으로 개선합니다.

대규모 언어 모델은 새로운 기업 작업에 적용하기 어렵습니다. 프롬프팅은 오류가 발생하기 쉽고 품질 향상이 제한적이며, 세부 조정은 대부분의 엔터프라이즈 작업에 사용할 수 없는 대량의 인간이 레이블링한 데이터를 필요로 합니다. 오늘, 우리는 레이블이 없는 사용 데이터만 필요로 하는 새로운 모델 튜닝 방법을 소개하고 있습니다. 이를 통해 기업들은 이미 가지고 있는 데이터만을 사용하여 AI의 품질과 비용을 개선할 수 있습니다. 우리의 방법론인 Test-time Adaptive Optimization (TAO)는 테스트 시간 계산 (o1과 R1에 의해 널리 알려짐)과 강화 학습 (RL)을 활용하여 모델이 과거 입력 예제만을 기반으로 작업을 더 잘 수행하도록 가르칩니다. 이는 인간의 라벨링 노력이 아닌 조정 가능한 튜닝 컴퓨팅 예산에 따라 확장됩니다. 중요한 점은, TAO가 테스트 시간 계산을 사용하긴 하지만, 이를 모델을 훈련시키는 과정의 일부로 사용합니다. 그 모델은 그 후 직접적으로 작업을 수행하며, 낮은 추론 비용(즉, 추론 시간에 추가 계산이 필요하지 않음)을 요구합니다. 놀랍게도, 라벨이 없는 데이터 없이도 TAO는 전통적인 미세 조정보다 더 나은 모델 품질을 달성할 수 있으며, 비용이 적은 오픈 소스 모델인 Llama를 비싼 독점 모델인 GPT-4o와 o3-mini와 같은 품질로 가져올 수 있습니다.

TAO는 우리 연구 팀의 데이터 인텔리전스 프로그램의 일부로, 기업이 이미 가지고 있는 데이터를 사용하여 AI가 특정 도메인에서 뛰어나게 하는 문제에 대한 것입니다. TAO를 통해 우리는 세 가지 흥미로운 결과를 얻었습니다:

  • 문서 질문 응답 및 SQL 생성과 같은 특수한 엔터프라이즈 작업에서 TAO는 수천 개의 레이블이 있는 예제에 대한 전통적인 세부 조정을 능가합니다. 이는 라벨이 필요 없는 상태에서 Llama 8B와 70B와 같은 효율적인 오픈 소스 모델을 GPT-4o와 o3-mini1 와 같은 비싼 모델과 비슷한 품질로 가져옵니다.
  • 또한 다중 작업 TAO를 사용하여 많은 작업에서 LLM을 넓게 개선할 수 있습니다. 라벨 없이 TAO를 사용하면, Llama 3.3 70B의 성능이 광범위한 기업 벤치마크에서 2.4% 향상됩니다.
  • 튜닝 시간에 TAO의 컴퓨팅 예산을 늘리면 동일한 데이터로 더 나은 모델 품질을 얻을 수 있으며, 튜닝된 모델의 추론 비용은 동일하게 유지됩니다.

그림 1은 TAO가 FinanceBench, DB Enterprise Arena, 그리고 BIRD-SQL (Databricks SQL 방언을 사용)²와 같은 세 가지 기업 작업에서 Llama 모델을 어떻게 개선하는지 보여줍니다. LLM 입력만 액세스할 수 있음에도 불구하고, TAO는 수천 개의 라벨이 있는 예제와 함께 전통적인 미세 조정(FT)을 능가하고 Llama를 비싼 독점 모델과 동일한 범위로 가져옵니다.

그림 1은 TAO가 FinanceBench, DB Enterprise Arena, 그리고 BIRD-SQL에서 Llama 모델을 어떻게 개선하는지 보여줍니다.

그림 1은 TAO가 FinanceBench, DB Enterprise Arena, 그리고 BIRD-SQL에서 Llama 모델을 어떻게 개선하는지 보여줍니다.

그림 1: 세 가지 기업 벤치마크에서 Llama 3.1 8B와 Llama 3.3 70B에 대한 TAO. TAO는 품질의 큰 개선을 가져다주며, 미세 조정을 능가하고 비싼 독점 LLM에 도전합니다.

TAO는 이제 Llama를 조정하고 싶어하는 Databricks 고객들에게 미리보기로 제공되며, 앞으로 출시될 여러 제품을 지원할 것입니다. 개인 미리보기의 일부로 작업에 시도하려는 관심을 표현하기 위해 이 양식 을 작성하세요. 이 포스트에서는 TAO가 어떻게 작동하는지와 그것을 사용한 결과에 대해 더 자세히 설명합니다.

TAO는 어떻게 작동합니까? 테스트 시간 컴퓨트와 강화 학습을 사용하여 모델 조정

사람이 주석을 단 출력 데이터를 필요로하는 대신, TAO의 핵심 아이디어는 테스트 시간 컴퓨트를 사용하여 모델이 작업에 대한 가능한 응답을 탐색하게 하고, 이러한 응답을 평가하여 LLM을 업데이트하는 데 강화 학습을 사용하는 것입니다. 이 파이프라인은 비싼 인간의 노력 대신 테스트 시간 컴퓨트를 사용하여 품질을 높일 수 있습니다. 더욱이, 이는 작업 특정 통찰력(예: 사용자 정의 규칙)을 사용하여 쉽게 맞춤화될 수 있습니다. 놀랍게도, 이런 스케일링을 고품질 오픈 소스 모델에 적용하면, 많은 경우에 인간 레이블보다 더 좋은 결과를 얻을 수 있습니다.

LIFT 파이프라인. LIFT는 추론 스케일링을 사용하여 작업에 대한 응답을 자동으로 생성하고 점수를 매기며, 노이즈가 있는 피드백을 기반으로 모델을 조정하는 방법을 배웁니다.
Figure 2: The TAO pipeline.

구체적으로, TAO는 네 단계로 구성됩니다:

  • 응답 생성: 이 단계는 작업에 대한 예제 입력 프롬프트 또는 쿼리를 수집하는 것으로 시작합니다. Databricks에서는 이러한 프롬프트를 AI Gateway를 사용하여 어떤 AI 애플리케이션에서든 자동으로 수집할 수 있습니다. 각 프롬프트는 다양한 후보 응답을 생성하는 데 사용됩니다. 여기에는 단순한 사고의 연결부터 정교한 추론 및 구조화된 프롬프팅 기법에 이르기까지 다양한 생성 전략이 적용될 수 있습니다.
  • 응답 점수 매기기: 이 단계에서는 생성된 응답이 체계적으로 평가됩니다. 점수 매기기 방법론에는 보상 모델링, 선호도 기반 점수 매기기, LLM 판사나 사용자 정의 규칙을 활용한 작업 특정 검증 등 다양한 전략이 포함됩니다. 이 단계는 각 생성된 응답이 품질과 기준과의 일치성에 대해 정량적으로 평가되도록 보장합니다.
  • 강화 학습(RL) 훈련: 마지막 단계에서는 RL 기반 접근법이 LLM을 업데이트하는 데 적용되어, 모델이 이전 단계에서 식별된 고점수 응답과 밀접하게 일치하는 출력을 생성하도록 안내합니다. 이 적응형 학습 과정을 통해, 모델은 예측을 개선하기 위해 예측을 세밀하게 조정합니다.
  • 지속적인 개선: TAO가 필요로 하는 유일한 데이터는 예제 LLM 입력입니다. 사용자는 LLM과 상호작용함으로써 이 데이터를 자연스럽게 생성합니다. LLM이 배포되자마자, 다음 라운드의 TAO를 위한 훈련 데이터를 생성하기 시작합니다. Databricks에서는, LLM이 더 많이 사용될수록 TAO 덕분에 더 좋아질 수 있습니다.

중요한 점은, TAO가 테스트 시간 계산을 사용하긴 하지만, 이를 훈련 하는 모델을 만드는 데 사용하며, 그 모델은 그 후 작업을 직접적으로 수행하며, 낮은 추론 비용을 요구합니다. 이는 TAO가 생성하는 모델들이 원래 모델과 동일한 추론 비용과 속도를 가지며, o1, o3 및 R1과 같은 테스트 시간 계산 모델보다 훨씬 적다는 것을 의미합니다. 우리의 결과가 보여주듯이, TAO로 훈련된 효율적인 오픈 소스 모델은 품질에서 선도적인 독점 모델에 도전할 수 있습니다.

TAO는 AI 모델을 조정하기 위한 도구 키트에 강력한 새로운 방법을 제공합니다. 느리고 오류가 발생하기 쉬운 프롬프트 엔지니어링과 비싸고 고품질의 인간 레이블을 생성하는 데 필요한 세부 조정과는 달리, TAO는 AI 엔지니어들이 그들의 작업의 대표적인 입력 예제를 단순히 제공함으로써 훌륭한 결과를 얻을 수 있게 합니다.

LLM 조정 방법의 비교.
Table 1: Comparison of LLM tuning methods.

TAO는 필요한 경우 맞춤화할 수 있는 매우 유연한 방법이지만, Databricks에서의 기본 구현은 다양한 엔터프라이즈 작업에서 상자 밖에서 잘 작동합니다. 우리의 구현의 핵심은 우리 팀이 개발한 새로운 강화 학습 및 보상 모델링 기법으로, 이를 통해 TAO는 탐색을 통해 학습하고, 그런 다음 RL을 사용하여 기본 모델을 조정할 수 있습니다. 예를 들어, TAO를 구동하는 요소 중 하나는 우리가 엔터프라이즈 작업을 위해 훈련시킨 사용자 정의 보상 모델인 DBRM으로, 이는 다양한 작업 범위에서 정확한 점수 신호를 생성할 수 있습니다.

TAO를 사용하여 작업 성능 개선

이 섹션에서는 우리가 TAO를 어떻게 사용하여 특수한 엔터프라이즈 작업에서 LLM을 조정하는지에 대해 더 깊이 들어가 보겠습니다. 우리는 인기 있는 오픈 소스 벤치마크와 우리가 도메인 인텔리전스 벤치마크 스위트(DIBS)의 일부로 개발한 내부 벤치마크를 포함한 세 가지 대표적인 벤치마크를 선택했습니다.

 표 2: 이 블로그에서 사용된 벤치마크 개요.
Table 2: Overview of benchmarks used in this blog.

각 작업에 대해 우리는 여러 접근법을 평가했습니다:

  • 오픈 소스 Llama 모델(Llama 3.1-8B 또는 Llama 3.3-70B)을 그대로 사용합니다.
  • Llama에 대한 세부 조정. 이를 위해 우리는 수천 개의 예제가 있는 크고 현실적인 입력-출력 데이터 세트를 사용하거나 생성했는데, 이는 일반적으로 미세 조정으로 좋은 성능을 달성하는 데 필요한 것입니다. 이에는 다음이 포함됩니다:
    • FinanceBench를 위한 SEC 문서에 대한 7200개의 합성 질문.
    • DB Enterprise Arena를 위한 4800개의 인간이 작성한 입력.
    • Databricks SQL 방언과 일치하도록 수정된 BIRD-SQL 훈련 세트에서 8137개의 예제.
  • 우리의 세부 조정 데이터셋에서 예제 입력만을 사용하고 출력은 사용하지 않고, 우리의 DBRM 엔터프라이즈 중심 보상 모델을 사용하여 Llama에 TAO를 적용합니다. DBRM 자체는 이러한 벤치마크에서 훈련되지 않습니다.
  • 고품질의 독점적인 LLMs - GPT 4o-mini, GPT 4o 및 o3-mini. 

표 3에 나타난 것처럼, 모든 세 가지 벤치마크와 두 가지 Llama 모델에 대해, TAO는 기본 Llama 성능을 크게 향상시키며, 심지어 세부 조정을 넘어섭니다. 

표 3: 세 가지 엔터프라이즈 벤치마크에서 Llama 3.1 8B와 Llama 3.3 70B에 대한 TAO.
Table 3: TAO on Llama 3.1 8B and Llama 3.3 70B across three enterprise benchmarks.

고전적인 테스트 시간 계산과 마찬가지로, TAO는 더 많은 계산을 할 수 있을 때 더 높은 품질의 결과를 생성합니다(예를 들어, 그림 3 참조). 그러나 테스트 시간 계산과는 달리, 이 추가 계산은 조정 단계 동안에만 사용되며, 최종 LLM은 원래 LLM과 동일한 추론 비용을 가집니다. 예를 들어, o3-mini는 우리의 작업에서 다른 모델보다 5-10배 더 많은 출력 토큰을 생성하여 추론 비용이 비례적으로 높아지지만, TAO는 원래의 Llama 모델과 동일한 추론 비용을 가집니다. 

그림 3: TAO는 조정 과정 중에 사용되는 테스트 시간 계산의 양에 따라 스케일링됩니다. 결과적으로 생성된 LLM을 사용하는 추론 비용은 원래의 LLM과 동일합니다.
Figure 3: TAO scales with the amount of test-time compute used during the tuning process. Inference cost to use the resulting LLM is the same as the original LLM.
TAO를 사용하여 다중 작업 인텔리전스 개선

지금까지, 우리는 TAO를 사용하여 SQL 생성과 같은 개별적인 좁은 작업에서 LLM을 개선하는 데 사용했습니다. 그러나, 에이전트가 더 복잡해짐에 따라, 기업들은 점점 더 많은 작업을 수행할 수 있는 LLM이 필요해집니다. 이 섹션에서는 TAO가 기업 작업 범위에 걸쳐 모델 성능을 넓게 개선할 수 있는 방법을 보여줍니다.

이 실험에서, 우리는 코딩, 수학, 질문 응답, 문서 이해, 채팅을 포함한 다양한 엔터프라이즈 작업을 반영하는 175,000개의 프롬프트를 모았습니다. 그런 다음 우리는 Llama 3.1 70B와 Llama 3.3 70B에 TAO를 실행했습니다. 마지막으로, 우리는 인기 있는 LLM 벤치마크(예를 들어, Arena Hard, LiveBench, GPQA Diamond, MMLU Pro, HumanEval, MATH) 및 엔터프라이즈에 관련된 여러 영역에서의 내부 벤치마크.

TAO는 두 모델의 성능을 의미있게 향상시킵니다[t][u]. Llama 3.3 70B와 Llama 3.1 70B는 각각 2.4 및 4.0 퍼센트 포인트 개선됩니다. TAO는 Llama 3.3 70B를 엔터프라이즈 작업에서 GPT-4o에 상당히 가깝게 만듭니다[v][w]. 이 모든 것은 인간 레이블링 비용 없이, 대표적인 LLM 사용 데이터와 우리의 생산 구현 TAO만을 사용하여 달성되었습니다. 품질은 코딩을 제외한 모든 하위 점수에서 개선되며, 성능은 고정되어 있습니다.

표 4: TAO를 사용한 다중 작업 엔터프라이즈 지능 개선
Table 4: Improving multitask enterprise intelligence using TAO

실제로 TAO 사용하기

TAO는 테스트 시간 계산을 활용하여 많은 작업에서 놀랍게도 잘 작동하는 강력한 조정 방법입니다. 자신의 작업에 성공적으로 사용하려면 다음이 필요합니다:

  • 충분한 예제 입력이 있는 작업 (수천 개), 배포된 AI 애플리케이션에서 수집 (예: 에이전트에게 보낸 질문)하거나 합성적으로 생성됩니다.
  • 충분히 정확한 점수화 방법: Databricks 고객의 경우, 여기서 강력한 도구 중 하나는 우리의 TAO 구현을 지원하는 사용자 정의 보상 모델인 DBRM이지만, 작업에 적용 가능한 사용자 정의 점수화 규칙이나 검증기로 DBRM을 보강할 수 있습니다.

TAO와 다른 모델 개선 방법을 가능하게 하는 최선의 실천 방법 중 하나는 AI 애플리케이션에 대한 데이터 플라이휠을 생성하는 것입니다. AI 애플리케이션을 배포하자마자, Databricks Inference Tables와 같은 서비스를 통해 입력, 모델 출력, 그리고 다른 이벤트를 수집할 수 있습니다. 그런 다음 입력만을 사용하여 TAO를 실행할 수 있습니다. 애플리케이션을 사용하는 사람이 많을수록, 그것을 조정하는 데 사용할 수 있는 데이터가 더 많아지고, - TAO 덕분에 - LLM이 더 좋아질 것입니다.

결론 및 Databricks에서 시작하기

이 블로그에서는 라벨이 없는 데이터를 필요로 하지 않고도 고품질 결과를 달성하는 새로운 모델 튜닝 기술인 Test-time Adaptive Optimization (TAO)를 소개했습니다. 우리는 TAO를 개발하여 기업 고객들이 직면한 주요 도전 과제를 해결하였습니다: 그들은 표준적인 미세 조정에 필요한 레이블이 지정된 데이터가 부족했습니다. TAO는 테스트 시간 계산과 강화 학습을 사용하여 기업들이 이미 가지고 있는 데이터, 예를 들어 입력 예제 등을 사용하여 모델을 개선하므로, 배포된 AI 애플리케이션의 품질을 향상시키고 더 작은 모델을 사용하여 비용을 줄이는 것이 간단해집니다. TAO는 특수한 AI 개발을 위한 테스트 시간 계산의 힘을 보여주는 매우 유연한 방법이며, 우리는 이것이 개발자들에게 강력하고 간단한 새로운 도구를 제공할 것이라고 믿습니다. 이는 프롬프팅과 미세 조정과 함께 사용됩니다.

Databricks 고객들은 이미 Llama에서 TAO를 사적인 미리보기로 사용하고 있습니다. 개인 미리보기의 일환으로 작업에 시도해 보고 싶다면 이 양식 을 작성하여 관심을 표현하세요. TAO는 또한 우리의 다가오는 AI 제품 업데이트와 출시에 많이 통합되고 있습니다 - 계속 주목해 주세요!


¹ 저자: Raj Ammanabrolu, Ashutosh Baheti, Jonathan Chang, Xing Chen, Ta-Chung Chi, Brian Chu, Brandon Cui, Erich Elsen, Jonathan Frankle, Ali Ghodsi, Pallavi Koppol, Sean Kulinski, Jonathan Li, Dipendra Misra, Jose Javier Gonzalez Ortiz, Sean Owen, Mihir Patel, Mansheej Paul, Cory Stephenson, Alex Trott, Ziyi Yang, Matei Zaharia, Andy Zhang, Ivan Zhou 

² 이 블로그에서는 o3-mini-medium을 사용합니다.
3 이것은 Databricks의 SQL 방언과 제품을 위해 수정된 BIRD-SQL 벤치마크입니다.

 

(이 글은 AI의 도움을 받아 번역되었습니다. 원문이 궁금하시다면 여기를 클릭해 주세요)

게시물을 놓치지 마세요

관심 있는 카테고리를 구독하고 최신 게시물을 받은편지함으로 받아보세요