TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

기술 요약: TAPIOCA – 작업 인식 가지치기가 OOD 모델 능력을 향상시키는 이유

문제 제기
최근 연구는 추론 시 대규모 언어 모델 (LLM) 의 계층을 제거하는 것, 특히 TALE(작업 인식 계층 제거) 과 같은 작업 인식 가지치기 방법을 통해 재학습 없이 특정 작업의 성능을 향상시킬 수 있음을 보여주었습니다. 그러나 이러한 향상을 이끄는 메커니즘은 여전히 불분명합니다. 핵심적인 질문은 이러한 개선이 일반적인 압축의 이점, 정규화 효과, 또는 모델의 내부 표현과 입력 분포 간의 특정 상호작용에서 비롯되는지 여부입니다. 또한 이러한 이득이 보편적인 것인지 아니면 특정 유형의 데이터 시프트에 국한된 것인지도 알려지지 않았습니다.

방법론
저자들은 통제된 합성 실험과 파인튜닝된 LLM 에 대한 대규모 평가를 결합하여 이 현상을 조사했습니다.

1. 통제된 회귀 설정: 연구는 컨텍스트 내 선형 회귀에 처음부터 훈련된 12 층 트랜스포머로 시작합니다. 훈련 분포는 명시적으로 정의됩니다 (계수가 $U(-1, 1)$에서 샘플링됨). 저자들은 분포 내 (ID) 데이터 (동일한 분포) 와 분포 외 (OOD) 데이터 (계수가 더 넓은 구간 $U(-\sigma, \sigma)$에서 샘플링됨, 여기서 $\sigma > 1$) 에 대해 모델을 평가합니다. 그들은 특정 분포에서의 검증 정확도를 최대화하기 위해 계층을 제거하는 탐욕적 가지치기 알고리즘인 TALE 를 적용합니다.
2. 대규모 LLM 평가: 이 발견들은 수학 추론과 코드 생성에 특화된 파인튜닝된 LLM(Llama 3.1 8B 및 GPT-OSS 120B) 으로 확장됩니다. 여기서 "적응된" 분포는 파인튜닝 작업 (예: MATH500) 이며, OOD 작업에는 MMLU 및 BoolQ 와 같은 일반 벤치마크가 포함됩니다.
3. 기하학적 및 인과적 분석: 가지치기가 작동하는 이유를 설명하기 위해 저자들은 숨겨진 상태의 기하학을 분석합니다. 그들은 계층별 표현 노름과 쌍별 토큰 거리를 측정합니다. 또한 계층 변환을 근사하기 위해 국소 선형 대리 모델을 사용하고, 기하학적 왜곡의 효과를 분리하기 위해 역 대리 맵을 삽입하고 잔차 업데이트를 재조정하는 것과 같은 인과적 개입을 수행합니다.

주요 발견

• OOD 특이적 이득: 주요 경험적 발견은 작업 인식 가지치기가 분포 시프트 하에서만 성능 향상을 가져온다는 것입니다. 모델이 적응된 분포 (ID) 로 평가될 때, TALE 는 성능을 향상시키지 못하거나 최적의 정밀도에 필요한 계층을 제거함으로써 성능을 능동적으로 저하시킵니다. 반면, OOD 데이터에서는 TALE 가 일관되게 정확도를 향상시키는 계층을 식별하고 제거합니다.
• 표현 기하학 불일치: 저자들은 OOD 입력이 ID 입력에 비해 왜곡된 표현 기하학을 유도함을 보여줍니다. 구체적으로, OOD 입력은 계층별 숨겨진 상태 노름과 쌍별 거리가 ID 데이터에서 관찰된 프로필과 체계적으로 벗어나게 만듭니다. 회귀 설정에서 OOD 입력은 계층 전반에 걸쳐 거리를 팽창시키는 반면, LLM 에서는 OOD 궤적이 중간에서 후기 계층에서 ID 궤적과 이격됩니다.
• 증폭 메커니즘: 이 연구는 특정 계층이 "분포 의존적 증폭기"로 작용함을 식별합니다. ID 입력에서 이러한 계층은 유용한 정교화를 수행합니다. 그러나 OOD 입력에서는 동일한 계층이 고이득, 저랭크 증폭을 나타내어 표현을 작업 적응 기하학에서 더 멀리 밀어냅니다. TALE 는 이러한 특정 계층을 제거함으로써 기하학적 왜곡의 증폭을 방지하여 성능을 향상시킵니다.
• 수축이 아닌 정렬: 가지치기는 단순히 활성화 노름을 억제하는 것이 아닙니다. 분포 시프트의 방향이 반전된 실험 (넓은 분포에서 훈련하고 좁은 분포에서 테스트) 에서 가지치기는 ID 프로필과 정렬되도록 표현 거리를 확장시킵니다. 핵심 메커니즘은 수축이든 확장이든 간에 적응 과정에서 학습된 기하학으로 OOD 표현을 재정렬하는 것입니다.
• 인과적 증거: 인과적 개입은 기하학적 왜곡이 성능 손실의 원인임을 확인합니다. 가지치기된 계층 뒤에 역 대리 맵을 삽입하여 (계층의 기하학적 효과를 취소) OOD 성능을 가지치기된 모델 수준으로 거의 복원합니다. 마찬가지로, 선택된 계층의 잔차 업데이트를 재조정하면 계층의 기여도가 감쇠됨에 따라 정확도가 지속적으로 향상되어 위상적 인공물이나 단순한 중복성이 유일한 원인이 아님을 배제합니다.

의의 및 주장
이 논문은 추론 시 가지치기의 성공에 대한 새로운 기하학적 설명을 제공한다고 주장합니다. 계층의 중요성은 고정된 것이 아니라 작업과 입력 분포 모두에 조건부라는 점을 논증합니다. 저자들은 사전 훈련된 모델이 각각 지역적 표현 기하학과 연관된 여러 작업 특화 기능을 인코딩한다고 가정합니다. ID 입력은 이 기하학에 부합하지만 OOD 입력은 이를 위반합니다.

이 연구의 의의는 다음과 같습니다:

1. 가지치기의 한계 명확화: 작업 인식 가지치기가 보편적인 성능 부스터가 아니라 분포 시프트를 위한 교정 메커니즘임을 확립합니다.
2. 기하학적 해석: 계층 중복성에 대한 이해를 정적 파라미터 수에서 동적이고 분포 의존적인 기하학적 왜곡으로 전환합니다.
3. 메커니즘적 통찰: 성능 향상이 모델 용량을 증가시키거나 노이즈를 줄이는 것이 아니라, 모델의 적응된 기하학에 맞춰 내부 표현의 궤적을 수정함으로써 발생한다는 인과적 증거를 제공합니다.

저자들은 가지치기가 입력 분포와 모델이 학습한 작업 기하학 간의 불일치를 증폭시키는 계층을 감쇠시킴으로써 OOD 성능을 향상시킨다고 결론지으며, 이는 효과적으로 모델의 내부 상태를 최적화된 조건과 재정렬한다고 주장합니다.