AI 스택을 직접 구축하기 (1부): BPE 토크나이저 만들기

현대 AI 스택을 처음부터 재구축하는 시리즈의 1부입니다. 이번 스프린트의 코드와 테스트는 ai-stack-sprints repository에 있습니다.

제가 이것을 하는 이유

저는 오랫동안 수동으로 코딩하지 않았다는 것을 깨달았습니다. AI 기반 개발 덕분에 훨씬 빠르게 움직일 수 있지만, 저는 여전히 직접 코딩하고 시스템을 이해하는 것이 가치 있는 기술이라고 믿습니다. 그리고 AI가 더 많은 개발 작업을 맡으면서 그 능력이 무뎌지고 있다는 느낌을 받았습니다.

또한 오랫동안 LLM(대규모 언어 모델)을 사용했지만, 그것들이 어떻게 작동하는지 이해하지 못했습니다. 그래서 저는 Mac mini에서 현대 AI 스택을 한 조각씩 재구축하고 있습니다. GPT-5를 구축할 필요는 없습니다. 제 목표는 제가 호출하여 응답을 받을 수 있는 무언가를 만드는 것입니다.

규칙은 간단합니다:

• 2주마다, 또는 시간이 되면 더 일찍 스프린트를 완료하고 그에 대한 게시물을 발행할 것입니다.
• AI가 스캐폴딩(scaffolding), 벤치마크 및 테스트를 설정해 주는데, 제가 주말에만 이것을 하기 때문입니다.
• 모든 구현 코드는 수동으로 작성해야 합니다: 자동 완성이나 AI 생성 코드는 안 됩니다.
• 구문은 웹을 사용할 수 있지만, 먼저 스스로 문제를 해결하려고 노력할 것입니다.

이것이 첫 번째 스프린트입니다: 토크나이저(tokenizer)입니다.

LLM은 프롬프트 작성기에 입력하는 글자들을 볼 수 없습니다. 토크나이저는 텍스트를 모델이 처리할 수 있는 토큰 ID로 변환하고, 모델이 응답할 때 토큰 ID를 다시 텍스트로 변환합니다.

제가 생각한 토크나이저의 역할

이것을 시작하기 전까지는 토큰화에 대해 크게 생각하지 않았습니다. LLM이 입력을 처리하는 데 도움이 된다는 것은 알고 있었지만, 자세히 살펴본 적은 없었습니다.

제가 만든 것 중 가장 비슷한 것은 멀티플레이어 3D 체스 게임에서였습니다. 그곳에서는 게임 상태 메시지가 JSON으로 도착했고, 이벤트 유형과 페이로드로 파싱되었으며, 업데이트된 상태로 변환되었습니다. 문제는 다르지만 모양은 비슷합니다: 시스템이 사용할 수 있도록 한 표현을 다른 표현으로 변환하는 것입니다.

제가 처음 생각했던 모델은 토크나이징(tokenization)이 텍스트를 구분자(delimiter)로 분할하고, 그 조각들을 효율성을 위해 변환한 다음, 모델의 핵심으로 보내는 것이었습니다.

하지만 이 구현 방식은 구분자로 시작하지 않습니다. 원시 바이트(raw bytes)에서 시작하여, 어떤 인접한 시퀀스들이 충분히 자주 발생하여 자체적인 단위가 되는지를 학습합니다.

v1 구축: 쌍 카운트 및 병합 반복 (count pairs, merge, repeat)

목표는 순환 과정(round trip)을 완성하는 것이었습니다:

text -> tokens -> original text

저는 두 개의 헬퍼 함수로 시작했습니다:

• get_stats: 토큰 ID 목록에서 모든 인접한 쌍(adjacent pair)의 개수를 셉니다. 이 쌍들은 중첩됩니다: [1, 2, 1]의 경우, (1, 2)와 (2, 1) 모두 카운트됩니다.
• merge_pair: 선택된 쌍이 비중첩적으로 발생하는 모든 경우를 새로운 토큰 ID로 대체한 새 목록을 반환합니다.

이 두 함수는 BPE의 두 가지 핵심 질문에 답합니다: 어떤 쌍이 가장 자주 발생하는가? 그리고 그것을 모든 곳에서 병합하면 어떻게 되는가?

학습 (Training)

학습은 텍스트를 UTF-8 바이트로 인코딩하면서 시작됩니다. 기본 어휘(base vocabulary)는 가능한 모든 256개 바이트 값, 즉 ID 0부터 255까지를 포함합니다. ASCII가 포함되지만 그것이 한계는 아닙니다: ASCII 범위를 벗어나는 문자는 UTF-8 바이트 시퀀스로 변환됩니다.

