Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

기술 요약: LLM 보정 미세 조정을 위한 기능 수준 불확실성 정량화

문제 정의
정확한 불확실성 정량화 (UQ) 는 의료 및 과학적 추론과 같은 안전이 중요한 분야에서 대규모 언어 모델 (LLM) 의 신뢰할 수 있는 배포에 필수적입니다. 그러나 미세 조정된 LLM 은 데이터가 제한된 작업에 적응할 때 종종 과도한 확신을 보이며, 출력의 실제 신뢰도를 반영하지 못합니다. 저랭크 적응 (LoRA) 과 같은 파라미터 효율적 미세 조정 (PEFT) 을 위한 기존 UQ 방법들은 대부분 사후 (post hoc) 적입니다. 이들은 모델이 훈련된 후 파라미터 공간 (예: 앙상블 또는 베이지안 근사를 통해) 을 분석하거나 프롬프트를 교란시켜 불확실성을 추정합니다. 이러한 접근법들은 과도한 확신의 근본 원인인 적응 과정 자체 중 미세 조정으로 인한 불확실성과 일반화 격차를 모델이 포착하지 못하는 문제를 해결하지 못합니다. 또한, 여러 순전달 (forward pass) 이나 광범위한 샘플링에 의존하는 방법들은 상당한 계산 오버헤드를 초래하여 확장성을 제한합니다.

방법론: UQ4CT
저자들은 보정 미세 조정을 위한 기능 수준 불확실성 정량화 (UQ4CT) 를 제안합니다. 이는 파라미터 공간에서 기능 공간 불확실성으로 초점을 이동시키는 프레임워크입니다. 훈련 후 학습된 파라미터를 분석하는 대신, UQ4CT 는 미세 조정 과정 중 에 프롬프트 의존적 LoRA 전문가들의 혼합으로 유도된 기능 공간에 걸쳐 불확실성을 보정합니다.

이 방법론은 세 가지 핵심 구성 요소로 이루어집니다:

1. MixLoRA 를 통한 계층적 기능 분해:
   모든 가능한 전문가 조합을 계층 전체에 걸쳐 모델링하는 데 따른 조합적 폭발을 피하기 위해, UQ4CT 는 계층적 분해를 채택합니다. 이는 각 계층에서 희소 라우터가 입력 은닉 상태에 기반하여 LoRA 전문가들의 부분집합을 동적으로 선택하는 전문가 혼합 (MoE) 아키텍처를 활용합니다. 이는 입력에 조건부로 다양한 기저 함수 (LoRA 전문가) 의 가중 합으로 모델 출력을 구성하는 유연한 기능 공간을 구축합니다. 기능 공간은 계층별 혼합의 재귀적 구성으로 표현되어 모델이 계산 경로를 적응적으로 구성할 수 있게 합니다.

2. 기능 수준 신뢰도 (FLC) 추정:
   모델은 전문가 선택의 집중도를 분석하여 불확실성을 정량화합니다. 구체적으로, 기능 수준 신뢰도 (FLC) 는 각 계층에서 라우터가 선택한 상위 $k$개 전문가의 라우팅 가중치에서 유도됩니다. 라우팅 가중치의 높은 집중도 (즉, 라우터가 특정 전문가를 선택하는 데 확신이 있음) 는 높은 기능적 신뢰도에 해당하며, 분산된 가중치는 더 높은 불확실성을 나타냅니다. 이는 여러 샘플링 순전달 없이 모델의 확신에 대한 원칙적이고 내부적인 신호를 제공합니다.

3. 보정 손실 ($L_{cal}$):
   핵심 혁신은 훈련 중 기능 수준 신뢰도를 예측 정확도와 정렬시키는 새로운 보정 손실입니다. 이 손실은 다음과 같이 정의됩니다:
   $$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$$
   여기서 $1\{\cdot\}$는 올바른 예측에 대한 지시 함수입니다. 이 목적 함수는 FLC 가 모델 예측이 정확할 확률을 근사하도록 강제합니다. 표준 교차 엔트로피 및 로드 밸런싱 손실과 함께 이 손실을 최소화함으로써, 모델은 정확할 가능성이 높을 때만 높은 신뢰도를 할당하고 그렇지 않을 때는 더 높은 불확실성을 할당하도록 학습합니다. 이는 라우터가 불확실성이 높을 때 대체 전문가 혼합을 탐색하도록 장려하여 일반화를 개선합니다.

