SPR-RAFT: Parameter-Efficient Regression-Aware Fine-Tuning for Biomedical LLM Regression

Motivation

• Known: LLM은 PubMed 등 생의학 코퍼스로 사전학습되어 풍부한 도메인 지식을 encoding하며, 텍스트로 직렬화된 이질적 근거를 입력받아 예측을 수행하는 범용 인터페이스로 활용될 수 있음이 알려져 있다. 또한 LoRA, Prefix-Tuning, P-Tuning 같은 PEFT 기법과 RAFT 같은 regression-aware fine-tuning 기법이 각각 파라미터 효율성과 회귀 정합성 측면에서 개별적으로 연구되어 왔다.
• Gap: LLM은 discrete tokenization과 cross-entropy objective로 인해 수치적 근접성(numerical proximity)을 인식하지 못하며, RAFT류 regression-aware 기법은 full-parameter fine-tuning에 의존해 계산 비용이 크고 catastrophic forgetting에 취약한 반면, LoRA/P-Tuning 같은 일반 PEFT 기법은 파라미터 효율적이지만 연속적 수치 타깃의 metric proximity를 다루는 메커니즘이 없다는 gap이 존재한다.
• Why: 임상시험 기간, 생물학적 나이, 분자 특성 예측 등 생의학 회귀 과제는 이질적이고 비정형적인 근거로부터 연속값을 정확히 예측해야 하는데, 데이터가 희소하고 형식이 불균일한 생의학 도메인에서 계산 비용이 낮으면서도 수치적으로 신뢰할 수 있는 LLM 기반 회귀기를 만드는 것은 실용적으로 매우 중요하다.
• Approach: frozen LLM에 정보 라우팅을 담당하는 학습 가능한 soft prompt와 값 추출을 담당하는 [REG] 앵커 토큰 기반 경량 회귀 헤드를 분리하여 배치하고, distribution 기반 텍스트 생성 손실과 representation 기반 강건 회귀 손실을 결합한 hybrid objective로 두 모듈을 공동 학습한다.

Achievement

Figure 3. Effectiveness of uncertainty-aware fusion (diagnostic). The dashed line marks the high-uncertainty region of t

1. 파라미터 효율성: 약 1.6M개의 학습 파라미터(4B 백본의 약 0.04%)만으로 backbone을 freeze한 채 고정밀 회귀를 수행함.
2. 일관된 성능 우위: 임상시험 기간 예측(TrialBench), 생물학적 나이 추정(NHANES), 생의학 의미 유사도(BIOSSES), 분자 특성 예측(MoleculeNet) 등 다양한 벤치마크에서 prompting 전략, 표준 fine-tuning, non-LLM 베이스라인을 모두 능가함.
3. 불확실성 인식 융합: distribution 기반 branch와 representation 기반 branch의 예측을 precision-weighted로 결합하여 수치적으로 보정된(calibrated) 최종 예측을 산출함(Figure 3에서 효과 입증).

How

Figure 1. Architecture of SPR-RAFT. SPR-RAFT adapts a frozen LLM for biomedical regression by decoupling information rou

• 입력 x(직렬화된 생의학 근거)에 대해 [THINK] 자유 서술 rationale과 [REG] 앵커 토큰, 그리고 결정론적으로 y에 매핑되는 표준 숫자 문자열을 생성하도록 정식화(response = [THINK]... [REG] str(y)).
• 학습 가능한 soft prompt P를 입력 앞에 prepend하여 frozen backbone f_θ의 attention을 조절, 흩어진 근거를 앵커 hidden state h_REG로 압축.
• distribution 기반 branch는 후보 문자열 확률로부터 (μ_text, σ²_text)를 추정하고 Huber loss로 학습.
• representation 기반 branch는 h_REG를 경량 헤드 g_ϕ에 통과시켜 heteroscedastic Gaussian NLL로 학습.
• 두 branch의 예측을 precision-weighted 방식으로 융합하여 최종 예측 ŷ(x)를 산출하며, gradient는 (P, ϕ)만 업데이트하고 θ는 고정.

Originality

• soft prompt tuning과 [REG] 전용 readout 토큰 기반 회귀 헤드를 결합한 최초의 프레임워크로, 계산 비용이 큰 RAFT류 regression-aware fine-tuning과 회귀 목표에 무관한(regression-agnostic) PEFT 기법(LoRA, P-Tuning) 사이의 간극을 메움.
• 정보 라우팅(soft prompt)과 값 추출(regression head)을 명시적으로 분리하는 dual-module 설계.
• distribution 기반 Huber loss와 representation 기반 heteroscedastic Gaussian NLL을 결합한 dual-branch hybrid objective 및 uncertainty-aware precision-weighted fusion 메커니즘 제안.

Limitation & Further Study

• 평가된 백본이 Qwen3, Gemma3, Llama-3.2 등 특정 모델군에 국한되어 있어 더 크거나 다른 아키텍처 계열에서의 일반화 가능성이 명확히 검증되지 않음.
• soft prompt 길이, [REG] 토큰 위치, Huber/Gaussian NLL 가중치 등 하이퍼파라미터에 대한 민감도 분석이 본문 발췌만으로는 충분히 드러나지 않음.
• 생의학 도메인 외 일반 회귀 과제나 극단적으로 분포가 치우친(long-tail) 타깃에 대한 강건성 검증이 추가로 필요함.
• 후속 연구로 멀티태스크 회귀나 few-shot/zero-shot 전이 성능에 대한 심층 분석이 요구됨.

Evaluation

총평: soft prompt와 전용 readout 토큰 기반 회귀 헤드를 결합한 아이디어가 참신하며, 극히 적은 파라미터로 다양한 생의학 회귀 벤치마크에서 일관된 성능 향상을 보여 실용적 가치가 높은 연구이다.