기술 요약: 분할 트리를 활용한 토큰화 (ToaST)

문제 제기

토큰화는 자연어 처리의 기초적인 단계이지만, 기존 방법들은 종종 압축이나 하류 작업 성능을 최적화하지 못할 수 있는 탐욕적 휴리스틱에 의존합니다. 바이트 페어 인코딩 (BPE) 은 탐욕적 병합 전략을 통해 토큰 수를 최소화하지만, 이는 휴리스틱 근사치에 불과합니다. 이전 연구들은 제한된 알파벳에서도 압축을 극대화하는 토크나이저를 찾는 문제가 NP-완전 문제임을 입증했습니다. furthermore, 압축과 하류 작업 성능 간의 관계는 여전히 논쟁의 여지가 있습니다. 일부 연구는 상관관계를 시사하지만, 다른 연구들은 압축만으로는 토큰화의 품질을 설명하기에 부족하다고 주장합니다. UnigramLM 이나 PathPiece 와 같은 기존 대안들은 다른 추론 절차 (예: 제거식 상향식 또는 최단 경로 추론) 를 활용하지만, 여전히 전역 최적화와 계산적 확장성 측면에서 어려움을 겪고 있습니다.

방법론

저자들은 **분할 트리를 활용한 토큰화 (ToaST)**를 소개합니다. 이는 후보 토큰의 구성과 최종 어휘의 선택을 분리하는 서브워드 토큰화 방법입니다. 이 접근법은 세 가지 주요 단계로 구성됩니다:

1. 분할 트리 구성

BPE 가 하향식으로 토큰을 구축하거나 UnigramLM 이 상향식으로 확률을 학습하는 것과 달리, ToaST 는 훈련 코퍼스 내의 모든 고유한 프리토큰 (pretoken) 에 대해 어휘와 무관하게 완전 이진 트리를 구성합니다.

