Tokenization for Molecular Foundation Models

Resumen Técnico: Tokenización para Modelos Fundacionales Moleculares

Planteamiento del Problema
La predicción precisa de las propiedades químicas es crítica para industrias que van desde el almacenamiento de energía hasta el descubrimiento farmacéutico. Si bien las arquitecturas Transformer han revolucionado el Procesamiento de Lenguaje Natural (NLP), su aplicación a los modelos fundacionales moleculares enfrenta un cuello de botella fundamental: la tokenización. Los modelos moleculares actuales dependen predominantemente de una tokenización de tipo "Atom-wise" (por átomo), donde las cadenas de SMILES (Simplified Molecular Input Line Entry System) se dividen en tokens a nivel de átomo utilizando vocabularios fijos.

La limitación principal de este enfoque es su incapacidad para cubrir completamente la especificación OpenSMILES. Los tokenizadores atom-wise tratan los átomos entre corchetes (que codifican isótopos, centros quirales, cargas y recuentos de hidrógeno explícitos) como tokens únicos e indivisibles. Para cubrir cada permutación posible de estas características, un vocabulario requeriría más de 28 billones de tokens. En consecuencia, los modelos existentes utilizan vocabularios de menos de 3,000 tokens, lo que resulta en brechas de cobertura significativas. Al encontrarse con átomos entre corchetes novedosos, estos tokenizadores de vocabulario cerrado recurren a un token de desconocido genérico [UNK], lo que potencialmente oscurece información química crítica como la quiralidad o la composición isotópica específica. Además, las alternativas de vocabulario abierto existentes (como los esquemas basados en BPE) a menudo sufren de ambigüedad, donde entidades químicas distintas (por ejemplo, un enlace azufre-carbono frente al elemento Escandio) se confunden en el mismo token.

Metodología
Los autores proponen un nuevo marco de tokenización, Smirk, y una variante comprimida, Smirk-GPE, diseñados para lograr una cobertura completa de la especificación OpenSMILES manteniendo la eficiencia computacional.

1. Tokenización Smirk: Este esquema emplea una descomposición de dos etapas a nivel de carácter de las cadenas SMILES basada en los glifos definidos por la especificación OpenSMILES.

   • Etapa 1: Descomposición en átomos (ej. OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Etapa 2: Descomposición de los átomos entre corchetes en glifos constituyentes (ej. [C@@H] $\rightarrow$ [ C @ @ H ]).
   • Este enfoque distingue entre secuencias ambiguas (ej. Sc como un enlace frente a [Sc] como Escandio) al tratar los corchetes y los símbolos internos como tokens distintos. El vocabulario resultante es fijo de 165 tokens, no requiere entrenamiento y garantiza que cualquier molécula codificada en OpenSMILES pueda ser tokenizada sin utilizar un token [UNK].

2. Smirk-GPE (Glyph Pair Encoding): Para abordar el aumento en la longitud de la secuencia (fertilidad) causado por la descomposición total de los átomos entre corchetes, los autores implementaron Smirk-GPE. Esta variante aplica una estrategia de compresión similar a BPE (Byte-Pair Encoding) específicamente a los tokens de glifos. A diferencia del BPE estándar que combina cadenas, Smirk-GPE aprende reglas de fusión sobre los IDs de los tokens, asegurando que las fusiones con significado químico (ej. combinar un glifo de azufre y uno de carbono) no creen ambigüedad con símbolos atómicos (ej. Escandio).

3. Marco de Evaluación:

   • Métricas Intrínsecas: Los autores evaluaron los tokenizadores utilizando fertilidad (longitud media de secuencia), entropía normalizada (eficiencia de compresión), desequilibrio de tokens y la frecuencia del token [UNK].
   • Proxy de Bajo Costo: Reconociendo que entrenar modelos Transformer completos para cada tokenizador es computacionalmente costoso, los autores utilizaron modelos n-grama como un proxy. Entrenaron modelos n-grama sobre 1.6 mil millones de cadenas SMILES y midieron la pérdida de entropía cruzada y la pérdida de información (vía divergencia KL) para estimar el rendimiento downstream.
   • Validación Extrínseca: Para validar el proxy de n-gramas, los autores preentrenaron 18 modelos RoBERTa de solo codificador (desde cero) utilizando 11 tokenizadores diferentes y tres codificaciones moleculares. Estos modelos fueron ajustados (fine-tuned) en seis tareas de regresión y siete de clasificación de MoleculeNet y tmQM.

Resultados Clave

• Cobertura: Smirk y Smirk-GPE son los únicos tokenizadores evaluados que logran una cobertura del 100% de la especificación OpenSMILES, eliminando el uso del token [UNK]. En contraste, los tokenizadores específicos de química existentes (incluyendo variantes de SPE, APE y varios tipos de BPE) emiten el token [UNK] con una frecuencia no despreciable (hasta ~50% en el conjunto de datos tmQM).
• Pérdida de Información: Los tokenizadores con cobertura limitada exhiben una pérdida de información sustancial, particularmente en conjuntos de datos ricos en metales de transición y estereoquímica (ej. tmQM). Por ejemplo, el tokenizador de MoLFormer incurre en una pérdida de 40.3 nats/molécula en tmQM debido a tokens desconocidos, mientras que las variantes de Smirk mitigan esta degradación.
• Correlación de Rendimiento: El estudio encontró una fuerte correlación lineal entre las métricas de n-gramas (entropía cruzada y pérdida de información) y el rendimiento downstream de los modelos basados en Transformer. Esto valida el uso de n-gramas como un proxy de bajo costo para evaluar la calidad del tokenizador.
• Impacto Downstream:
  • Smirk mostró un efecto positivo en la calidad del preentrenamiento y el rendimiento downstream en el conjunto de datos tmQM.
  • En las tareas de MoleculeNet (dominadas por moléculas orgánicas pequeñas), Smirk se desempeñó de manera similar a la tokenización atom-wise estándar.
  • Los tokenizadores con mala cobertura (SPE/APE) afectaron negativamente tanto el preentrenamiento como el rendimiento downstream en comparación con el baseline.
  • Se encontró que la elección de la codificación molecular (SMILES vs. SELFIES) tiene un impacto insignificante en comparación con la elección del tokenizador.

Significancia y Reivindicaciones
El artículo argumenta que un modelo fundacional para la química debe codificar toda la amplitud del espacio químico para evitar oscurecer características críticas. Los autores afirman que los tokenizadores actuales oscurecen inadvertidamente la información a nivel de átomo (como la quiralidad en Cisplatino o isótopos específicos), lo que desencadena una pérdida de información potencialmente significativa que no es meramente teórica, sino que impacta en moléculas clínicamente o industrialmente relevantes.

La significancia de este trabajo reside en:

1. Robustez: Demostrar que los tokenizadores de vocabulario abierto (Smirk/Smirk-GPE) proporcionan una cobertura robusta del espacio químico, evitando la pérdida de información asociada con los tokens desconocidos.
2. Eficiencia: Establecer que los modelos n-grama pueden servir como un proxy fiable y de bajo costo para evaluar el rendimiento de los tokenizadores, reduando la carga computacional de la selección de modelos y el ajuste de hiperparámetros.
3. Interpretabilidad: Destacar que Smirk permite a los investigadores manipular directamente el contenido rico en información de los átomos entre corchetes, expandiendo los beneficios de interpretabilidad de la tokenización atom-wise al tiempo que elimina el riesgo de errores de vocabulario fuera de rango (out-of-vocabulary).

Los autores concluyen que, si bien los benchmarks actuales (como MoleculeNet) pueden no exponer plenamente las deficiencias de los tokenizadores de cobertura limitada debido a la falta de diversidad en elementos y estereoquímica, la transición hacia tokenizadores capaces de codificar la totalidad del espacio químico es necesaria para modelos fundacionales moleculares fiables. Instan a la comunidad a evaluar rigurosamente el alcance de los benchmarks y a expandir los conjuntos de datos para incluir características químicas diversas.