🤖 AI

Learning What Matters: Probabilistic Task Selection via Mutual Information for Model Finetuning

이 논문은 상호 정보량과 행동 발산(behavioral divergences)을 사용하여 에너지 기반 마르코프 무작위장을 통해 태스크 간의 상호작용을 모델링함으로써, 대규모 언어 모델의 성능을 향상시키기 위해 태스크 커버리지와 중복성 사이의 균형을 맞추고 지도 미세 조정 데이터 혼합을 최적화하는 프레임워크인 TaskPGM을 소개한다.

원저자: Prateek Chanda, Saral Sureka, Parth Pratim Chatterjee, Krishnateja Killamsetty, Nikhil Shivakumar Nayak, Ganesh Ramakrishnan

게시일 2026-06-05

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Prateek Chanda, Saral Sureka, Parth Pratim Chatterjee, Krishnateja Killamsetty, Nikhil Shivakumar Nayak, Ganesh Ramakrishnan

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 매우 배고프고 매우 똑똑한 로봇(거대 언어 모델, LLM)을 위한 완벽한 "테이스팅 메뉴"를 만들려는 수석 셰프라고 상상해 보세요. 당신의 창고에는 수백 가지의 다양한 식재료(태스크)가 가득 차 있습니다. 이 식재료는 매콤한 추론 퍼즐부터 달콤한 시 쓰기, 짭짤한 코딩 챌린지에 이르기까지 다양합니다.

문제는, 모든 요리를 다 만들 만큼 시간이나 돈("학습 예산")이 충분하지 않다는 것입니다. 당신은 로봇에게 먹일 특정 재료의 조합을 골라내야 합니다.

과거의 방식: 추측과 무작위성
전통적으로 셰프들(AI 연구자들)은 다음과 같은 단순한 규칙을 사용하여 재료의 혼합 비율을 정했습니다:

"균등 배분" 규칙: 아주 작은 허브든 거대한 스테이크든 상관없이, 모든 식재료에 정확히 똑같은 양의 공간을 할당합니다.
"큰 냄비" 규칙: 가지고 있는 재료 중 가장 많은 것을 그냥 많이 사용합니다. 만약 밀가루 한 포대가 아주 크다면, 로봇에게 그것이 필요하지 않더라도 밀가루를 많이 사용합니다.

이 논문은 이러한 방법들이 결함이 있다고 주장합니다. 이 방법들은 로봇을 더 똑똑하게 만들 수 있는 희귀하고 특별한 풍미를 놓치거나, 맛이 똑같은 재료에 돈을 낭비하는 경우가 많습니다.

새로운 방식: TASKPGM ("풍미 지도" 접근법)
저자들은 TASKPGM이라는 새로운 시스템을 소개합니다. 이 시스템은 초스마트 풍미 지도처럼 작동합니다. 단순히 추측하는 대신, 수학적으로 완벽한 레시피를 찾아냅니다.

작동 원리는 다음과 같습니다 (쉬운 비유를 사용합니다):

1. 태스크의 "사회적 네트워크"

모든 태스크(예: "수학 문제 풀기" 또는 "시 쓰기")를 파티에 온 사람이라고 상상해 보세요.

단항 포텐셜 (Unary Potentials, "인기도" 점수): 어떤 사람들은 본래부터 인기가 많고 그 자체로 파티에 큰 가치를 가져다줍니다. 시스템은 이러한 태스크에 더 높은 기본 점수를 부여합니다.
쌍항 포텐셜 (Pairwise Potentials, "우정" 점수): 이것이 마법 같은 부분입니다. 시스템은 이 "사람들"이 어떻게 상호작용하는지 살펴봅니다.
- 만약 두 태스크가 **절친(Best Friends)**이라면 (즉, 로봇에게 정확히 똑같은 것을 가르친다면), 시스템은 이렇게 말합니다. "둘 다 초대하지 마세요! 공간 낭비입니다." 그리고 중복성에 대해 벌점을 부여합니다.
- 만약 두 태스크가 서로 보완하는 사이인 낯선 이들이라면 (예를 들어, 한 명은 논리를 가르치고 다른 한 명은 창의성을 가르친다면), 시스템은 이렇게 말합니다. "둘 다 초대하세요! 두 사람이 함께하면 파티가 더 좋아집니다." 그리고 다양성에 대해 보상을 줍니다.

2. 메뉴판을 읽지 않고도 "우정"을 측정하는 법

