Introducción: procesar menos tokens para acelerar los modelos de IA
Investigadores de Stanford y la Universidad de California en Berkeley (UC Berkeley) han presentado TokenGate, un método para acelerar el procesamiento de contextos largos en modelos de lenguaje de gran tamaño (LLMs) mediante una técnica de selección aprendida de tokens.
La propuesta aborda uno de los principales costes computacionales de los transformers: procesar todos los tokens de entrada a través de todas las capas del modelo, incluso cuando una parte del contenido puede dejar de ser necesaria para las operaciones posteriores.
TokenGate introduce un mecanismo que aprende a identificar qué tokens deben continuar hacia las capas superiores y cuáles pueden descartarse. A diferencia de una selección fija, el sistema utiliza una pequeña red neuronal para tomar esta decisión a partir de las representaciones internas del modelo.
Según los resultados publicados por sus investigadores, el sistema consigue aceleraciones de entre 1,4 y 11,3 veces en el tiempo hasta el primer token (TTFT) y disminuye el consumo máximo de memoria entre 1,5 y 3,4 veces, sin pérdida de precisión en los resultados generales reportados.
El proyecto se describe en el artículo científico TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning, enviado a arXiv el 19 de octubre de 2026.
Contexto: por qué el prefill consume tantos recursos
El funcionamiento de un modelo de lenguaje incluye una etapa denominada prefill, durante la cual el sistema procesa los tokens del prompt antes de comenzar a generar la respuesta.
En los transformers tradicionales, todos los tokens atraviesan las diferentes capas de procesamiento. Además, el coste computacional de la atención crece de manera cuadrática respecto a la longitud de la secuencia.
Esto resulta especialmente relevante cuando los modelos trabajan con documentos extensos o contextos que contienen decenas de miles de tokens.
TokenGate plantea reducir ese coste haciendo que determinadas capas del modelo procesen únicamente una selección de tokens.
Su diseño permite añadir esta capacidad a modelos previamente entrenados mediante un módulo integrable (drop-in), sin sustituir completamente su arquitectura.
La idea central es que las primeras capas mantengan el procesamiento completo, mientras que las capas posteriores trabajen con una representación reducida de la entrada.
Cómo funciona TokenGate: tres fases para seleccionar los tokens
El mecanismo utiliza una capa de selección, identificada como L, que determina el punto en el que se reduce la cantidad de tokens procesados.
El procedimiento se divide en tres fases.
<h5 id="1-procesamiento-completo-en-las-primeras-capas">1. Procesamiento completo en las primeras capas</h5>Las capas situadas por debajo de L reciben y procesan todos los N tokens de entrada.
En esta etapa inicial no se aplica ninguna reducción. El modelo trabaja con el contenido completo antes de llegar al mecanismo de selección.
<h5 id="2-una-red-mlp-aprende-que-tokens-conservar">2. Una red MLP aprende qué tokens conservar</h5>En la capa L, una pequeña red neuronal multicapa, o MLP, analiza la representación residual posterior a la atención de cada token.
A partir de esa representación, la red asigna una puntuación individual a cada token.
Después, un selector top-K conserva los K tokens con mayor puntuación.
De esta manera, TokenGate aprende a determinar qué partes de la entrada deben seguir participando en el procesamiento.
<h5 id="3-las-capas-superiores-procesan-unicamente-los-tokens-seleccionados">3. Las capas superiores procesan únicamente los tokens seleccionados</h5>Una vez realizada la selección, las capas situadas por encima de L reciben solamente los K tokens conservados.
Los tokens descartados dejan de participar en las operaciones posteriores y no generan entradas en la memoria caché de claves y valores (KV cache) de esas capas.
Esta reducción afecta tanto a la cantidad de tokens procesados como a los recursos de memoria necesarios durante el prefill.
Existe una restricción importante en el diseño: solo puede podarse la parte del prompt correspondiente al documento; la pregunta nunca se elimina.
Así, el mecanismo de selección actúa sobre el contenido documental, mientras que la pregunta permanece disponible para el modelo.
Cómo consigue entrenar una selección top-K
Uno de los desafíos técnicos del método es que la operación top-K no es diferenciable.
Esto dificulta utilizar directamente el procedimiento habitual de entrenamiento mediante gradientes para aprender qué tokens deben seleccionarse.
Para resolverlo, TokenGate incorpora un mecanismo denominado token-gated attention, inspirado en técnicas utilizadas por los routers de arquitecturas Mixture-of-Experts (MoE).
En lugar de utilizar las puntuaciones únicamente para decidir qué tokens permanecen, el método convierte la puntuación de cada token conservado en un sesgo que se aplica a los logits de atención.
El sistema incorpora dos componentes aprendibles:
- Una ganancia por cabeza de atención, que ajusta la influencia de las puntuaciones sobre cada cabeza.
- Una temperatura aprendible, que permite modificar cómo intervienen esas puntuaciones en la atención.
Este diseño permite que las puntuaciones de selección participen en el aprendizaje del modelo.
<h5 id="una-sola-puerta-para-diferentes-presupuestos-de-tokens">Una sola puerta para diferentes presupuestos de tokens</h5>TokenGate utiliza la pérdida habitual de predicción del siguiente token (next-token loss) durante su entrenamiento.
Sin embargo, combina dos condiciones: la pérdida obtenida con el contexto completo y la pérdida calculada utilizando un presupuesto aleatorio de tokens conservados.
Ambas pérdidas se promedian durante el entrenamiento.
El resultado es una puerta aprendida que puede utilizarse con diferentes presupuestos de tokens sin necesitar una puerta independiente para cada cantidad.
Resultados: hasta 11,3 veces más rápido y 3,4 veces menos memoria
Los investigadores evaluaron TokenGate utilizando Qwen3-1.7B y diferentes tareas de procesamiento de contextos largos.
Los experimentos incluyen mediciones del tiempo hasta el primer token, consumo máximo de memoria y precisión.
Según el artículo, TokenGate alcanza mejoras frente al procesamiento completo y también frente a otros métodos diseñados para reducir el coste del prefill.
<h5 id="comparativa-de-rendimiento-reportado">Comparativa de rendimiento reportado</h5>| Escenario | Aceleración | Reducción de memoria pico | Resultado de precisión |
|---|---|---|---|
| Resultados generales de TokenGate | 1,4× a 11,3× en TTFT | 1,5× a 3,4× | Sin pérdida de precisión reportada |
| Frente a la siguiente mejor línea base | 1,2× a 6,4× en TTFT | Hasta 2,1× | — |
| 2WikiMultihopQA | 3,0× en prefill | 2,7× | 95 % de la precisión con contexto completo |
| NIAH con contexto de 16k tokens | 11,3× | 3,42× | — |
| Casos con aceleración más reducida | 1,4× | — | — |
Fuente: resultados reportados en el paper oficial de TokenGate. Las cifras corresponden a los escenarios indicados y no representan una mejora uniforme en todas las tareas.
Una de las comparativas más relevantes procede de 2WikiMultihopQA, una tarea que requiere trabajar con información distribuida en el contexto.
En ese escenario, TokenGate consiguió un prefill tres veces más rápido y un consumo máximo de memoria 2,7 veces menor, manteniendo el 95 % de la precisión obtenida con el contexto completo.
Por su parte, la prueba needle-in-a-haystack (NIAH), utilizada para evaluar la recuperación de información dentro de contextos extensos, produjo uno de los resultados más destacados.
Con una entrada de 16.000 tokens, el sistema alcanzó una aceleración de 11,3 veces y utilizó 3,42 veces menos memoria pico.
El rendimiento varía según la tarea: en algunos experimentos, la aceleración se redujo a 1,4 veces.
Contextos de hasta 128k tokens: una misma puerta mantiene la aceleración
Los investigadores también analizaron cómo se comporta TokenGate cuando aumenta la longitud del contexto.
En particular, utilizaron una puerta entrenada con contextos de 16k tokens y la aplicaron a secuencias más extensas.
Los resultados reportados para ese escenario fueron los siguientes:
| Longitud del contexto | Aceleración reportada |
|---|---|
| 32k tokens | 12,8× |
| 64k tokens | 13,0× |
| 128k tokens | 13,0× |
Estas mediciones corresponden específicamente al escenario de contextos más largos con una puerta entrenada a 16k tokens.
Por tanto, la aceleración de 13,0 veces es un resultado específico de esta evaluación, distinto del intervalo general de 1,4 a 11,3 veces comunicado para los experimentos principales.
El resultado muestra que la puerta aprendida puede utilizarse en longitudes de contexto superiores a la empleada durante su entrenamiento.
Qué tareas y alternativas se utilizaron en la evaluación
Las pruebas realizadas con Qwen3-1.7B incluyeron cinco conjuntos de tareas:
- RepoBench-P: una de las tareas utilizadas para evaluar el método.
- 2WikiMultihopQA: evaluación con recuperación de información distribuida.
- MultiNews: tarea incluida en los experimentos.
- BioDEX: tarea incluida en los experimentos.
- Needle-in-a-haystack (NIAH): evaluación de recuperación de información en contextos extensos.
Además del procesamiento completo, TokenGate se comparó con seis alternativas:
- FastKV
- SlimInfer
- DASH
- GemFilter
- FlexPrefill
- MInference
Las comparaciones utilizaron el mismo modelo base con LoRA, con el objetivo de evaluar los métodos sobre un entorno común.
En estas pruebas, TokenGate consiguió mejoras de entre 1,2 y 6,4 veces en TTFT frente a la siguiente mejor línea base, junto con reducciones de memoria de hasta 2,1 veces.
TokenGate también funciona con un modelo híbrido: Qwen3.5-2B
Además del modelo denso Qwen3-1.7B, los investigadores aplicaron el mecanismo a Qwen3.5-2B, que combina componentes Gated DeltaNet con capas de atención completa.
Esta arquitectura presenta un escenario diferente para la reducción de tokens.
Los experimentos con FastKV se limitan a la poda en capas de atención. TokenGate, en cambio, permite situar su puerta en diferentes capas del modelo híbrido.
Los resultados publicados para NIAH con contextos de 32k tokens muestran cómo cambia el equilibrio entre velocidad y precisión dependiendo de la posición de la puerta.
| Configuración en Qwen3.5-2B | Aceleración | Precisión |
|---|---|---|
| TokenGate en capa 3 | 3,8× frente al prefill completo; 2,1× frente a FastKV | Igual a la del prefill completo |
| TokenGate en capas 1 o 2 | Hasta 7,8× | Dentro del 6 % de la precisión del prefill completo |
Con la puerta situada en la tercera capa, TokenGate igualó la precisión del procesamiento completo y redujo de manera significativa el tiempo de prefill.
Al colocar la puerta en las capas primera o segunda, la aceleración aumentó hasta 7,8 veces, con un resultado de precisión dentro del 6 % respecto al procesamiento completo.
Estos resultados muestran la relación entre la posición del mecanismo de poda, el rendimiento computacional y la precisión en el modelo híbrido.
Implementación: una puerta MLP y adaptadores LoRA
TokenGate está diseñado para adaptarse a modelos que ya han sido entrenados.
Su implementación incorpora una puerta MLP de tamaño reducido junto con adaptadores LoRA de rango 16 (r=16).
Los experimentos publicados utilizan dos modelos:
- Qwen3-1.7B, empleado en las evaluaciones principales con cinco tareas.
- Qwen3.5-2B, utilizado para estudiar el comportamiento del método en una arquitectura híbrida.
Las mediciones de rendimiento del artículo se realizaron utilizando una GPU NVIDIA A100 con un tamaño de lote (batch size) de 1.
Para el entrenamiento, los investigadores utilizaron cuatro GPU A100.
<h5 id="coste-de-entrenamiento">Coste de entrenamiento</h5>El método introduce un coste adicional por paso respecto al entrenamiento habitual con LoRA.
Según los datos publicados, cada paso de entrenamiento con TokenGate es 1,3 veces más lento que un paso convencional de LoRA.
El artículo también reporta tiempos medios de entrenamiento de 1,8 frente a 2,2 horas A100.
Estos datos distinguen el coste de cada paso del tiempo medio de entrenamiento registrado.
Código abierto: requisitos, parámetros y modelos disponibles
Los investigadores publicaron la implementación en el repositorio oficial de TokenGate en GitHub, bajo licencia Apache-2.0.
La documentación del proyecto establece los siguientes requisitos de software:
- Python 3.10 o posterior.
- PyTorch 2.9 o posterior con soporte CUDA.
La implementación incluye un kernel opcional de FlashAttention-2 con sesgo de clave, utilizado para acelerar la ejecución con el modelo denso Qwen3.
Cuando ese kernel no está disponible, el sistema utiliza FlexAttention.
<h5 id="como-identifica-el-contenido-que-puede-podarse">Cómo identifica el contenido que puede podarse</h5>La interfaz del repositorio utiliza un marcador especial, {state}, para señalar la parte documental del prompt sobre la que puede actuar la poda.
Este marcador permite distinguir el contenido susceptible de reducción de la pregunta, que permanece intacta.
Por otra parte, la función generate acepta el parámetro max_kept_prefill_tokens, que establece el presupuesto máximo de tokens que se conservarán durante el prefill.
Este parámetro está relacionado con una de las características centrales del método: una misma puerta entrenada puede trabajar con diferentes presupuestos de tokens.
<h5 id="checkpoints-publicados">Checkpoints publicados</h5>El repositorio incluye checkpoints para las cinco tareas evaluadas con Qwen3-1.7B y una puerta entrenada para NIAH sobre Qwen3.5-2B.
Entre los ejemplos disponibles se encuentra:
liana313/tokengate-niah-qwen3-1.7b
La publicación de estos checkpoints acompaña al código del mecanismo de selección y a su implementación con los modelos utilizados en la investigación.
Qué aporta TokenGate al procesamiento de contextos largos
La contribución principal de TokenGate es incorporar un mecanismo de selección aprendida que reduce la cantidad de tokens procesados por las capas superiores de un LLM.
El sistema combina tres elementos: una puerta neuronal que puntúa tokens, un mecanismo de atención que permite entrenar esa selección y un procedimiento de aprendizaje que utiliza presupuestos variables.
Frente a los métodos que mantienen la totalidad de los tokens durante todo el procesamiento, TokenGate evita que los elementos descartados continúen ocupando recursos en las capas posteriores.
Los resultados publicados muestran mejoras tanto en el tiempo hasta el primer token como en el consumo máximo de memoria, con diferencias según el modelo, la tarea, la longitud del contexto y la posición de la puerta.
Las evaluaciones con Qwen3-1.7B y Qwen3.5-2B también muestran que el método puede aplicarse tanto a modelos densos como a una arquitectura híbrida que combina Gated DeltaNet y atención completa.
En contexto: Stanford, UC Berkeley y el proyecto TokenGate
TokenGate es un proyecto de investigación presentado por investigadores de Stanford y UC Berkeley.
El trabajo científico lleva por título TokenGate: Adapting LLMs for Fast and Memory-Efficient Prefill with Learned Token Pruning y fue enviado a arXiv el 19 de octubre de 2026.
Sus autores son Liana Patel, Rishabh Ranjan, Ion Stoica, Matei Zaharia y Carlos Guestrin.
El proyecto se centra en añadir poda aprendida de tokens a modelos de lenguaje previamente entrenados para reducir el coste de procesamiento y memoria durante el prefill.
El artículo científico está disponible en arXiv, mientras que la implementación, publicada con licencia Apache-2.0, se encuentra en GitHub.
Imagen de portada: ilustración generada con IA.
