Tokenization for Molecular Foundation Models

Resumo Técnico: Tokenização para Modelos de Fundação Molecular

Definição do Problema
A previsão precisa de propriedades químicas é crítica para indústrias que variam desde o armazenamento de energia até a descoberta farmacêutica. Embora as arquiteturas transformer tenham revolucionado o Processamento de Linguagem Natural (PLN), sua aplicação a modelos de fundação molecular enfrenta um gargalo fundamental: a tokenização. Os modelos moleculares atuais dependem predominantemente de uma tokenização "por átomo" (Atom-wise), onde strings SMILES (Simplified Molecular Input Line Entry System) são divididas em tokens de nível atômico usando vocabulários fixos.

A limitação primária desta abordagem é sua incapacidade de cobrir totalmente a especificação OpenSMILES. Tokenizadores por átomo tratam átomos entre colchetes (que codificam isótopos, centros quirais, cargas e contagens explícitas de hidrogênio) como tokens únicos e indivisíveis. Para cobrir todas as permutações possíveis desses recursos, um vocabulário exigiria mais de 28 trilhões de tokens. Consequentemente, os modelos existentes utilizam vocabulários de menos de 3.000 tokens, resultando em lacunas significativas de cobertura. Ao encontrar átomos entre colchetes inéditos, esses tokenizadores de vocabulário fechado recorrem a um token de desconhecido genérico [UNK], o que pode obscurecer informações químicas críticas, como quiralidade ou composição isotópica específica. Além disso, alternativas de vocabulário aberto existentes (como esquemas baseados em BPE) frequentemente sofrem de ambiguidade, onde entidades químicas distintas (por exemplo, uma ligação enxofre-carbono versus o elemento Escândio) são confundidas no mesmo token.

Metodologia
Os autores propõem um novo framework de tokenização, Smirk, e uma variante comprimida, Smirk-GPE, projetados para alcançar a cobertura completa da especificação OpenSMILES mantendo a eficiência computacional.

1. Tokenização Smirk: Este esquema emprega uma decomposição em dois estágios, ao nível de caractere, de strings SMILES baseada nos glifos definidos pela especificação OpenSMILES.

   • Estágio 1: Decomposição em átomos (ex: OC[C@@H][OH] $\rightarrow$ O C [C@@H] [OH]).
   • Estágio 2: Decomposição de átomos entre colchetes em glifos constituintes (ex: [C@@H] $\rightarrow$ [ C @ @ H ]).
   • Esta abordagem distingue sequências ambíguas (ex: Sc como uma ligação versus [Sc] como Escândio) ao tratar os colchetes e os símbolos internos como tokens distintos. O vocabulário resultante é fixo em 165 tokens, não requer treinamento e garante que qualquer molécula codificada em OpenSMILES possa ser tokenizada sem o uso de um token [UNK].

2. Smirk-GPE (Glyph Pair Encoding): Para abordar o aumento no comprimento da sequência (fertilidade) causado pela decomposição total dos átomos entre colchetes, os autores implementaram o Smirk-GPE. Esta variante aplica uma estratégia de compressão semelhante ao Byte-Pair Encoding (BPE) especificamente aos tokens de glifos. Diferente do BPE padrão, que mescla strings, o Smirk-GPE aprende regras de mesclagem nos IDs de tokens, garantindo que mesclagens quimicamente significativas (ex: combinar um glifo de enxofre e um de carbono) não criem ambiguidade com símbolos atômicos (ex: Escândio).

