Pitch: Cansei do BPE. Criei um tokenizador contínuo e livre de vocabulário usando OpenAI Triton (FuzzyTok)
A forma como alimentamos texto nos Large Language Models (LLMs) hoje me incomoda. Quase todos os modelos modernos dependem de algoritmos de tokenização discreta de subpalavras, como Byte Pair Encoding (BPE) ou SentencePiece. Embora sejam a norma, eles trazem problemas graves:
- Cegueira morfológica: Quebram palavras de forma arbitrária (ex: "pre-processamento" vira sub-tokens desconexos).
- Fragilidade: Um único erro de digitação muda completamente os tokens gerados, quebrando a semântica para o modelo.
- Sobrecarga de memória: As tabelas de embedding escalam linearmente com o tamanho do vocabulário, comendo gigabytes de VRAM e parâmetros que poderiam ser mais bem aproveitados.
Para tentar resolver isso de forma nativa e eficiente, eu desenvolvi o FuzzyTok.
Ele elimina completamente o vocabulário discreto. O FuzzyTok lê bytes UTF-8 brutos diretamente e os mapeia para um espaço latente contínuo e denso em tempo real através de uma pilha de convoluções causais dilatadas 1D.
Como funciona a arquitetura?
Em vez de buscar um ID em uma tabela estática de 32 mil tokens, o FuzzyTok recebe o fluxo de bytes e calcula uma representação contínua. Usando dilatações exponenciais (1, 2 e 4), conseguimos um campo receptivo de até 43 caracteres. Isso permite ao modelo capturar prefixos, sufixos e radicais dinamicamente.
A causalidade temporal é estrita: a saída no passo t depende apenas dos bytes até o passo t.
O gargalo do hardware (Triton ao resgate)
Processar texto caractere por caractere expande o comprimento da sequência, o que costuma explodir o uso de memória e tempo de GPU. Fazer isso com convoluções padrão do PyTorch em contextos longos gera um gargalo absurdo de VRAM por conta do carregamento repetido de tensores da memória global (VRAM) para a memória local (SRAM).
Escrevi então um kernel de GPU customizado usando OpenAI Triton:
- Ladrilhamento de Canais 2D (2D Channel Tiling): Dividi as dimensões dos canais em blocos de 32x32. Os dados são carregados na SRAM rápida, onde as acumulações parciais rodam direto nas Tensor Cores nativas da GPU com
tl.dot. Isso reduziu o tráfego de memória global em 32x. - Prevenção de Register Spilling: Mantive os blocos de processamento limitados a 32 para evitar que a GPU transbordasse dados para a VRAM lenta durante as operações.
- Máscara Causal Dinâmica: Ajustei o cálculo de ponteiros para evitar segmentation faults quando o índice retrocede no tempo, resolvendo isso com
tl.whereno nível do registrador.
Resultados e Impacto (NVIDIA Tesla T4)
Abaixo estão os resultados que obtive comparando o FuzzyTok com baselines equivalentes treinados para predição do próximo caractere por 40 épocas:
- Perplexidade: FuzzyTok alcançou uma perplexidade final de 1.51 (perda de 0.4110), comparado a 8.38 do baseline BPE Conv1D e 11.64 de um Transformer causal tradicional.
- Consumo de VRAM: Redução de até 38.1% comparado ao PyTorch Conv1D em contexto médio, e até 95.7% de economia contra o Transformer clássico em contextos longos (392 MB vs 9.24 GB).
- Vazão (Throughput): O kernel Triton otimizado alcançou um ganho de velocidade de até 4.46x sobre a implementação inicial do Triton que sofria com register spilling.
O projeto é open source
Queria muito o feedback da comunidade sobre a arquitetura do modelo e o design do kernel Triton. Se você trabalha com otimização de modelos de linguagem ou computação de baixo nível em GPU, ficaria muito grato se desse uma olhada no código e nos benchmarks.
O repositório do projeto, contendo também o artigo acadêmico completo e um notebook interativo para rodar os testes, está aberto aqui: