Faster Superword Tokenization

1. El Problema

El algoritmo de codificación de pares de bytes (BPE) es el estándar en la tokenización de modelos de lenguaje grandes (LLM). Sin embargo, BPE tiene una limitación fundamental: sus tokens no pueden cruzar los límites de los "pretokens" (segmentos iniciales como palabras completas, puntuación o espacios). Esto impide la formación de tokens que abarquen frases completas (superpalabras), limitando la compresión y la eficiencia de la representación.

Algoritmos recientes como BoundlessBPE y SuperBPE intentaron resolver esto introduciendo "superpalabras" (combinaciones de pretokens adyacentes). No obstante, sus implementaciones originales eran imprácticas para el entrenamiento debido a una ineficiencia crítica:

• Mantenían copias completas de todos los documentos en memoria para calcular el contexto de las superfusiones.
• Esto resultó en tiempos de entrenamiento extremadamente largos (ej. 4.7 días de CPU para entrenar en solo 1 GB de datos), en comparación con los segundos que tarda el BPE estándar.
• La falta de agregación de frecuencias de los candidatos a superfusiones hacía que el tiempo de entrenamiento escalara linealmente con el número total de pretokens en lugar de con los únicos.

2. Metodología

Los autores proponen una nueva arquitectura que desacopla y optimiza el proceso de entrenamiento mediante dos fases principales y una nueva estrategia de agregación:

A. Agregación de Candidatos a Superfusiones (Supermerge Candidates)

En lugar de almacenar documentos completos, el método identifica y agrupa los candidatos a superfusiones (dos o más pretokens consecutivos elegibles) por su frecuencia.

• Proceso: Se identifican los pretokens elegibles (aquellos que son una sola letra Unicode y ya tokenizados como un solo token en la fase 1). Se agrupan secuencias adyacentes de estos pretokens elegibles.
• Beneficio: Esto permite trabajar con un conjunto de datos comprimido en memoria, similar a como BPE agrupa pretokens, eliminando la necesidad de escanear documentos completos repetidamente.

B. Tokenización Pre-tokenización Específica por Guion (Script-Specific)

Para idiomas sin espacios (como chino o japonés), donde todo el texto podría ser un solo pretoken, se introduce una regex específica que divide cada carácter Unicode en su propio pretoken de 3 bytes.

• Objetivo: Evitar la creación de tokens inválidos de UTF-8 (bytes parciales de caracteres diferentes) y permitir que las superfusiones combinen caracteres completos en frases según su frecuencia.

C. Formulación de Dos Fases (Two-Phase BoundlessBPE)

Se presenta una reformulación de BoundlessBPE que produce resultados idénticos al método original pero con una estructura de dos fases:

1. Fase 1: Entrenamiento de un tokenizador BPE estándar sobre los pretokens agregados hasta alcanzar el tamaño de vocabulario deseado.
2. Fase 2: Se comparan las reglas de fusión regular de la Fase 1 con los mejores candidatos a superfusiones (agregados en el paso anterior). Se utiliza un algoritmo (Algoritmo 1) que selecciona dinámicamente entre fusiones regulares y superfusiones basándose en la cuenta de pares, sin necesidad de especificar un hiperparámetro manual para la cantidad de superfusiones (a diferencia de SuperBPE).

D. Optimización de Inferencia

Se introduce una optimización en la inferencia: si un pretoken ya existe en el vocabulario como un token único, se deja intacto sin aplicar reglas de fusión. Esto evita el procesamiento innecesario de la mayoría de los pretokens.

3. Contribuciones Clave

1. Implementación de Alta Velocidad: Lograron reducir el tiempo de entrenamiento de 4.7 días de CPU a ~600 segundos (un aumento de velocidad de más de 600x) para el mismo conjunto de datos de 1 GB.
2. Equivalencia Teórica y Práctica: Demostraron que su versión de dos fases de BoundlessBPE es matemáticamente equivalente al método original, pero mucho más eficiente.
3. Unificación de Enfoques: Mostraron la equivalencia casi total entre la versión de dos fases de BoundlessBPE y SuperBPE, permitiendo que BoundlessBPE determine automáticamente el número óptimo de superfusiones en lugar de requerir un ajuste manual.
4. Código de Código Abierto: Liberaron implementaciones de referencia en Python y una versión optimizada en Rust que produce resultados idénticos, facilitando la adopción práctica.
5. Manejo de Idiomas No Espaciados: Introdujeron una estrategia robusta para la pre-tokenización de scripts asiáticos que previene la generación de tokens semánticamente sin sentido (bytes parciales).

4. Resultados

• Velocidad de Entrenamiento:
  • BoundlessBPE: De 4.7 días a 603 segundos (Python) y 593 segundos (Rust) en 1 GB de datos.
  • SuperBPE: Logró tiempos similares de ~593 segundos en la misma configuración.
  • Comparado con BPE estándar (que tarda ~59s en 1 GB), la sobrecarga es mínima (aprox. 10x más lento que BPE, pero 600x más rápido que la implementación anterior de BoundlessBPE).
• Desglose de Tiempo: La mayor parte del tiempo de entrenamiento se gasta en la lógica procedural de fusiones y actualización de contadores. La implementación en Rust reduce significativamente este tiempo en comparación con Python.
• Inferencia: La optimización de "atajo" (verificar si el pretoken está en el vocabulario) proporcionó una aceleración de 2.8x en Python y 1.6x en Rust.
• Escalabilidad: El tiempo de entrenamiento crece de forma sublineal con el tamaño del vocabulario, ya que las fusiones posteriores se aplican a pares cada vez más raros.

5. Significado e Impacto

Este trabajo elimina la principal barrera de adopción para los métodos de tokenización basados en superpalabras (BoundlessBPE y SuperBPE): la viabilidad computacional.

• Viabilidad Práctica: Permite entrenar estos modelos avanzados en conjuntos de datos mucho más grandes (típicos de LLMs) en tiempos razonables, algo que antes era imposible.
• Mejora de Modelos: Al permitir el uso de superpalabras, se mejora la compresión, la eficiencia de Rényi y la distribución de tokens, lo que se ha correlacionado previamente con mejoras en tareas posteriores (hasta un 4% de mejora en precisión).
• Accesibilidad: Al proporcionar implementaciones rápidas en Rust y Python, democratizan el uso de estas técnicas para la comunidad de investigación y práctica, sin requerir infraestructuras masivas de CPU para el entrenamiento.

En resumen, el artículo transforma algoritmos teóricamente superiores pero computacionalmente prohibitivos en herramientas prácticas y eficientes, manteniendo la calidad del modelo y acelerando drásticamente el proceso de entrenamiento.