보통 사람들은 태스크의 제목이나 설명(예: "이것은 수학 태스크이다")을 보고 판단합니다. 하지만 TASKPGM은 더 똑똑합니다. 그것은 라벨에 신경 쓰지 않고 행동에 집중합니다.

시스템은 단 하나의 태스크만을 위해 아주 작고 임시적인 로봇을 훈련시킵니다. 그런 다음 다음과 같이 묻습니다. "만약 내가 이 로봇에게 수학 문제를 주면 어떻게 반응할까? 만약 내가 시 문제를 주면 어떻게 반응할까?"

만약 로봇이 두 태스크 모두에 유사하게 반응한다면, 그 태스크들은 "중복된 것"(두 사람이 똑같은 농담을 하는 것과 같음)입니다.
만약 로봇이 다르게 반응한다면, 그 태스크들은 "상호 보완적인 것"(한 명은 농담을 하고 다른 한 명은 이야기를 들려주는 것과 같음)입니다.

시스템은 이러한 반응을 정밀하게 측정하기 위해 수학(Jensen-Shannon Divergence 및 Pointwise Mutual Information)을 사용합니다.

3. 완벽한 레시Recipe

시스템은 누가 누구와 친구인지, 그리고 누가 고유한 가치를 가져오는지 지도를 그린 후, 복잡한 수학 퍼즐("에너지 최소화" 문제)을 해결합니다.

시스템은 완벽한 균형을 찾습니다. 즉, 중복되지 않으면서 가능한 한 다양한 "풍미"를 커버하는 혼합물을 찾아냅니다.
그리고 정확한 비율이 담긴 쇼핑 리스트를 출력합니다: "수학 태스크 15%, 시 쓰기 10%, 코딩 5% 등을 사용하세요" 등과 같이 말이죠.

4. 왜 더 나은가 (결과)

저자들은 두 가지 인기 있는 로봇 뇌(Qwen2-7B 및 Llama-2-7B)를 대상으로 서로 다른 양의 "음식"(25,000개 및 50,000개의 예시)을 사용하여 테스트했습니다.

결과: TASKPGM의 "완벽한 혼합물"을 먹은 로봇들은 "균등 배분"이나 "큰 냄비" 혼합물을 먹은 로봇들보다 어려운 테스트(추론 및 논리 퍼즐)에서 일관되게 더 좋은 성적을 거두었습니다.
효율성: 다른 고급 방법들은 개별 예시를 하나하나 고르려고 시도하며 엄청난 컴퓨팅 파워와 시간을 소모하지만, TASKPGM은 초기 "풍미 지도"를 그린 후 몇 초 만에 전체 혼합 비율을 계산해 냅니다.
해석 가능성: 이 시스템은 단순히 숫자만 주는 것이 아니라, 그 '이유'를 보여줍니다. 시스템은 어떤 태스크가 "끌개(Attractor, 폭넓고 유용한 태스크)"이고 어떤 것이 "전문가(Specialist, 틈새 태스크)"인지 보여주는 지도를 그릴 수 있으며, 이를 통해 인간이 로봇의 학습 과정을 이해하도록 돕습니다.

요약하자면
TASKPGM은 단순히 가진 재료가 얼마나 많은지에 따라 냄비에 재료를 던져 넣는 마스터 셰프가 아닙니다. 대신, 모든 재료의 맛을 보고, 그들이 서로 어떻게 상호작용하는지 이해하며, 로봇을 더 똑똑하고 빠르고 효율적으로 만들 수 있는 과학적으로 균형 잡힌 메뉴를 만들어내는 마스터 셰프입니다. 단 한 조각의 데이터도 낭비하지 않고 말이죠.

기술 요약: TaskPGM – 태스크 친화도를 통한 학습 태스크 혼합

문제 정의

대규모 언어 모델(LLM)을 위한 지도 미세 조정(SFT)은 이질적인 지시 태스크 집합에 대한 학습 예산의 분포에 크게 의존합니다. 현재의 관행은 균등 샘플링이나 데이터셋 크기에 비례하는 가중치 부여와 같은 단순한 휴리스틱에 의존하는 경우가 많습니다. 이러한 방법들은 태스크 간 상호작용을 고려하지 못하여 다음과 같은 두 가지 주요 비효전성을 초래합니다:

중복성(Redundancy): 거의 동일한 행동적 업데이트를 유도하는 태스크에 과도한 가중치를 부여하여 계산 예산을 낭비함.
전이 기회 상실(Missed Transfer): 보완적인 정보나 교차 도메인 전이 이점을 제공하는 태스크를 충분히 활용하지 못함.

