En Gemma 4, la predicción de varios tokens (MTP) es la arquitectura específica que se usa para habilitar la decodificación especulativa altamente eficiente. La decodificación especulativa es una técnica para acelerar la inferencia en modelos de lenguaje grandes. En lugar de depender únicamente del modelo objetivo grande para generar tokens de forma autorregresiva (generando un token a la vez, donde cada token nuevo depende de los anteriores), un "modelo de borrador" más pequeño y rápido predice varios tokens por adelantado. Luego, el modelo de destino verifica estos tokens preliminares en paralelo. Si el modelo objetivo rechaza un token borrador, igual produce el token correcto para esa posición (lo que garantiza que no se desperdicie ese paso), y el modelo borrador reanuda la predicción a partir de ese nuevo token correcto.
Gemma 4 implementa MTP extendiendo el modelo base con este modelo de borrador más pequeño y rápido. Este modelo de borrador no es independiente, ya que comparte la tabla de incorporación de entrada con el modelo objetivo y se basa directamente en las activaciones de su última capa. Esto genera aceleraciones significativas en la velocidad de decodificación y, al mismo tiempo, garantiza la misma calidad que la generación autorregresiva estándar, lo que hace que estos puntos de control sean perfectos para aplicaciones de baja latencia e integradas en el dispositivo.
La decodificación especulativa funciona redactando varios tokens y verificándolos en un solo pase hacia adelante. En el caso de los modelos densos, se usan los mismos pesos para cada token, por lo que verificar varios tokens borrador agrega una sobrecarga mínima. Los modelos de mezcla de expertos (MoE), como Gemma 4 26B A4B, funcionan de manera diferente. Cada token puede activar diferentes expertos, por lo que verificar los tokens borrador puede requerir cargar pesos de expertos adicionales desde la memoria, lo que compensa las ganancias de la creación de borradores. Con tamaños de lote más grandes, suele haber más superposición en los expertos activados en las secuencias, lo que mejora la reutilización de los pesos cargados. Con un tamaño de lote de 1, esta superposición es limitada, por lo que es posible que el borrador de A4B de 26B no genere aceleraciones en plataformas de hardware sin un buen paralelismo.
Mejoras en MTP
Gemma 4 introduce varias mejoras en la canalización de decodificación especulativa estándar para mejorar la calidad de los tokens redactados y la eficiencia:
- Incorporaciones de entrada compartidas: El modelo de borrador comparte la tabla de incorporaciones de entrada con el modelo objetivo.
- Activaciones objetivo: El modelo de borrador usa las activaciones de la última capa del modelo objetivo, las concatena con las incorporaciones de tokens y las proyecta hacia abajo en la dimensión del modelo de borrador.
- Efficient Embedder: Para evitar la costosa operación de predecir en todo el vocabulario, el modelo agrupa los tokens similares en clústeres. Primero, identifica los clústeres más probables y, luego, restringe sus cálculos finales solo a los tokens dentro de esos clústeres seleccionados (solo para E2B y E4B).