El Prism Transformer introduce un horario progresivo de cabezas que varía el número de cabezas de atención entre capas, comenzando con menos cabezas más anchas en las primeras capas y aumentando la cantidad monótonamente con la profundidad. Este enfoque desafía la asignación uniforme estándar al abordar las necesidades estructurales distintas de las primeras versus las últimas capas sin añadir sobrecarga arquitectónica.

  • Las primeras capas utilizan subespacios amplios por cabeza (dh=256) para una captura rica de patrones locales, mientras que las últimas capas usan muchas cabezas estrechas para descomposición especializada.
  • Las matrices de peso conservan formas estándar dmodel×dmodel, manteniendo neutral el conteo de parámetros.
  • Los FLOPs totales permanecen matemáticamente invariantes al conteo de cabezas, asegurando neutralidad computacional.
  • Las dimensiones de cabeza de potencia de 2 (dh ∈ {256, 128}) preservan la alineación con Tensor Core para neutralidad de rendimiento.
  • Los resultados muestran menor pérdida de validación en cada escala en comparación con las bases uniformes con tokens/segundo y tiempo de pared idénticos.
  • El modelo logra ganancias o paridad en benchmarks incluyendo PIQA, WinoGrande, HellaSwag y ARC-Easy.

El análisis de distancia de atención por capa confirma que la ganancia es estructural, ya que las primeras capas de Prism atienden más localmente antes de cambiar a integración global. La implementación requiere solo un cambio de una línea por capa de atención para hacer el conteo de cabezas dependiente de la capa.