인컨텍스트 러닝 (In-Context Learning)에서 태스크 벡터 (Task Vectors) 설계를 위한 기준으로서의 분포 정렬

인컨텍스트 러닝 (In-Context Learning, ICL)은 대규모 언어 모델 (Large Language Models, LLMs)이 예시 (demonstrations)를 통해 새로운 태스크에 적응할 수 있게 해주지만, 컨텍스트 길이 (context length)가 길어짐에 따라 추론 비용 (inference costs)이 급격히 증가하는 문제를 겪습니다. 태스크 벡터 (task vectors)는 예시를 압축된 은닉 상태 표현 (hidden-state representations)으로 압축함으로써 유망한 대안을 제시하지만, 그 품질은 오직 다운스트림 태스크 정확도 (downstream task accuracy)를 통해서만 평가되어 왔습니다. 이러한 간접적인 기준은 어떻게 하면 더 효과적인 태스크 벡터 추출 방법을 설계할 수 있는지에 대해 제한적인 통찰만을 제공합니다. 본 논문에서 우리는 태스크 벡터를 사용한 추론이 그 예측 분포 (predictive distribution)를 ICL의 분포와 정렬 (align)해야 한다고 가정합니다. 이를 정량화하기 위해, 우리는 태스크 벡터 기반 추론과 ICL 기반 추론 사이의 다음 토큰 확률 (next-token probabilities) 차이를 측정하는 지표인 $d_{\text{NTP}}$를 도입합니다. 우리의 실증적 분석 결과, $d_{\text{NTP}}$는 성능 프록시 (performance proxy) 역할을 하며 다운스트림 정확도와 강한 음의 상관관계 (negative correlation)를 보임을 확인했습니다. 이에 착안하여, 우리는 회귀 (regression)를 통해 예시의 효과를 추정하는 폐형 선형 매핑 (closed-form linear mapping)을 통해 $d_{\text{NTP}}$를 최소화하도록 설계된 방법론인 선형 태스크 벡터 (Linear Task Vector, LTV)를 개발했습니다. 8개의 분류 벤치마크 (classification benchmarks)와 5개의 LLM을 대상으로 실험한 결과, LTV는 추론 지연 시간 (inference latency)을 줄이면서 평균 정확도를 9.2% 향상시켜 기존의 태스크 벡터 베이스라인 (baselines)들을 일관되게 능가했습니다. 나아가 우리는 LTV가 회귀 태스크 (regression tasks)에서도 베이스라인보다 뛰어난 성능을 보임을 입증했습니다. 또한, 태스크 벡터 연구에서 아직 미미한 영역으로 남아 있는 부분인 서로 다른 모델 규모 (model scales) 간의 LTV 전이 가능성 (transferability)을 조사했습니다. 구체적으로, 우리는 더 큰 모델에서 추출한 태스크 벡터가 작은 모델의 성능을 6.4% 향상시킬 수 있음을 실증적으로 보여주었으며, 이는 추출된 태스크 표현 (task representations)의 새로운 유용성을 시사합니다.