Improving LLM First-Token Predictions in Multiple-Choice Question Answering via Output Prefilling

1. 문제 제기 (Problem)

대형 언어 모델 (LLM) 을 객관식 질문 답변 (MCQA) 태스크에서 평가할 때, 현재 가장 널리 사용되는 방법은 첫 번째 토큰 확률 (First-Token Probability, FTP) 입니다. 이 방법은 모델이 생성하는 첫 번째 토큰이 정답 옵션 (예: 'A', 'B', 'C', 'D') 중 하나일 확률을 계산하여 정답을 선택합니다.

그러나 FTP 는 다음과 같은 두 가지 주요 취약점으로 인해 신뢰성이 낮아질 수 있습니다:

1. 첫 번째 토큰 불일치 (First-Token Misalignment): 모델이 기대된 심볼릭 형식 (단순한 옵션 문자) 대신 완전한 문장 (예: "I believe the answer is A") 이나 모호한 문자열을 생성하여, 첫 번째 토큰이 유효한 옵션이 아닌 경우가 발생합니다.
2. 첫 번째 토큰 오해석 (First-Token Misinterpretation): 첫 번째 토큰이 유효한 옵션 문자 (예: 'A') 라 하더라도, 그것이 정답을 의미하는 것이 아니라 문법적 역할을 하는 경우입니다 (예: "A possible answer could be C"라는 문장에서 'A'가 첫 번째 토큰이지만 정답은 'C'인 경우).

이러한 문제들은 평가의 정확도, 보정 (calibration), 그리고 모델 간 비교의 공정성을 해칩니다.

2. 방법론 (Methodology)

저자들은 AI 안전 분야에서 '공격 (attack)' 기법으로 사용되던 출력 프리필링 (Output Prefilling) 을 MCQA 평가의 신뢰성을 높이는 해결책으로 재해석하여 제안했습니다.

• 핵심 아이디어: 모델의 응답 시작 부분에 구조화된 자연어 접두사 (prefix) 를 삽입하여 모델이 유효한 옵션을 첫 번째 토큰으로 생성하도록 유도합니다.
• 구현 방식:
  • 지시형 (Instruction-tuned) LLM 의 대화 형식 (예: <|assistant|>) 직후에 고정된 프레이즈를 주입합니다.
  • 예시 템플릿: "Given the question and the possible options, my answer is: "
  • 이 접두사는 모델이 생성할 첫 번째 토큰을 심볼릭 옵션 ('A', 'B' 등) 으로 강제적으로 유도하며, 모델의 파라미터를 수정하거나 추가적인 디코딩 (full decoding) 을 필요로 하지 않습니다.
• 평가 지표:
  • 정확도 (Accuracy): FTP 와 프리필링을 적용한 FTP 의 성능 비교.
  • 보정 (Calibration): Adaptive Calibration Error (ACE), Brier Score, Log Loss 등을 통해 모델의 확신도 (confidence) 와 실제 정확도의 일치도를 측정.
  • 유효성 및 일관성: 전체 어휘부 (full vocabulary) 에서 첫 번째 토큰이 유효한 옵션일 확률 (FTVR) 과 이어서 나오는 두 번째 토큰의 다양성 (Continuation Diversity, CD) 분석.

3. 주요 기여 (Key Contributions)

1. 새로운 평가 기법 제안: AI 안전 공격 기법이었던 '프리필링'을 MCQA 평가의 신뢰성 향상을 위한 제로-비용 (zero-cost) 전략으로 재정의했습니다.
2. FTP 의 한계 극복: 첫 번째 토큰 불일치와 오해석 문제를 기계적으로 해결하여, 외부 분류기나 오픈 엔디드 생성 (open-ended generation) 없이도 고품질의 평가를 가능하게 합니다.
3. 광범위한 검증: Llama-3.1, Qwen, Gemma, Mistral, Phi 등 다양한 아키텍처와 크기의 최신 오픈소스 LLM 과 MMLU, RACE, ARC 등 14 개 이상의 다양한 MCQA 벤치마크를 통해 방법론의 효과성을 입증했습니다.
4. 보정성 개선: 단순히 정확도만 높이는 것이 아니라, 모델의 확신도 (confidence) 가 실제 정확도를 더 잘 반영하도록 보정 (calibration) 성능을 크게 향상시켰습니다.

4. 실험 결과 (Results)

• 정확도 향상: 프리필링을 적용한 FTP 는 표준 FTP 보다 모든 모델과 벤치마크에서 일관되게 높은 정확도를 보였습니다. 특히 Gemma-2-9B는 평균 25.9% 포인트, Gemma-7B는 14.9% 포인트, Phi-4-14B는 7.3% 포인트의 큰 향상을 기록했습니다.
• 강력한 베이스라인 대비 우위:
  • 프롬프트 엔지니어링 (입력 측 지시) 만으로는 일관된 향상을 보이지 못했으나, 출력 측 프리필링은 훨씬 강력하고 견고한 성능 향상을 가져왔습니다.
  • 오픈 엔디드 생성 + 외부 분류기 (GPT-3.5, Llama-3-70B, xFinder) 를 사용한 평가와 비교했을 때, 프리필링을 적용한 FTP 는 유사하거나 때로는 더 높은 정확도를 보여주면서도 계산 비용이 훨씬 낮습니다.
• 보정성 (Calibration) 개선: ACE, Brier Score, Log Loss 모든 지표에서 프리필링 적용 후 값이 감소하여, 모델이 자신의 예측에 대해 더 현실적인 확신을 갖게 됨을 확인했습니다.
• 형식 일관성: 프리필링은 모델이 첫 번째 토큰 이후에도 일관된 심볼릭 형식을 유지하도록 도와주어, 문법적 오해석 오류를 크게 줄였습니다.

5. 의의 및 결론 (Significance)

이 연구는 LLM 의 객관식 평가에서 단순하고 효율적인 출력 프리필링 기법이 모델의 신뢰성을 획기적으로 높일 수 있음을 증명했습니다.

• 실용성: 추가적인 학습 (fine-tuning) 이나 복잡한 외부 모델 없이, 프롬프트 구조만 변경하여 즉시 적용 가능한 방법입니다.
• 신뢰성: 기존 FTP 방식의 취약점을 해결하여, LLM 의 추론 능력을 더 정확하게 측정할 수 있는 표준적인 평가 프로토콜을 제시합니다.
• 미래 방향: 이는 LLM 평가의 신뢰성을 높이는 동시에, 모델의 출력 형식을 제어하는 새로운 패러다임을 제시하며, 향후 다양한 태스크와 언어로 확장될 수 있는 기초를 마련했습니다.

요약하자면, 이 논문은 **"출력 프리필링"**이라는 간단한 기법을 통해 LLM 의 객관식 평가 정확도와 신뢰성을 획기적으로 개선할 수 있음을 체계적으로 입증한 중요한 연구입니다.