기술 요약: SRA (대규모 언어 모델 증류를 위한 스펜 표현 정렬)

문제 정의

본 논문은 대형 "교사" 언어 모델과 소형 "학생" 모델이 서로 다른 토크나이저와 어휘를 사용하는 크로스-토크나이저 지식 증류 (CTKD) 의 과제를 다룹니다. 기존 CTKD 방법들은 주로 토큰 수준 또는 로그이트 (logit) 수준의 정렬 전략 (예: 편집 거리, 최적 수송, 또는 프로젝션 레이어 사용) 에 의존합니다. 저자들은 이러한 접근 방식이 종종 취약하며, 서로 다른 어휘 체계 간에 의미적으로 대응되지 않을 수 있는 이산적이고 원자적인 단위를 정렬하려고 시도하기 때문에 토크나이저 간의 차이점에 민감하다고 주장합니다. 또한, 현재 방법들은 토크나이저가 다를 때 트랜스포머 아키텍처 내의 특징 진화의 근본적인 역학을 포착하는 데 실패하는 경우가 많습니다.

방법론: SRA 프레임워크

제안된 SRA 프레임워크는 정렬의 기본 단위를 개별 토큰에서 강건하고 토크나이저에 구애받지 않는 스펜 (span) 으로 변경함으로써 CTKD 를 재정의합니다. 이 접근법은 트랜스포머를 레이어를 통해 진화하는 상호작용하는 입자 (토큰) 의 시스템으로 보는 다중 입자 동역학 시스템 (MPDS) 의 물리적 유추에 기반합니다.

1. 스펜 구성 및 매핑

• 스펜 정의: 토큰을 직접 정렬하는 대신, SRA 는 토큰들을 원시 입력의 동일한 텍스트 단위에 해당하는 "스펜"으로 그룹화합니다.
• LCS 를 통한 정렬: 교사 및 학생 모델 간의 스펜을 일치시키기 위해, 저자들은 두 토크나이저가 생성한 문자 오프셋의 최장 공통 부분 수열 (Longest Common Subsequence, LCS) 을 계산합니다. 이는 두 모델에서 정렬된 스펜의 시작 및 끝 인덱스를 정의하는 입력 텍스트의 일치된 세그먼트를 식별하며, 직접적인 토큰 대 토큰 매핑을 요구하지 않습니다.

2. 질량 중심 (CoM) 을 통한 스펜 표현

MPDS 관점에서 각 스펜은 입자 군집으로 간주됩니다. 이 군집의 상태는 질량 중심 (Center of Mass, CoM) 으로 표현됩니다:

• 가중 풀링: 스펜의 표현은 구성 토큰의 숨겨진 상태 (hidden states) 에 대한 어텐션 가중 평균으로 계산됩니다.
• 어텐션 가중치: 토큰에 대한 가중치 $w_i$는 모든 어텐션 헤드를 통해 시퀀스의 마지막 토큰으로부터 받은 어텐션 점수에서 유도됩니다. 이는 물리학에서 "무거운" (의미적으로 더 중요한) 입자가 군집의 중심을 지배한다는 개념을 모방합니다.
• 공식: $C_i = \sum w_t H_t$, 여기서 $H_t$는 숨겨진 상태이고 $w_t$는 어텐션에서 유도된 가중치입니다.

3. 증류 목적 함수

SRA 는 스펜 수준에서 지식을 전달하기 위해 결합된 손실 함수를 사용합니다:

• 스펜 수준 숨겨진 상태 전달 ($L_{HS}^{Span}$):

  • 코사인 유사도 손실을 사용하여 학생 및 교사 스펜의 CoM 표현을 정렬합니다.
  • 기하학적 정규화자 ($L_{Geo}$): 표현 공간의 구조적 무결성을 보존하기 위해, 교사 모델의 스펜 쌍 간의 상대적 거리 (코사인 거리) 와 일치하도록 학생 모델의 스펜 쌍 간의 상대적 거리를 유지하는 정규화자가 추가됩니다.
  • 가중 손실: 손실은 스펜 내 토큰 가중치의 합에서 유도된 스펜의 의미적 중요도에 따라 가중되며, 이를 통해 학생이 가장 정보량이 많은 영역에 집중하도록 보장합니다.