요청되는 어휘 크기는 최소 256이어야 합니다. 학습된 병합(merge)은 다음 사용 가능한 ID를 받으며, 이는 256부터 시작합니다.

반복 과정은 다음과 같습니다:

1. 현재 토큰 목록에서 인접한 쌍의 개수를 셉니다.
2. 가장 빈번하게 발생하는 쌍을 선택합니다.
3. self.merges에 (왼쪽, 오른쪽) -> 새_ID 형식으로 기록합니다.
4. 그 쌍을 토큰 목록 전체에서 대체합니다.
5. 목표 어휘 크기에 도달하거나 병합하기에 충분히 빈번한 쌍이 더 이상 없을 때까지 반복합니다.

어휘 크기는 병합 예산(merge budget)입니다: 256은 학습된 병합을 허용하지 않으며, 512는 최대 256개의 병합을 허용합니다.

마지막에 self.merges는 학습된 규칙서가 됩니다. 인코딩은 이를 새로운 텍스트에 적용하고; 디코딩은 토큰 ID를 바이트로 확장한 다음 UTF-8로 디코드합니다.

문제가 발생한 지점: 디코딩은 단일 통과(one-pass) 작업이 아니다

디코딩은 토크나이저에서 가장 어려웠던 부분이었습니다.

인코딩은 비교적 간단했습니다. 학습된 순서대로 병합(merge)을 반복하며 각 규칙에 대해 merge_pair를 호출하면 되었습니다. 즉, 규칙당 토큰 시퀀스를 한 번 통과하는 방식이었습니다.

저의 첫 디코딩 시도 역시 반복적이었습니다. 병합 결과로 나온 입력 토큰 각각에 대해, 저는 이를 생성한 두 개의 ID로 한 번 대체했습니다. 한 번의 과정을 거친 후에는 문자열로 다시 변환할 수 있는 바이트를 얻을 것이라고 기대했습니다.

문제를 드러낸 예시는 다음과 같습니다:

tok = BPETokenizer()
tok.train(
    "the quick brown fox jumps over the lazy dog. " * 30,
...

최소 빈도에 도달하지 못한 나머지 쌍이 없어 학습은 44개의 병합(merges)에서 중단되었습니다. 일부 병합 규칙은 다음과 같았습니다:

{
    (116, 104): 256,  # "th"
    (256, 101): 257,  # "the"
...

전체 문장은 단 두 개의 토큰으로 인코딩되었습니다:

[258, 293]

하지만 이 토큰들을 각각 확장했을 때 다음과 같이 나왔습니다:

[257, 32, 292, 103]

이것은 바이트 목록이 아니었습니다. 토큰 257은 여전히 "the"를 나타냈고, 토큰 292는 문장의 거의 전체 나머지 부분을 나타냈습니다. 257과 292가 255보다 크기 때문에 아예 bytes 객체에 넣을 수 없었습니다. 저는 부분적으로 확장된 시퀀스를 손에 쥐게 되었습니다.

재귀(Recursion)가 빠진 아이디어였습니다

병합된 토큰은 또 다른 병합된 토큰을 포함할 수 있습니다. 따라서 한 번 확장하는 것만으로는 작업이 완료되지 않으며, 단지 다음 계층을 드러낼 뿐입니다.

해결책은 모든 병합된 토큰을 작은 트리로 취급하는 것이었습니다. private __decode_single 함수는 256 미만의 ID를 바이트로 반환합니다. 그렇지 않은 경우, 해당 ID를 생성한 쌍을 찾아 두 자식 노드를 재귀적으로 확장하여 모든 리프(leaf)가 바이트가 될 때까지 진행했습니다.

그러면 decode는 모든 입력 토큰에서 바이트를 연결하고 이를 UTF-8로 디코딩합니다. 위의 문장의 경우, 재귀를 통해 결국 원래의 43바이트가 생성됩니다.

버그를 통해 얻은 교훈은 명확했습니다: 토큰이 단순히 한 번 대체되었다고 해서 디코딩되는 것이 아닙니다. 그 아래의 모든 값이 바이트에 도달했을 때 디코딩됩니다.

숫자들

저는 동일한 작은 코퍼스에서 세 가지 어휘 크기를 벤치마킹했습니다. 단지 4,646바이트만으로도 이 결과는 예상되는 생산 성능보다 트레이드오프의 형태를 더 잘 보여줍니다.

corpus: 4,644 chars / 4,646 bytes
vocab= 256  merges=  0  tokens= 4,646  ratio=100.00%  bytes/token=1.00  train=0.00s  encode=0.000s
vocab= 384  merges=128  tokens= 2,346  ratio=50.50%  bytes/token=1.98  train=0.10s  encode=0.041s
...

어휘 크기가 256일 때는 학습된 병합(merges)이 없으므로 모든 바이트가 자체 토큰으로 남아 있습니다: 4,646개 토큰.

어휘 크기가 512일 때는 동일한 텍스트가 1,882개의 토큰을 사용합니다. 이는 바이트 레벨 기준선보다 59.5% 적은 토큰 수입니다. 시퀀스가 짧다는 것은 요청당 모델이 처리해야 할 입력이 적다는 것을 의미합니다.

감소하는 수익률(diminishing returns) 또한 흥미로웠습니다:

• 첫 128개의 병합으로 토큰 수가 2,300개 감소했습니다.
• 다음 128개의 병합으로는 단지 464개만 감소시켰습니다.

BPE는 탐욕적(greedy)입니다: 빈번한 쌍이 먼저 병합되므로, 나중의 병합은 더 희귀한 패턴을 추구합니다. 어휘 크기는 '클수록 항상 좋다'는 문제가 아니라 트레이드오프입니다.

왜 딕셔너리인가 — 그리고 이 버전이 느린 이유

저는 두 가지 핵심 위치에서 딕셔너리를 사용했습니다.

get_stats는 인접한 쌍(adjacent pair)을 딕셔너리 키로 사용하여 시퀀스를 스캔하면서 발생 횟수를 셀 수 있습니다. self.merges는 동일한 쌍을 키로 사용하고 이를 새로운 토큰 ID에 매핑합니다.

merges 딕셔너리는 두 번째로, 그리고 덜 명확한 이점을 제공합니다: Python 딕셔너리는 삽입 순서를 보존하기 때문입니다. merges는 학습되는 순서대로 삽입되므로, 삽입 순서는 곧 그들의 순위이기도 합니다. 인코딩 중에는 이 딕셔너리가 규칙집이자 우선순위 목록 역할을 동시에 수행합니다.

이러한 트레이드오프는 디코딩에서 나타납니다. self.merges는 한 쌍(pair)에서 토큰 ID로 앞으로 향합니다. 디코딩은 역방향을 필요로 합니다. 제 구현은 필요한 순간에 그 쌍을 검색합니다. 저는 역 딕셔너리나 미리 계산된 토큰-바이트 테이블을 만들 수도 있지만, 이는 추가 메모리와 중복 표현의 비용으로 조회 속도를 빠르게 만듭니다. 첫 번째 패스에서는 이 부분을 생략했습니다.

학습 역시 비슷한 트레이드오프를 가집니다. 각 라운드마다 제 코드는 모든 쌍(pair)을 처음부터 다시 계산하고, 카운트를 정렬하여 승자를 찾습니다. M이 merges의 수이고 N이 현재 시퀀스 길이일 때, 반복적인 전체 스캔은 대략 M × N, 그리고 정렬 비용이 발생합니다.

이는 필요 이상입니다. 한 쌍이 병합될 때, 병합 지점 바로 주변의 쌍들만 변경되고 다른 곳의 카운트는 유효하게 유지됩니다. 최적화된 구현은 이러한 이웃(neighbor)들만 업데이트할 수 있겠지만, 변화하는 카운트와 승자를 정확하게 추적하기 위해 훨씬 더 많은 장부 정리(bookkeeping)가 필요합니다.

인코더 역시 같은 단순성 트레이드오프를 가집니다: 학습된 모든 merges 규칙에 대해 전체 토큰 시퀀스를 한 번 스캔합니다.

마지막으로, 더 큰 어휘(vocabulary)는 시퀀스를 짧게 만들지만, 압축 이득은 점차 감소하는 반면, 학습 작업량, 모델의 임베딩 및 출력 레이어 크기, 그리고 배워야 하는 희귀 토큰의 수는 꾸준히 증가합니다.

이러한 트레이드오프들은 4,646바이트 분량의 학습 코퍼스와 수백 개의 merges에 대해서는 허용 가능합니다. 하지만 프로덕션 규모에서는 그대로 유지될 경우 허용되지 않을 것입니다.

다음: 파트 2

이 버전은 정확하지만, 의도적으로 순진합니다.

파트 2에서는 이를 최적화할 것입니다: 최대값을 찾기 위해 모든 쌍의 개수를 정렬하는 것을 피하고, 디코딩에 필요한 조회(lookup)를 미리 계산하며, 실제로 무엇이 개선되는지 확인하기 위해 동일한 테스트와 벤치마크를 다시 실행할 것입니다. 또한 완전한 v1 대 v2 비교 자료를 바탕으로 프로덕션 규모에서의 트레이드오프(trade-offs) — 그리고 제가 처음 시작했을 때 오해했던 부분들—을 재검토할 것입니다.