3. Framework de Avaliação:

   • Métricas Intrínsecas: Os autores avaliaram os tokenizadores usando fertilidade (comprimento médio da sequência), entropia normalizada (eficiência de compressão), desequilíbrio de tokens e a frequência do token [UNK].
   • Proxy de Baixo Custo: Reconhecendo que treinar modelos transformer completos para cada tokenizador é computacionalmente caro, os autores utilizaram modelos n-grama como um proxy. Eles treinaram modelos n-grama em 1,6 bilhão de strings SMILES e mediram a perda de entropia cruzada e a perda de informação (via divergência KL) para estimar o desempenho downstream.
   • Validação Extrínseca: Para validar o proxy n-grama, os autores pré-treinaram 18 modelos RoBERTa apenas de encoder (do zero) usando 11 tokenizadores diferentes e três codificações moleculares. Esses modelos foram ajustados (fine-tuned) em seis tarefas de regressão e sete de classificação do MoleculeNet e tmQM.

Principais Resultados

• Cobertura: Smirk e Smirk-GPE são os únicos tokenizadores avaliados que alcançam 100% de cobertura da especificação OpenSMILES, eliminando o uso do token [UNK]. Em contraste, os tokenizadores específicos de química existentes (incluindo SPE, APE e várias variantes de BPE) emitem o token [UNK] com frequência não desprezível (até ~50% no dataset tmQM).
• Perda de Informação: Tokenizadores com cobertura limitada exibem perda substancial de informação, particularmente em datasets ricos em metais de transição e estereoquímica (ex: tmQM). Por exemplo, o tokenizador do MoLFormer incorre em uma perda de 40,3 nats/molécula no tmQM devido a tokens desconhecidos, enquanto as variantes Smirk mitigam essa degradação.
• Correlação de Desempenho: O estudo encontrou uma forte correlação linear entre métricas n-grama (entropia cruzada e perda de informação) e o desempenho downstream de modelos baseados em transformer. Isso valida o uso de n-gramas como um proxy de baixo custo para avaliar a qualidade do tokenizador.
• Impacto Downstream:
  • O Smirk mostrou um efeito positivo na qualidade do pré-treinamento e no desempenho downstream no dataset tmQM.
  • No MoleculeNet (dominado por pequenas moléculas orgânicas), o Smirk teve desempenho similar à tokenização padrão por átomo.
  • Tokenizadores com baixa cobertura (SPE/APE) impactaram negativamente tanto o pré-treinamento quanto o desempenho downstream em relação ao baseline.
  • Descobriu-se que a escolha da codificação molecular (SMILES vs. SELFIES) tem um impacto negligenciável comparado à escolha do tokenizador.

Significância e Alegações
O artigo argumenta que um modelo de fundação para química deve codificar toda a amplitude do espaço químico para evitar a ocultação de características críticas. Os autores alegam que os tokenizadores atuais obscurecem inadvertidamente informações ao nível atômico (como quiralidade em Cisplatina ou isótopos específicos), desencadeando uma perda de informação potencialmente significativa que não é meramente teórica, mas que impacta moléculas clinicamente e industrialmente relevantes.

A significância deste trabalho reside em:

1. Robustez: Demonstrar que tokenizadores de vocabulário aberto (Smirk/Smirk-GPE) fornecem cobertura robusta do espaço químico, prevenindo a perda de informação associada a tokens desconhecidos.
2. Eficiência: Estabelecer que modelos n-grama podem servir como um proxy confiável e de baixo custo para avaliar o desempenho de tokenizadores, reduzindo o ônus computacional de ajuste de hiperparâmetros e seleção de modelos.
3. Interpretabilidade: Destacar que o Smirk permite que pesquisadores manipulem diretamente o conteúdo rico em informações dos átomos entre colchetes, expandindo os benefícios de interpretabilidade da tokenização por átomo ao mesmo tempo em que remove o risco de erros de vocabulário fora do padrão (out-of-vocabulary).

Os autores concluem que, embora os benchmarks atuais (como o MoleculeNet) possam não expor totalmente as deficiências de tokenizadores de cobertura limitada devido à falta de diversidade de elementos e estereoquímica, a transição para tokenizadores capazes de codificar a totalidade do espaço químico é necessária para modelos de fundação molecular confiáveis. Eles incentivam a comunidade a avaliar rigorosamente o escopo dos benchmarks e a expandir os datasets para incluir diversas características químicas.