• 전처리: 정규 표현식을 사용하여 훈련 데이터를 프리토큰으로 분할합니다.
• 분할 논리: 주어진 프리토큰에 대해 알고리즘은 이를 두 개의 비어 있지 않은 부분으로 재귀적으로 분할합니다. 분할 지점은 코퍼스에서 사전 계산된 n-그램 통계에 기반하여 결과적으로 생성되는 왼쪽 및 오른쪽 바이트 n-그램의 최소 카운트를 최대화하도록 결정적으로 선택됩니다 ($\max \min(c_p, c_{p'})$).
• 구조: 이 과정은 잎 노드가 단일 바이트가 될 때까지 계속됩니다. 이러한 트리의 노드들은 모든 후보 토큰의 집합을 형성합니다. 이 구성은 최종 어휘와 무관하며 반복적인 훈련이 필요하지 않습니다.

2. 재귀적 추론

어휘 $V$가 주어지면, 토큰화는 분할 트리의 재귀적 하강을 통해 수행됩니다:

• 루트 (전체 프리토큰) 에서 시작하여 알고리즘이 현재 노드가 $V$에 있는지 확인합니다.
• 노드가 $V$에 있으면, 이를 토큰으로 방출합니다.
• 그렇지 않으면, 알고리즘은 두 자식 노드를 재귀적으로 처리합니다.
• 특성: 이 절차는 모든 필요한 단일 바이트를 포함하는 어휘라면 유효한 토큰화를 보장합니다. 어휘에서 토큰을 제거하는 것은 단순히 추론이 트리 내에서 더 깊이 하강하도록 강제할 뿐, 병합 기반 방법과 달리 다른 토큰에 연쇄적인 영향을 미치지 않습니다.

3. 정수 프로그래밍을 통한 어휘 선택

핵심 혁신은 훈련 코퍼스 전체의 총 토큰 수를 최소화하기 위해 어휘 선택을 **정수 계획법 (Integer Program, IP)**으로 공식화하는 것입니다.

• 변수: 이진 변수 $x_i$는 후보 토큰 $i$가 어휘에 포함되는지 여부를 나타냅니다. 이진 변수 $z_{jk}$는 분할 트리 $j$의 노드 $k$가 추론 중에 방출되는지 여부를 나타냅니다.
• 목적 함수: 방출된 토큰의 총 가중 카운트를 최소화합니다: $\min \sum_{j} c_j \sum_{k} z_{jk}$.
• 제약 조건:
  • 어휘 크기: 선택된 토큰의 합은 목표 크기 $m$과 같아야 합니다.
  • 커버리지: 잎에서 루트까지의 모든 경로에 대해 정확히 하나의 노드가 선택되어야 합니다 (전체 프리토큰이 커버되도록 보장).
  • 일관성: 노드가 방출될 수 있으려면 ($z_{jk}=1$) 어휘에 포함되어야 하며 ($x_i=1$), 그 조상 노드 중 어느 것도 포함되어서는 안 됩니다.
• 해법: IP 는 **선형 계획법 (Linear Program, LP)**으로 완화됩니다. 실제로 LP 완화는 "거의 정수성 (near-integral)"을 가지며, 이는 해에 소수 값이 매우 적게 포함된다는 것을 의미합니다. 간단한 반올림 휴리스틱을 적용하여 전역 최적점에 근접한 (전역 최적값의 무시할 수 있는 분수 이내) 정수 해를 얻습니다.

주요 기여

1. 새로운 토큰화 프레임워크: ToaST 는 BPE 의 탐욕적 근사치와 UnigramLM 의 특정 추론 제약을 피하는 분할 트리 기반 접근법을 도입합니다.
2. 전역 최적화: 어휘 선택을 LP 완화로 밀접하게 연결된 IP 로 공식화함으로써, ToaST 는 국소적인 탐욕적 결정에 의존하는 대신 주어진 분할 트리 집합에 대해 토큰 수를 전역적으로 최소화하는 어휘를 찾습니다.
3. 확장성: 이 방법은 분할 트리의 수에 대해 이차적으로 확장되어, 실현 가능한 훈련 시간 내에 대규모 코퍼스 (프리토큰 카운트의 99% 포함) 를 처리할 수 있습니다.
4. 분리된 설계: 코퍼스 통계에 기반한 트리 구성과 최적화에 기반한 어휘 선택을 분리함으로써, 최적화 모델을 재훈련하지 않고도 형태소 또는 문자 수준의 선호도를 통합하는 것과 같은 유연한 확장이 가능합니다.

실험 결과

저자들은 CulturaX 데이터셋을 사용하여 영어 텍스트에 대해 ToaST 를 평가하고, BPE, WordPiece, UnigramLM 과 비교했습니다.

• 압축: 40,960 이상의 어휘 크기에서 ToaST 는 모든 베이스라인 대비 11% 이상 토큰 수를 감소시킵니다. 이는 토큰당 바이트 수로 측정됩니다.
• 토큰 분포: ToaST 는 어휘 크기 65,536 에서 베이스라인 대비 "Leaf" 토큰 (단일 바이트 폴백) 의 빈도를 14~19 배 크게 감소시킵니다. 반대로, "Root" 토큰 (전체 프리토큰) 을 훨씬 더 많이 생성하여 의미 단위의 더 나은 유지를 나타냅니다.
• Rényi 효율성: ToaST 는 더 높은 Rényi 효율성 ($\alpha=2.5$) 을 달성하여, 더 균일한 토큰 분포와 더 적은 고빈도 토큰을 나타내며, 이는 종종 더 나은 하류 작업 성능과 상관관계가 있습니다.
• 하류 작업 성능: 15 억 파라미터 언어 모델 (nanochat) 훈련 실험에서 ToaST 는 22 개 벤치마크의 평균인 가장 높은 CORE 점수를 달성하여 베이스라인보다 2.6% 에서 7.6% 더 높은 성능을 보였습니다. 이는 BPE 와 UnigramLM 을 크게 능가했으며, 22 개 개별 작업 중 13 개에서 최상의 성능을 기록했습니다.
• 최적화 품질: LP 완화는 매우 밀접한 것으로 밝혀졌습니다. 128 개의 문제 인스턴스 전반에 걸쳐, 반올림된 정수 해와 LP 상한 사이의 최대 상대적 격차는 $5.76 \times 10^{-7}$이었으며, 많은 인스턴스는 완전히 정수였습니다.

중요성 및 주장

이 논문은 ToaST 가 탐욕적 휴리스틱에 의존하는 대신 전역 최적화 프레임워크를 사용하여 압축을 위해 토큰화를 직접 최적화하는 것의 가치를 입증한다고 주장합니다. 저자들은 우수한 압축이 더 큰 유효 컨텍스트 윈도우로 이어지고, 결과적으로 향상된 하류 작업 성능을 가져온다고 주장합니다.

이 연구의 중요성은 다음과 같습니다:

• 실용성: 최적 토큰화의 이론적 NP-완전성에도 불구하고, LP 완화는 이 접근법을 대규모 훈련에 대해 계산적으로 실행 가능하고 실용적으로 만듭니다.
• 성능: 실험 결과는 토큰 수의 11% 감소가 언어 모델 성능 (CORE 점수) 의 측정 가능한 향상으로 이어진다는 것을 시사하며, 더 나은 압축이 더 나은 모델 유용성과 상관관계가 있다는 가설을 지지합니다.
• 유연성: 이 아키텍처는 어휘 최적화 논리를 변경하지 않고도 트리 구성 계층을 수정함으로써 다중 바이트 문자나 형태소 분할을 처리하는 것과 같은 미래의 확장을 허용합니다.

저자들은 한계를 지적했는데, 특히 실험이 영어 텍스트에서만 수행되었으며, 공백으로 구분되지 않는 언어의 경우 상위 집계된 프리토큰으로 훈련하는 효과는 다를 수 있다고 언급했습니다. 또한 코드 구현 및 그림 생성에 LLM 을 사용했음을 공개했습니다.