Text-to-Distribution Prediction with Quantile Tokens and Neighbor Context

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 대규모 언어 모델 (LLM) 을 활용한 회귀 작업은 주로 단일 값 (점 추정) 을 예측하는 데 집중되어 왔습니다. 그러나 가격 예측, 수요 예측, 리스크 평가 등 실제 응용 분야에서는 불확실성을 고려한 전체 확률 분포를 이해하는 것이 필수적입니다.
• 기존 방법의 한계: 최근 연구 (Vedula et al., 2025) 는 공유된 최종 은닉 상태에 여러 개의 선형 헤드를 연결하여 다양한 분위수 (quantile) 를 예측하는 방식을 사용했습니다. 그러나 이 방식에는 세 가지 주요 문제가 있습니다.
  1. 간접적인 연결: 모든 분위수 예측이 동일한 표현 병목 (representation bottleneck) 에서 파생되어, 입력과 특정 분위수 출력 간의 직접적인 연결이 부족합니다.
  2. 국소적 근거 부재: 모델이 쿼리 텍스트만 기반으로 분포를 예측하므로, 유사한 사례들의 분포 정보를 활용하지 못해 분산과 꼬리 (tail) 행동을 정확히 포착하기 어렵습니다.
  3. 제한된 감독 신호: 기존 검색 기반 방법은 이웃 데이터에 단일 점 레이블만 제공하여, 전체 경험적 분포를 활용한 감독 학습이 부족합니다.

2. 제안된 방법론 (Methodology)

저자들은 두 가지 핵심 기법을 도입하여 Quantile Token Regression 아키텍처를 제안했습니다.

A. 분위수 토큰 회귀 (Quantile Token Regression)

• 개념: 입력 시퀀스의 끝에 학습 가능한 전용 분위수 토큰 (⟨Qτ1⟩, ..., ⟨QτQ⟩) 을 삽입합니다.
• 작동 원리: 각 분위수 토큰은 트랜스포머의 자기 주의 (self-attention) 메커니즘을 통해 입력 시퀀스의 다른 부분들과 직접 상호작용합니다.
• 장점:
  • 직접적인 입력 - 출력 경로: 공유된 은닉 상태를 거치지 않고, 각 분위수 토큰이 해당 분위수에 최적화된 정보를 직접 수집하여 예측합니다.
  • 해석 가능성: 각 토큰이 입력의 어떤 특징에 주의를 기울이는지 분석하여 예측 근거를 파악할 수 있습니다.
  • 일관성: 모든 분위수 토큰이 동일한 어텐션 계산 내에서 생성되므로, 분위수 간의 모순 (예: 90 백분위수가 80 백분위수보다 작음) 을 줄이고 더 일관된 분포를 생성합니다.

B. 검색 기반 분포 추정 (Retrieval-Augmented Distribution Estimation)

• 개념: 쿼리 텍스트와 의미적으로 유사한 이웃 (Neighbor) 인스턴스를 검색하여 입력에 포함시킵니다.
• 차별점: 기존 검색 기반 방법은 이웃의 단일 값만 제공했으나, 본 방법은 이웃의 전체 경험적 분포 (여러 개의 분위수) 를 함께 제공합니다.
• 효과: 유사한 사례들의 분포 증거를 통해 모델이 분산과 꼬리 행동을 더 잘 학습하도록 돕습니다.

C. 손실 함수 및 이론적 분석

• 손실 함수 비교: 다양한 손실 함수 (Pinball loss, Wasserstein distance 등) 를 이론적으로 분석했습니다.
• 결론: 경험적 분위수 (empirical quantiles) 를 감독 신호로 사용할 때, Wasserstein 거리 기반의 ℓ1 손실 함수가 가장 효과적입니다. 이는 목표 분위수에 대한 Fisher 일관성 (Fisher-consistency) 을 보장하며, Pinball loss 를 적용할 때 발생하는 편향 (bias) 과 분포 정보 손실을 방지합니다.

3. 실험 및 결과 (Results)

데이터셋: Inside Airbnb (부동산 가격 예측) 와 StackSample (Stack Overflow 답변 시간 예측) 데이터셋을 사용했습니다.
모델: Qwen3 (1.7B ~ 14B 파라미터) 모델을 파인튜닝하여 평가했습니다.

주요 성과:

1. 성능 향상: 제안된 방법 (Quantile Tokens + Retrieval) 은 모든 모델 크기와 데이터셋에서 베이스라인 (기존 공유 표현 방식) 을 일관되게 능가했습니다.
   • 정확도: 평균 MAPE(평균 절대 백분율 오차) 가 약 4 포인트 감소했습니다.
   • 예측 구간: 예측 구간 폭 (RCIW) 이 베이스라인 대비 2 배 이상 좁아졌습니다.
2. 작은/어려운 데이터셋에서의 효과: 데이터가 적고 불확실성이 높은 StackSample 데이터셋에서 효과가 특히 두드러졌습니다.
   • 기존 베이스라인은 매우 넓은 신뢰 구간 (RCIW@99: 4.55×10⁴) 을 보였으나, 제안 방법은 이를 346.90 으로 획기적으로 줄였습니다.
   • 평균 MAPE 는 266.65 에서 84.30 으로 68% 감소했습니다.
3. 검색 기반 증강의 효과: 검색된 이웃 (K=8) 을 추가하면 모델 크기가 작을수록 성능 향상이 큽니다. 이는 제한된 데이터 환경에서 유사 사례의 분포 정보가 예측을 grounding 하는 데 결정적인 역할을 함을 보여줍니다.
4. 손실 함수 검증: 이론적 분석과 일치하게, ℓ1 Wasserstein 손실을 사용한 모델이 가장 낮은 MAPE 와 가장 날카로운 (narrow) 예측 구간을 달성했습니다.

4. 주요 기여 및 의의 (Significance)

1. 새로운 아키텍처: LLM 기반 분포 예측을 위해 입력 시퀀스에 전용 분위수 토큰을 삽입하는 최초의 작업으로, 각 분위수별 직접적인 학습 경로를 제공합니다.
2. 분포 기반 검색 증강: 단순한 점 값이 아닌, 이웃의 전체 분포 정보를 활용하여 분포 예측의 정확도와 불확실성 추정을 개선했습니다.
3. 이론적 통찰: 경험적 분위수 감독 하에서 최적의 손실 함수가 무엇인지에 대한 이론적 분석을 제공하여, 향후 연구에 중요한 지침을 제시했습니다.
4. 실용적 가치: 가격 예측, 리스크 관리 등 불확실성이 중요한 실제 비즈니스 시나리오에서 LLM 의 활용 가능성을 크게 확장시켰습니다. 특히 데이터가 부족한 환경에서도 강력한 성능을 발휘하여 확장성 있는 솔루션을 제시합니다.

요약

이 논문은 LLM 이 텍스트를 입력받아 단순한 숫자가 아닌 정교한 확률 분포를 예측할 수 있도록 하는 혁신적인 프레임워크를 제시합니다. 분위수 토큰을 통한 직접적인 학습 경로와 이웃 분포 정보를 활용한 컨텍스트 증강, 그리고 Wasserstein 손실 함수의 적절한 적용을 통해, 기존 방법론보다 훨씬 정확하고 신뢰할 수 있는 분포 예측을 가능하게 했습니다.