Investigadores de Stanford y UC Berkeley lanzaron TokenGate, un módulo para prellenado de contexto largo. TokenGate aprende la poda de tokens que permite a un LLM existente escanear su entrada. Solo los tokens que importan reciben procesamiento completo. El equipo de investigación informa 1.4x a 11.3x más rápido tiempo-hasta-el-primer-token (TTFT) sin pérdida de precisión. La memoria de prellenado pico disminuye 1.5x a 3.4x. El código y los puntos de control están en GitHub bajo Apache-2.0.
TL;DR
- MLP gate: Una pequeña puerta MLP y LoRA (r=16) agregadas a Qwen3-1.7B o Qwen3.5-2B. Probado en contextos de hasta 128k tokens.
- Se ejecuta en: Los tiempos del papel utilizan 1 A100 con tamaño de lote 1. El entrenamiento utilizó 4 A100s.
- Rendimiento: Óptimo de Pareto en 5 tareas contra 6 líneas base.
- Mejor: NIAH 16k a 11.3x más rápido TTFT y 3.42x menos memoria. La ganancia de TTFT alcanza 13.0x a 128k.
- Peor: La aceleración disminuye a 1.4x en algunas tareas.
¿Qué es TokenGate?
TokenGate es un componente integrado que agrega poda de tokens trainable a LLMs preentrenados. Los transformadores empujan cada token a través de 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 altas tasas de poda. Los cambios arquitectónicos a menudo necesitan preentrenamiento desde cero.
¿Cómo funciona la poda de tokens aprendida de TokenGate?
El pase hacia adelante se divide en 3 fases:
- Pase hacia adelante estándar: las capas debajo de la capa de la puerta L procesan los N tokens como de costumbre.
- Módulo de puerta: un MLP puntuación cada token desde su residual de post-atención en la capa L. Un selector de top-K mantiene los tokens con la puntuación más alta.
- Pase hacia adelante podado: las capas por encima de L procesan solo esos K tokens y nunca crean entradas KV para los tokens podados.
Solo la parte del documento de la solicitud es podable. La pregunta nunca se poda.
¿Qué es la atención con puerta de token?
La selección de top-K no es diferenciable, lo que bloquea el entrenamiento de extremo a extremo. TokenGate toma prestada la trampa del enrutador de Expertos en Mezcla. La puntuación de cada token conservado se convierte en un sesgo en los logits de atención por encima de la puerta. Ese sesgo da a las gradientes una ruta de regreso al puntuador.
Una ganancia por cabeza aprendible y una temperatura aprendible refinan esto. La temperatura corrige para cada consulta que ve menos claves. El entrenamiento utiliza una pérdida de token de próximo token plano. Cada paso promedia una pérdida de contexto completo y una pérdida de presupuesto aleatorio. Entonces, 1 puerta entrenada funciona en varios presupuestos.
¿Cómo se desempeña TokenGate en las pruebas de referencia?
El equipo de investigación probó Qwen3-1.7B en RepoBench-P, 2WikiMultihopQA, MultiNews, BioDEX y needle-in-a-haystack (NIAH). Las líneas de base fueron FastKV, SlimInfer, DASH, GemFilter, FlexPrefill y MInference, en el mismo esqueleto LoRA.
- En comparación con el prellenado completo: 1,4 a 11,3 veces más rápido TTFT y 1,5 a 3,4 veces menos memoria.
- En comparación con la mejor línea de base: 1,2 a 6,4 veces más rápido TTFT y hasta 2,1 veces menos memoria.
- 2WikiMultihopQA: al 95% de la precisión de contexto completo, el prellenado se ejecuta 3,0 veces más rápido con 2,7 veces menos memoria.
- Contextos más largos: una puerta entrenada a 16k alcanza 12,8 veces a 32k y 13,0 veces a 64k y 128k. Los ahorros de memoria se mantienen cerca de 3,43 veces.
El selector aprendido supera la heurística de FastKV en la misma capa. Incluso la selección de random-K obtiene una puntuación más alta en el esqueleto TokenGate. El equipo de investigación atribuye la coadaptación del esqueleto.
¿Funciona en modelos de atención híbrida?
Sí, Qwen3.5-2B combina Gated DeltaNet y capas de atención completa. En NIAH 32k, TokenGate en la capa 3 coincidió con la precisión de prellenado completo. Se ejecutó 3,8 veces más rápido que el prellenado completo y 2,1 veces más rápido que FastKV. La puerta en las capas 1 o 2 alcanzó 7,8 veces, dentro del 6% de precisión. 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 token | Puerta MLP aprendida | Atención desde las últimas consultas | Atención de bloque, multi-etapa | Atención de última consulta en la capa de filtro | Estabilidad del estado oculto |
| Entrenamiento necesario | Sí (puerta + LoRA r=16) | No | No | No | No |
| Capa de poda en TokenGate’s Qwen3-1.7B | 1 | 13 (predeterminado) | 8, 17, 25 | 12 | 11 |
| Poda en capas Gated DeltaNet | Sí | No (solo capas de atención) | No divulgado | No divulgado | No divulgado |
| Mejor beneficio de prellenado autoinformado (configuración propia) | 11.3x TTFT, 3.4x memoria | 1.82x prellenado | 2.53x TTFT (LLaMA3.1-8B, RTX 4090) | 2.4x velocidad, 30% menos memoria de GPU en comparación con SOTA | No divulgado |
| Código | GitHub | GitHub | GitHub | GitHub | Anunciado |
Los beneficios autoinformados utilizan configuraciones diferentes, por lo que no son directamente comparables. Las capas de poda provienen de la Apéndice C.2 de TokenGate.
¿Cómo desplegar TokenGate?
El repositorio de GitHub requiere Python 3.10+ y CUDA PyTorch 2.9+. Un kernel de clave-sesgo FlashAttention-2 opcional acelera la densa Qwen3. Cuesta 1.0x a 1.1x una llamada simple. Sin él, la inferencia vuelve a FlexAttention.
Un marcador de posición {state} marca el documento podable. generate toma un presupuesto max_kept_prefill_tokens_. Los puntos de control cubren las 5 tareas en Qwen3-1.7B, más una puerta NIAH Qwen3.5-2B. Ejemplo: liana313/tokengate-niah-qwen3-1.7b.
El entrenamiento es un espejo de una receta TRL/PEFT. Cada paso es 1.3x más lento que el LoRA simple. Las corridas convergieron más rápido en 4 de 5 tareas. El tiempo de entrenamiento promedio fue 1.8 versus 2.2 horas A100.
Puntos clave
- TokenGate agrega la poda de tokens aprendida a los LLM existentes a través de una pequeña puerta y LoRA.
- Las corridas de prellenado son 1.4x a 11.3x más rápidas con 1.5x a 3.4x menos memoria pico.
- La atención con puerta de tokens hace que la selección de top-K sea trainable, al estilo de MoE-router.
- Superó la línea de base de poda más cercana hasta 6.4x en TTFT.
- Funciona en híbrido Qwen3.5-2B, donde FastKV no puede podar capas DeltaNet.
Revisa el Paper, repositorio de GitHub y página del paper en Hugging Face. También, no dudes en seguirnos en Twitter y no te olvides de unirte a nuestro SubReddit de 150k+ML y suscribirte a nuestro boletín informativo. ¡Espera! ¿Estás en Telegram? ahora también puedes unirte a nosotros en Telegram.
¿Necesitas asociarte con nosotros para promocionar tu repositorio de GitHub OR página de Hugging Face OR lanzamiento de producto OR webinar, etc.? Conéctate con nosotros
La publicación Investigadores de Stanford y UC Berkeley presentan TokenGate: Token-Gated Attention hace que el prellenado de contexto largo sea hasta 11.3x más rápido apareció por primera vez en MarkTechPost.