기존의 자동화된 접근 방식(예: LESS와 같은 영향력 기반 검색 또는 SMART와 같은 부가 함수적 선택)은 연속적인 혼합 비율보다는 이산적인 부분 집합을 최적화하거나, 모델의 다운스트림 동작을 무시하는 의미론적 유사성에 의려하거나, 비용이 많이 드는 반복 학습을 요구합니다. 핵심 과제는 고정된 예산 하에서 다운스트림 성능을 극대화하기 위해, 동시에 대표성(높은 효용)과 다양성(낮은 중복성)을 갖춘 연속적인 태스크 혼합을 자동으로 결정하는 것입니다.

방법론: TaskPGM

저자들은 TaskPGM(확률적 그래프 모델링을 통한 태스크 혼합)을 제안하며, 이는 SFT 데이터 혼합을 조밀한 마르코프 무작위장(MRF) 상의 에너지 최소화 문제로 정식화합니다.

1. 확률적 그래프 모델 정식화

태스크들은 MRF 내의 노드로 모델링됩니다. 목표는 에너지 함수 $E(p)$ 를 최대화하는 최적의 샘플링 분포 $p^* \in \Delta_n$ (여기서 $\Delta_n$ 은 확률 심플렉스)를 찾는 것입니다. 에너지 함수는 두 가지 구성 요소로 이루어집니다:

단항 포텐셜 ( $\Psi_{un}$ ): 개별 태스크의 효용 또는 대표성을 인코딩합니다. 이는 전체 컬렉션에 대한 해당 태스크의 총 유사성 질량으로부터 유도됩니다.
쌍방 포텐셜 ( $\Psi_{pair}$ ): 중복성과 다양성을 인코딩합니다. 이 항들은 혼합물 내에서 매우 유사한 태스크들이 공존하는 것을 억제합니다.

최적화 목적 함수는 다음과 같이 정식화됩니다:
$\max_{p \in \Delta_n} E(p) = \Psi_{un}^\top p - \frac{1}{2} p^\top \Psi_{pair} p$
이는 특정 조건 하에서 볼록 이차 함수(convex quadratic function)를 심플렉스 제약 조건 하에서 최소화하는 것과 같습니다.

2. 행동적 태스크 친화도 지표

TaskPGM의 핵심 혁신은 쌍방 태스크 친화도( $S_{ij}$ )를 계산하는 방법입니다. 의미론적 유사성(임베딩)이나 메타데이터에 의존하는 대신, TaskPGM은 단일 태스크로 미세 조정된 모델의 예측 분포로부터 도출된 **행동적 발산(behavioral divergences)**을 기반으로 친화도를 계산합니다.

점별 상호 정보량 (Pointwise Mutual Information, PMI): 각자의 데이터셋에 대해 평가될 때 두 예측 분포 사이의 상호 정보를 측정합니다.
젠슨-샤논 발산 (Jensen–Shannon Divergence, JSD): 교차 태스크 예시에 대한 두 모델의 예측 분포 간 불일치를 나타내는 대칭적 척도입니다.

이러한 지표들은 태스크가 모델의 행동을 어떻게 형성하는지를 포착하며, 이는 태스크 외(out-of-task) 예시에 대한 태스크 간 전이 및 중복성을 직접적으로 다룹니다.

3. 최적화 및 솔버

연속적 혼합: 저자들은 카루슈-쿤-터커(KKT) 조건을 사용하여 연속적 혼합 가중치에 대한 **폐쇄형 최적해(closed-form optimum)**를 도출합니다.
- 만약 쌍방 유사도 행렬이 양의 준정부호(PSD)가 아닌 경우, 이들은 스펙트럼 시프팅(대각 성분에 상수를 더함)을 적용하여 볼록성을 보장하며, 이는 비-PSD 구조가 심할 경우 해를 균등 분포 쪽으로 편향시키는 정규화 항을 도입합니다.
- 역행렬이 불량하거나 비음수 제약 조건이 위반되는 경우, 심플렉스 상에서 **투영 경사 하강법(Projected Gradient Descent, PGD)**을 사용합니다.
이산/예산 기반 변형: 특정 부분 집합(top- $k$ )을 선택해야 하는 시나리오를 위해, 저자들은 이 연속적 최적화를 이산 집합 함수로 격상시킵니다. 그들은 이 함수가 표준적인 정칙성 조건(제한적 강볼록성 및 매끄러움) 하에서 **단조적(monotonic)**이고 **약한 부가 함수적(weakly submodular)**임을 증명하여, 이론적 보증을 갖춘 효율적인 그리디(greedy) 구성을 가능하게 합니다.

