LLM은 자신의 실수를 얼마나 잘 수정할까? Keras와 TPU를 이용한 챗봇 아레나 실험

TL;DR

Hugging Face는 LLM이 캘린더 도구를 위한 API 호출을 생성하고, 짧은 영어 피드백을 바탕으로 실수를 수정해야 하는 챗봇‑아레나 실험을 진행했습니다. Gemma 2 9B가 오류를 안정적으로 수정하며 가장 우수한 성능을 보였고, Llama 3.1 8B가 그 뒤를 이었으나 추가 프롬프트가 필요했습니다. 더 작은 모델(1-2B)과 Vicuna 같은 구형 모델은 종종 실패하거나 깨진 출력을 생성했습니다.

실험 설정

테스트에는 두 개의 함수를 가진 Python API가 사용되었습니다: action.add_calendar_entry(title, date="YYYY-MM-DD", time="HH:MM", duration=m) 및 action.remove_calendar_entry(title, date, time). 시스템 프롬프트는 모델이 유용한 음성 비서 역할을 수행하고 실행 가능한 코드 한 줄로 응답하도록 지시했습니다. 대화는 연도가 누락된 부정확한 요청(“Add a meeting with Fred on Nov 11 at 5PM”)으로 시작되어, 모델이 실수를 저지르게 유도한 뒤 “The current year is 2024.”라는 후속 문구로 이를 수정할 수 있도록 했습니다. 이후의 턴에서는 이벤트를 추가하고, 기간을 변경하고, 항목을 취소하는 등 실수 수정의 기회를 여러 번 제공했습니다.

기술적 구현: TPU, JAX, Keras 및 모델 샤딩

아레나는 Gradio를 사용하여 Hugging Face Spaces에 구축되었습니다. 이 시스템은 코어당 8개의 코어와 16GB의 RAM(총 128GB)을 제공하는 TPU v5e 2x4에서 실행되었습니다. 이 메모리 덕분에 저자들은 모든 코어에 모델을 샤딩(sharding)하여 여러 모델을 동시에 로드할 수 있었습니다. Keras(현재 JAX에서 네이티브로 실행됨)를 사용하여 bfloat16 형식으로 최대 5개의 ~8B 파라미터 모델과 3개의 ~2B 파라미터 모델을 로드했으며, 총 7개의 LLM이 메모리에 상주했습니다. 모델 병렬성은 각 모델을 8개의 TPU 코어에 샤딩하기 위해 Keras의 내장 레이아웃 맵(예: keras_hub.models.Llama3Backbone.get_layout_map(device_mesh))에 의존했습니다. 저자들은 안정적인 로딩을 위해 디버깅과 간헐적인 레이아웃 맵 조정이 필요했다고 언급했습니다.

평가된 모델

실험은 10B 파라미터 미만의 인스트럭션 튜닝(instruction-tuned) 모델에 집중되었습니다. 이는 모델들이 TPU 메모리에 함께 들어갈 수 있고 작업이 충분히 간단했기 때문입니다. 테스트된 제품군은 Gemma, Llama 3, Mistral, Vicuna, CodeGemma였습니다. 구체적인 체크포인트는 다음과 같습니다:

  • Gemma 2 9B-instr
  • Llama 3.1 8B-instr
  • Llama 3.2 3B-instr
  • Llama 3.2 1B-instr
  • Gemma 2B-instr
  • CodeGemma 7B-instr
  • Vicuna 1.5 7B-instr
  • Mistral 7B-instr 모든 모델은 게시물에 나열된 Hugging Face Hub URL을 통해 액세스되었습니다.

첫 번째 질문에 대한 신뢰도

각 모델에게 초기 프롬프트를 다섯 번씩 요청했습니다. 결과(✓ = 올바른 API 호출, 🍄 = 대부분 맞지만 실수가 있음, 🔥 = 쓰레기/인식 가능한 호출 없음):

  • Gemma 2 9B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.1 8B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 3B-instr: ✓ ✓ ✓ ✓ ✓
  • Llama 3.2 1B-instr: 🔥 🍄 🔥 🔥 🔥
  • Gemma 2B-instr: 🍄 🍄 🍄 🍄 ✓
  • CodeGemma 7B-instr: ✓ ✓ ✓ ✓ ✓
  • Vicuna 1.5 7B-instr: ✓ 🔥 🔥 ✓ 🔥
  • Mistral 7B-instr: ✓ ✓ 🍄 ✓ 🍄 1-2B 모델과 구형 Vicuna 모델만 지속적으로 실패했으며, 더 큰 모델들은 매번 올바르게 답변했습니다.

전체 대화 – 실수 수정