• 정렬된 스펜 로그이트 증류 ($L_{KD}^{Span}$):

  • 어휘 불일치 ($V_{tea} \neq V_{stu}$) 를 해결하기 위해 지식 전달을 공유 부분 공간 $\tilde{V} = V_{tea} \cap V_{stu}$로 제한합니다.
  • 공유 어휘 내의 토큰 로그이트에서 스펜 수준 로그이트를 집계하고, 이 공유 공간에 대한 교사 및 학생의 분포 간에 쿨백 - 라이블러 (KL) 발산을 계산합니다.

• 전체 목적 함수: 최종 손실은 실제 레이블에 대한 표준 교차 엔트로피 손실과 숨겨진 상태 및 로그이트의 스펜 수준 증류 손실을 하이퍼파라미터 $\alpha$로 균형을 맞추어 결합합니다.

주요 기여

1. 새로운 개념적 프레임워크: 본 논문은 트랜스포머의 다중 입자 동역학 시스템 관점에 기반한 CTKD 프레임워크를 소개하며, 스펜을 입자 군집으로 모델링하고 안정적인 가법적 표현을 위해 그들의 질량 중심을 활용합니다.
2. 스펜 수준 증류: 숨겨진 표현과 로그이트를 모두 포함하는 스펜 수준에서 지식을 증류하는 방법을 제안합니다. 여기에는 중요한 스펜을 우선시하기 위한 어텐션에서 유도된 가중치 체계와 표현 공간의 상대적 기하학을 보존하기 위한 기하학적 정규화자가 포함됩니다.
3. 실증적 검증: 저자들은 SRA 가 여러 교사 - 학생 쌍 및 벤치마크에서 최첨단 CTKD 베이스라인 (ULD, DSKD, MinED, MultilevelOT 포함) 보다 현저히 우수한 성능을 달성함을 입증했습니다.

실험 결과

저자들은 지시 따르기 데이터셋 (Dolly, Vicuna, SelfInst, S-NI, DialogSum) 을 사용하여 다양한 크로스-아키텍처 및 크로스-토크나이저 쌍 (예: Qwen1.5/2.5, Mistral-7B, GPT-2, OPT, TinyLLaMA) 에서 SRA 를 평가했습니다.

• 성능: SRA 는 강력한 베이스라인과 비교하여 모든 설정에서 일관되게 가장 높은 ROUGE-L 점수와 최고의 평균 성능을 달성했습니다.
• 의미적 품질: GPT-4o-mini 를 판정자로 사용한 의미 평가는 SRA 가 의미적으로 더 충실하고 문맥적으로 정렬된 응답을 생성함을 확인했습니다.
• 애블레이션 연구:
  • 기하학적 정규화자나 가중 풀링 메커니즘을 제거하면 성능이 저하되어 구조적 관계의 보존과 중요한 스펜에 대한 초점의 필요성이 확인되었습니다.
  • 이 방법은 동일한 토크나이저 시나리오에서도 견고성을 보이며, 거리 기반 및 기타 CTKD 방법보다 우수한 성능을 발휘했습니다.
• 효율성: SRA 는 스펜 매칭 및 집계로 인해 일부 베이스라인보다 약간 더 많은 GPU 메모리를 필요로 하지만, 비교된 방법들 중 가장 낮은 평균 단계당 훈련 시간을 달성했습니다.

중요성 및 주장

본 논문은 SRA 가 토큰 수준 정렬의 취약성을 극복하는 강건하고 물리적으로 기반을 둔 지식 증류 접근법을 제공한다고 주장합니다. 정렬 단위를 스펜으로 변경하고 질량 중심 개념을 활용함으로써, SRA 는 동적인 레이어 간 특징 진화의 전달을 자연스럽게 지원합니다. 저자들은 그들의 방법이 최종 출력뿐만 아니라 시스템 기하학적 진화의 역학에 초점을 맞춰야 한다는 원칙을 검증하며, 이질적인 토크나이저 환경에서 효율적인 모델을 배포하기 위한 확장 가능한 솔루션을 제공한다고 주장합니다.

이 작업은 더 안정적이고 포괄적인 지식 전달을 향한 한 걸음으로 제시되지만, 저자들은 현재 로그이트 매핑의 정적 성질과 스펜 정렬에 필요한 온더플라이 (on-the-fly) 교사 추론의 계산 비용과 같은 한계를 인정합니다.