Investigadores de Stanford y UC Berkeley han lanzado TokenGate, un módulo para prefill de contexto largo. La poda de tokens aprendida de TokenGate permite que un LLM existente examine su entrada de forma rápida. Solo los tokens que importan reciben procesamiento completo. El equipo de investigación informa 1.4x a 11.3x más rápido en tiempo hasta el primer token (TTFT) sin pérdida de precisión. La memoria máxima de prefill se reduce 1.5x a 3.4x. El código y los checkpoints están en GitHub bajo Apache-2.0.
TL;DR
- Compuerta MLP: Una pequeña compuerta MLP y LoRA (r=16) añadidas a Qwen3-1.7B o Qwen3.5-2B. Probado en contextos de hasta 128k tokens.
- Se ejecuta en: Los tiempos del artículo usan 1 A100 con batch size 1. El entrenamiento usó 4 A100.
- Rendimiento: Pareto-óptimo en 5 tareas frente a 6 líneas base.
- Mejor caso: NIAH 16k con 11.3x de TTFT más rápido y 3.42x menos memoria. La ganancia de TTFT alcanza 13.0x en 128k.
- Peor caso: La aceleración cae a 1.4x en algunas tareas.
¿Qué es TokenGate?
TokenGate es un componente que se integra directamente y añade poda de tokens entrenable a LLM preentrenados. Los Transformers hacen pasar cada token por cada capa. El costo de la atención crece cuadráticamente con la longitud. La poda sin entrenamiento es barata pero pierde precisión a tasas de poda altas. Los cambios arquitectónicos a menudo requieren preentrenamiento desde cero.
¿Cómo funciona la poda de tokens aprendida de TokenGate?
La pasada hacia adelante se divide en 3 fases:
- Pasada hacia adelante estándar: las capas por debajo de la capa de compuerta L procesan todos los N tokens como de costumbre.
- Módulo de compuerta: un MLP puntúa cada token a partir de su residual posterior a la atención en la capa L. Un selector top-K conserva los tokens con mayor puntuación.
- Pasada hacia adelante podada: las capas por encima de L procesan solo esos K tokens y nunca crean entradas KV para los podados.
Solo la parte del documento del prompt es podable. La pregunta nunca se poda.
¿Qué es la atención con control de acceso por tokens (token-gated attention)?
La selección Top-K no es diferenciable, lo que bloquea el entrenamiento de extremo a extremo. TokenGate toma prestado el truco del router de Mixture-of-Experts. La puntuación de cada token conservado se convierte en un sesgo en los logits de atención por encima de la compuerta. Ese sesgo da a los gradientes un camino de regreso hacia el scorer.
Una ganancia aprendible por cabeza y una temperatura aprendible refinan esto. La temperatura corrige el hecho de que cada consulta ve menos claves. El entrenamiento usa la pérdida simple de siguiente token. Cada paso promedia una pérdida de contexto completo y una pérdida de presupuesto aleatorio. Así, 1 compuerta entrenada funciona en todos los presupuestos.
¿Cómo rinde TokenGate en los benchmarks?
El equipo de investigación probó Qwen3-1.7B en RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX y needle-in-a-haystack (NIAH). Los baselines fueron FastKV, SlimInfer, DASH, GemFilter, FlexPrefill y MInference, sobre el mismo backbone LoRA.
- Frente al prefill completo: TTFT de 1.4x a 11.3x más rápido y memoria de 1.5x a 3.4x menor.
- Frente al siguiente mejor baseline: TTFT de 1.2x a 6.4x más rápido y memoria hasta 2.1x menor.
- 2WikiMultihopQA: al 95% de la precisión de contexto completo, el prefill corre 3.0x más rápido con 2.7x menos memoria.
- Contextos más largos: una compuerta entrenada a 16k alcanza 12.8x a 32k y 13.0x a 64k y 128k. Los ahorros de memoria se mantienen cerca de 3.43x.
El selector aprendido supera la heurística de FastKV en la misma capa. Incluso la selección random-K puntúa más alto en el backbone de TokenGate. El equipo de investigación lo atribuye a la co-adaptación del backbone.
¿Funciona en modelos de atención híbrida?
Sí, Qwen3.5-2B mezcla capas Gated DeltaNet y de atención completa. En NIAH 32k, TokenGate en la capa 3 igualó la precisión del prefill completo. Corrió 3.8x más rápido que el prefill completo y 2.1x más rápido que FastKV. La aplicación de compuertas en las capas 1 o 2 alcanzó 7.8x, con una precisión dentro del 6%. FastKV solo puede podar capas de atención.
TokenGate vs FastKV, SlimInfer, GemFilter y DASH
| Característica | TokenGate | FastKV | SlimInfer | GemFilter | DASH |
|---|---|---|---|---|---|
| Selección de tokens | Puerta MLP aprendida | Atención de las últimas consultas | Atención por bloques, multip Etapa | Atención de última consulta en la capa de filtro | Estabilidad de estados ocultos |
| Entrenamiento necesario | Sí (puerta + LoRA r=16) | No | No | No | No |
| Capa de poda en el Qwen3-1.7B de TokenGate ejecuta | 1 | 13 (predeterminado) | 8, 17, 25 | 12 | 11 |
| Poda en las capas Gated DeltaNet | Sí | No (solo capas de atención) | No divulgado | No divulgado | No divulgado |
| Mejor ganancia de prefill autoinformada (configuración propia) | 11.3x TTFT, 3.4x memoria | 1.82x prefill | 2.53x TTFT (LLaMA3.1-8B, RTX 4090) | 2.4x velocidad, 30% menos memoria de GPU frente al SOTA | No divulgado |
| Código | GitHub | GitHub | GitHub | GitHub | Anunciado |
Las mejoras autoinformadas usan configuraciones diferentes, por lo que no son directamente comparables. Las podas de capas provienen de TokenGate’s Apéndice C.2.
¿Cómo se despliega TokenGate?
El Repositorio de GitHub necesita Python 3.10+ y CUDA PyTorch 2.9+. Un kernel opcional de key-bias de FlashAttention-2 acelera el Qwen3 denso. Cuesta entre 1.0x y 1.1x una llamada simple. Sin él, la inferencia recurre a FlexAttention.
Un {state} marca el documento podable. generate toma un max_kept_prefill_tokens presupuesto. Los checkpoints cubren las 5 tareas en Qwen3-1.7B, además de una puerta NIAH de Qwen3.5-2B. Ejemplo: liana313/tokengate-niah-qwen3-1.7b.
El entrenamiento refleja una receta de TRL/PEFT. Cada paso es 1.3x más lento que un LoRA simple. Las ejecuciones convergieron más rápido en 4 de 5 tareas. El tiempo medio de entrenamiento fue de 1.8 frente a 2.2 horas de A100.
Conclusiones clave
- TokenGate añada poda aprendida de tokens a los LLM existentes mediante una pequeña puerta y LoRA.
- El prefill corre entre 1.4x y 11.3x más rápido con una memoria pico entre 1.5x y 3.4x menor.
- La atención con puerta por token hace que la selección top-K sea entrenable, al estilo de un router MoE.
- Supera la siguiente mejor línea base de poda por hasta 6.4x en TTFT.
- Funciona en el Qwen3.5-2B híbrido, donde FastKV no puede podar las capas DeltaNet.
Echa un vistazo al Artículo, Repositorio de GitHub y página del artículo en Hugging Face. Además, no dudes en seguirnos en Twitter y no olvides unirte a nuestro ML SubReddit de más de 150k y Suscríbete a nuestro Boletín. ¡Espera! ¿estás en telegram? ahora también puedes unirte a nosotros en telegram.
¿Necesitas asociarte con nosotros para promocionar tu Repo de GitHub O Página de Hugging Face O Lanzamiento de Producto O Webinar, etc.? Conéctate con nosotros
La entrada Investigadores de Stanford y UC Berkeley Presentan TokenGate: La Atención con Compuertas por Token Hace el Prefill de Contexto Largo Hasta 11.3x Más Rápido apareció primero en MarkTechPost.