전체 6턴 대화를 실행했을 때, 저자들은 각 실수 후에 피드백을 주었고 모델이 이를 수정할 수 있는지 기록했습니다(🥦 = 수정 성공). 주요 결과는 다음과 같습니다:

  • Gemma 2 9B-instr 및 Llama 3.1 8B-instr는 한두 번의 사소한 실수만 저지르며 대화를 완료했습니다. Llama는 브로콜리(성공)를 얻기 위해 추가적인 "fix it" 프롬프트가 필요했습니다.
  • 온라인 Gemini 실행(훨씬 더 큰 모델)은 API 호출을 출력하기 위해 특별한 프롬프트가 필요했으며 여전히 여러 실수를 저질렀습니다. 이는 크기만으로는 이 작업에서 더 나은 성능을 보장하지 않음을 보여줍니다.
  • 작은 모델 중에서는 Gemma 2B-instr만이 대화를 마칠 수 있었으나, 수다스러운 텍스트를 추가하거나(“Sure, here’s the updated code…”) 날짜/시간을 혼동하는 경향이 반복되었습니다. 하지만 요청을 받으면 실수를 수정할 수 있었습니다.
  • Vicuna 1.5 7B-instr는 남은 오류를 수정하기 전에 종종 반복되거나 무의미한 출력(🔥🔥)으로 퇴보했습니다.
  • Mistral 7B-instr는 많은 실수를 저질렀지만 수정할 수 있었으며, 수많은 교정 턴 끝에 6개의 브로콜리 기호를 획득했습니다.
  • CodeGemma 7B-instr는 고질적인 반복 실수를 보였습니다. 연도에 공백을 삽입(“20 24”)했으며 교정 후에도 이를 제거할 수 없었습니다.

추가 실수 수정 테스트

수정 능력을 분리하기 위해, 저자들은 다른 모델이 생성한 오류가 있는 출력을 각 모델에 입력하고 모델이 이를 복구할 수 있는지 확인했습니다. 결과(🥦 = 올바른 수정, 🍄 = 오류 지속, 🔥 = 다수의 오류):

  • Gemma 2 9B-instr: 잘못된 시간과 "API only" 프롬프트를 완벽하게 수정했습니다. "Wrong API" 프롬프트에는 어려움을 겪으며 올바른 호출과 함께 사과 문구를 출력했습니다(이는 "API call only please"라고 요청하여 제거할 수 있는 반복적인 실수입니다).
  • Llama 3.1 8B-instr: Gemma와 유사하게 대부분의 오류를 수정했지만, 잘못된 API 호출을 안정적으로 수정하는 데 어려움을 겪었습니다.
  • CodeGemma 7B-instr: Gemma와 유사하게 사과 패턴을 유지하며 행동했습니다.
  • Mistral 7B-instr: 시간 오류는 잘 수정했지만, API-only 및 wrong-API 프롬프트에는 일관성을 보이지 않았습니다.
  • 모든 작은 모델(Llama 3.2 3B, Llama 3.2 1B, Gemma 2B)과 Vicuna는 많은 🔥 또는 🍄 결과를 생성하여, 제공된 실수를 안정적으로 복구할 수 없음을 나타냈습니다.

요약 및 시사점

저자들은 단순한 두 개의 API 설정이 모든 모델에게 쉬울 것으로 예상했으나, 결과는 명확한 차이를 보여주었습니다. Gemma 2 9B-instr는 최소한의 수정 프롬프트로 거의 완벽한 정확도로 전체 대화를 수행한 유일한 모델이었습니다. Llama 3.1 8B-instr가 약간 더 많은 가이드가 필요하며 근소한 차이로 2위를 차지했습니다. 모델 크기가 도움이 되긴 했지만, 훨씬 더 큰 Gemini 모델은 포맷팅 문제로 인해 성능이 저하되었으며, 이는 지시 이행 능력(instruction fidelity)과 토큰 수준의 제어가 원시 규모(raw scale)보다 더 중요하다는 것을 나타냅니다. 작은 모델(≤2B)과 구형 모델은 유효한 API 호출을 생성하는 데 자주 실패했으며, 생성하더라도 광범위한 프롬프트 없이는 교정 피드백을 통합하는 데 어려움을 겪었습니다. 이 실험은 사용자가 평문 영어로 실수를 수정할 수 있는 단순한 도구 사용 시나리오에서, 현재의 중간 크기 인스트럭션 튜닝 LLM이 유용할 수 있지만 신뢰성은 여전히 모델 규모, 학습 데이터 및 불필요한 텍스트를 억제하는 능력에 달려 있음을 보여줍니다.

독자들은 공개적으로 사용 가능한 Keras Chatbot Arena Space를 사용하여 이 연구를 재현하거나 확장할 수 있으며, 실수 수정 동작을 개선하기 위한 미세 조정(fine-tuning) 실험을 권장합니다.

Sources