주요 기여

• 기능 수준 관점: 이 논문은 파라미터 공간에서 기능 공간 불확실성 정량화로의 전환을 도입하여, 미세 조정 중에 실현할 수 있는 입력 - 출력 매핑에 대한 모델의 신뢰도를 명시적으로 보정합니다.
• 보정 인식 미세 조정: 프롬프트 의존적 LoRA 혼합 가중치를 예측 정확도와 동적으로 정렬하는 새로운 보정 손실이 제안되어, 정확도를 저하시키지 않으면서 과도한 확신을 완화합니다.
• 효율적인 아키텍처: 이 방법은 최소한의 계산 오버헤드로 다양한 기능적 관계를 포착하는 계층적 MixLoRA 아키텍처를 활용하며, 추론 시 단일 순전달만 필요합니다.
• 이론적 보장: 저자들은 데이터 분포에 걸쳐 보정 위험을 최적화하면 예측 정확도의 실제 확률에 해당하는 최적의 FLC 를 산출한다는 증명을 제공합니다.

실험 결과
저자들은 Llama-3.1-8B 와 Mistral-7B 를 사용하여 네 가지 객관식 벤치마크 (ARC-E, ARC-C, OBQA, ClimateQA) 와 두 가지 개방형 생성 QA 작업 (TruthfulQA, TriviaQA) 에서 UQ4CT 를 평가했습니다.

• 분포 내 성능: UQ4CT 는 강력한 베이스라인 (LoRA, MC Dropout, Deep Ensembles, Bayesian LoRA 포함) 대비 벤치마크 전반에 걸쳐 기대 보정 오차 (ECE) 를 25% 이상 감소시켰으며, 경쟁력 있거나 우수한 정확도를 유지했습니다. 특히 UQ4CT 는 단일 순전달로 이러한 결과를 달성한 반면, BLoB (N=10) 와 같은 방법들은 열 번의 순전달이 필요함에도 불구하고 더 높은 ECE 를 보였습니다.
• 분포 이동: 분포 이동 하에서 (CS, 공학, 법률, 건강 분야의 MMLU 하위 작업에서 테스트), UQ4CT 는 베이스라인에 비해 분포 외 입력에 대한 향상된 견고성을 보여주며 우수한 보정과 경쟁력 있는 정확도를 유지했습니다.
• 생성 작업: 개방형 생성에서 UQ4CT 는 TruthfulQA 와 TriviaQA 에서 가장 낮은 ECE 를 달성하여, 기능 공간 보정 관점이 분류 스타일 설정을 넘어 확장됨을 시사했습니다.
• 효율성: 메모리 분석에 따르면 UQ4CT 의 피크 GPU 메모리 사용량은 표준 LoRA 보다 약 7.5% 높을 뿐이며, 대규모 배포를 위해 계산적으로 실행 가능합니다.

의의 및 주장
이 논문은 UQ4CT 가 현재 미세 조정 관행의 근본적인 한계인 적응 과정 중 불확실성 인식의 부재를 해결한다고 주장합니다. 기능 수준 신호를 통해 불확실성 보정을 미세 조정 목적 함수에 직접 통합함으로써, 이 방법은 모델의 신뢰도가 실제 예측 신뢰성을 더 잘 반영하도록 보장합니다. 저자들은 이 작업을 데이터가 부족한 환경에서 과도한 확신이 심각한 위험을 초래하는 신뢰할 수 있는 미세 조정된 LLM 을 향한 기초적인 단계로 위치시킵니다. 그들은 F1 기반 프록시를 사용하여 개방형 생성으로 확장되지만, 장형 추론이나 다단계 생성으로의 확장은 향후 연구 방향이라고 겸손하게 언급합니다.