주요 기여

태스크 혼합 설계를 위한 확률적 프레임워크: SFT 데이터 혼합을 대표성(단항 항)과 다양성(쌍방 항)을 명시적으로 절충하는 MRF 상의 에너지 최소화로 정식화한 새로운 프레임워크를 제시했습니다.
행동 기반 친화도 지표: 정적인 의미론적 유사성을 넘어 기능적 태스크 의존성을 포착하기 위해, 단일 태스크 미세 조정 모델의 예측 분포 발산(JSD/PMI)을 사용하여 태스크 관계를 계산하는 방법을 제안했습니다.
이론적 근거를 갖춘 최적화:
- KKT 조건을 통한 연속적 혼합을 위한 폐쇄형 솔루션 도출.
- 그리디 선택에 이론적 보증을 지원하는 이산 예산 변형에 대한 약한 부가 함수성 확립.
실험적 검증: Qwen2-7B, Llama-2-7B 모델과 25K 및 50K 인스턴스 예산에 대해 309개의 지시 태스크를 대상으로 수행한 실험에서 휴리스틱 및 최신 상태(SOTA) 선택 베이스라인 대비 일관된 성능 향상을 입증했습니다.

실험 결과

프레임워크는 Qwen2-7B 및 Llama-2-7B 모델을 사용하여 25K 및 50K 인스턴스 예산으로 309개의 지시 태스크에 대해 평가되었습니다.

성능 향상: TaskPGM(PMI 및 JSD 변형 모두 사용)은 균등 샘플링, 크기 비례(EPM) 및 랜덤 베이스라인보다 일관되게 우수한 성능을 보였습니다.
- Qwen2-7B (25K 예산)에서, PMI 변형은 복잡한 추론 벤치마크에서 휴리스틱 베이스라인보다 성능을 거의 3% 향상시켰습니다.
- Llama-2-7B에서, PMI 유도 혼합물은 BBH에서 Uniform 및 EPM보다 2%에서 4% 이상 높은 성능을 보였습니다.
SOTA와의 비교: TaskPGM은 LESS(강력한 영향력 함수 기반 베이스라인)와 대등하거나 더 우수한 성능을 유지했습니다. 특히, LESS는 각 벤치마크를 위해 별도의 타겟 지향적 모델을 구축해야 하는 반면, TaskPGM은 단일 범용 체크포인트를 통해 이를 달성합니다.
효율성: TaskPGM은 LESS보다 계산 효율성이 높습니다. 유사도 계산은 일회성으로 분할 상환되는 비용(~16–34 GPU-hours)이며, 이후 모든 예산에 대한 혼합물은 폐쇄형 솔루션을 통해 몇 초 만에 도출될 수 있습니다. 반면, LESS는 각 타겟에 대한 그래디언트 데이터스토어를 구축하기 위해 약 320 GPU-hours가 필요합니다.
확장성: 예산을 25K에서 50K로 늘렸을 때, 휴리스틱 방법(Uniform, EPM)은 파편화되고 비강건한 동작을 보인 반면, TaskPGM은 일관된 이득을 얻었습니다.

의의 및 주장

본 논문은 TaskPGM이 원칙적이고, 해석 가능하며, 효율적인 SFT 데이터 혼합 접근 방식을 제공한다고 주장합니다.

해석 가능성: 학습된 가중치와 유사도 히트맵은 태스크 클러스터, 중복성, 그리고 "끌개(attractor)" 태스크를 드러내어, 왜 특정 소스가 도움이 되는지에 대한 통찰을 제공합니다.
일반화: 특정 타겟을 위해 최적화되는 방식과 달리, TaskPGM은 광범위한 다운스트림 태스크(추론, 언어 이해, 지식) 전반에서 성능을 개선하는 단일 혼합물을 생성합니다.
강건성: 성능 향상은 유리한 샘플링 분산에 의한 것이 아니라, 모델 행동에 기반한 기능적 태스크 선택에 의해 주도됨이 다양한 베이스 모델과 유사도 지표를 통해 입증되었습니다.

저자들은 원칙적이고 유사도 기반의 혼합이 확장 가능한 지시 튜닝에 필수적이며, 종종 휴리스틱 혼합보다 적은 데이터로도 그와 일치하거나 이를 능가한다고 결론짓습니다. 저자들은 고정된 태스크 풀과 친화도 추정의 계산 비용에 관한 한계를 인정하며, 향-확장 가능한 근사법을 향후 연구 방향으로 제시합니다.