SpatialLM-Mini aceleró el entrenamiento con un contexto de 16K en 2,47 veces en una RTX 3050

El contexto largo resulta costoso incluso antes de que el modelo empiece a aportar valor. Con atención completa, cada nuevo token puede atender a todos los anteriores y el volumen de cálculo crece de forma cuadrática. El experimento SpatialLM-Mini restringió ese acceso: conservó las posiciones cercanas y unas pocas conexiones de largo alcance. En una RTX 3050 con 8 GB de memoria, el prototipo se entrenó 2,47 veces más rápido que el Transformer de referencia con una ventana de 16 384 tokens.

Para los desarrolladores, hay una pista concreta: reducir el número de conexiones puede acelerar los experimentos con secuencias largas incluso en una tarjeta gráfica económica. Sin embargo, el resultado corresponde a un modelo pequeño, un único corpus y un kernel de cálculo específico. La utilidad del método propiamente «espacial» para calcular la atención sigue siendo una cuestión abierta.

Cómo eligen los tokens a quién escuchar

Imaginemos un documento largo. Con la atención causal habitual, un token situado al final puede tener en cuenta directamente cualquier posición anterior. En SpatialLM-Mini, la lista de posiciones accesibles se define de antemano: el propio token, varios predecesores cercanos y posiciones a distancias de 1, 2, 4, 8 y así sucesivamente. No se permite atender a posiciones futuras.

Después, el modelo decide cuánta atención dedicar a cada vecino permitido. Para ello, transforma el estado oculto en coordenadas tridimensionales aprendibles de consulta y clave. Cuanto menor sea la distancia entre ellas, mayor puede ser el peso de la conexión; otros parámetros aprendibles ajustan este cálculo.

Las coordenadas sirven aquí para calcular los pesos. El modelo no busca los tokens más cercanos por todo el documento ni utiliza los núcleos RT de la tarjeta gráfica. La lista de vecinos es fija, mientras que la distribución de la atención dentro de ella depende del texto.

Este grafo reduce el cálculo, pero limita el acceso directo a datos lejanos. Por eso, el autor también probó un híbrido: dos bloques espaciales se alternan con dos bloques de atención completa. Así, algunas capas pueden atender a todo el contexto anterior.

Qué mostró la comparación con 8 GB

Se entrenaron desde cero tres modelos de unos 9,3 millones de parámetros cada uno con Orgullo y prejuicio, de Jane Austen. Cada modelo tenía cuatro capas, la misma longitud de contexto y tres ejecuciones con distintas semillas aleatorias. De GPT-2 solo se utilizó el tokenizador.

Con un presupuesto de 180 segundos de entrenamiento medido, el Transformer de referencia procesó 13 126 tokens por segundo; el modelo espacial, 32 364; y el híbrido, 18 725. El híbrido obtuvo la menor pérdida media en el conjunto de prueba. Sin embargo, durante ese tiempo los modelos más rápidos pudieron procesar más tokens y realizar más actualizaciones.

Al comparar los resultados tras los mismos 100 pasos, la diferencia en la pérdida se redujo considerablemente. El modelo espacial necesitó 50,64 segundos; el Transformer, 125,10 segundos. La principal mejora confirmada en este experimento es la velocidad de las actualizaciones con una secuencia larga.

Con la memoria ocurrió algo distinto: el pico de memoria asignada por PyTorch aumentó de unos 499 a 671 MiB. La implementación dispersa crea tensores intermedios al recopilar los vecinos, mientras que la atención completa utiliza un kernel optimizado. Estas cifras no reflejan toda la memoria de vídeo ocupada en el equipo.

Qué comprobar antes de incorporarlo a un proyecto

El Transformer de referencia utilizó PyTorch Efficient Attention. FlashAttention no estaba disponible en la compilación de Windows empleada. Por tanto, el factor de aceleración no puede extrapolarse automáticamente a otro sistema ni a una comparación con otros kernels.

También hay una laguna más importante: se cambiaron simultáneamente el conjunto de conexiones y el método para calcular sus pesos. Para determinar la aportación de las coordenadas, hace falta utilizar el mismo grafo disperso con un producto escalar convencional o con pesos fijos. Por ahora, la ventaja podría deberse principalmente a la reducción del número de interacciones.

Aceptar 16 mil tokens tampoco demuestra que el modelo sepa recuperar el dato pertinente del principio del documento. Eso requiere una prueba independiente. No se midió la velocidad de generación, el prototipo carece de caché KV y un entrenamiento breve con una sola novela no permite valorar su idoneidad para tareas de producción.

El siguiente paso práctico es repetir la comparación con un corpus y un equipo propios y añadir después una referencia con el mismo grafo, pero sin coordenadas. Según la descripción del autor, el código, las configuraciones y las mediciones originales están disponibles en el repositorio SpatialLM-Mini. Para un equipo que investigue el entrenamiento económico con contextos largos, es un punto de partida útil para realizar su propio experimento.