SpatialLM-Mini accélère l'entraînement sur un contexte de 16K d'un facteur 2,47 sur une RTX 3050

Un long contexte coûte cher avant même que le modèle commence à être utile. Avec l'attention complète, chaque nouveau token peut s'adresser à tous les tokens précédents, et le volume de calcul croît de façon quadratique. Dans l'expérience SpatialLM-Mini, cet accès a été restreint : on a conservé les positions les plus proches et quelques liens lointains. Sur une RTX 3050 dotée de 8 Go de mémoire, le prototype s'est entraîné 2,47 fois plus vite que le Transformer de référence avec une fenêtre de 16 384 tokens.

Pour un développeur, il y a ici une piste concrète : réduire le nombre de liens peut accélérer les expériences sur de longues séquences, même sur une carte graphique peu coûteuse. Mais le résultat concerne un petit modèle, un seul corpus et un noyau de calcul précis. L'intérêt de la méthode « spatiale » de calcul de l'attention reste, à ce stade, une question ouverte.

Comment les tokens choisissent qui écouter

Imaginez un long document. Dans l'attention causale classique, un token situé à la fin peut tenir compte directement de n'importe quelle position précédente. Dans SpatialLM-Mini, la liste des positions accessibles est fixée à l'avance : le token lui-même, quelques prédécesseurs immédiats et les positions situées à des distances de 1, 2, 4, 8, et ainsi de suite. Les accès au futur sont interdits.

Ensuite, le modèle décide de l'attention à accorder à chaque voisin autorisé. Pour cela, l'état caché est transformé en coordonnées tridimensionnelles apprenables pour la requête et la clé. Plus la distance entre elles est faible, plus le poids du lien peut être élevé ; des paramètres apprenables supplémentaires ajustent ce calcul.

Les coordonnées servent ici de moyen de calculer les poids. Le modèle ne cherche pas les tokens les plus proches dans l'ensemble du document et n'utilise pas les cœurs RT de la carte graphique. La liste des voisins est fixe, tandis que la répartition de l'attention au sein de cette liste dépend du texte.

Un tel graphe réduit les calculs, mais limite l'accès direct aux faits éloignés. C'est pourquoi l'auteur a aussi testé un hybride : deux blocs spatiaux alternent avec deux blocs d'attention complète. Cela permet à une partie des couches d'accéder à tout le contexte précédent.

Ce qu'a montré la comparaison sur 8 Go

Trois modèles d'environ 9,3 millions de paramètres chacun ont été entraînés à partir de zéro sur « Orgueil et Préjugés » de Jane Austen. Chacun avait quatre couches, le même contexte et trois exécutions avec des états aléatoires initiaux différents. De GPT-2, seul le tokenizer a été utilisé.

Avec un budget de 180 secondes d'entraînement mesuré, le Transformer de référence traitait 13 126 tokens par seconde, le modèle spatial 32 364 et l'hybride 18 725. L'hybride a obtenu l'erreur moyenne la plus basse sur la partie de test. Toutefois, dans ce laps de temps, les modèles rapides ont pu traiter plus de tokens et effectuer davantage de mises à jour.

Lorsqu'on compare après le même nombre de 100 pas, l'écart d'erreur se réduit nettement. Le modèle spatial a eu besoin de 50,64 secondes, le Transformer de 125,10 secondes. Le principal gain confirmé de cette expérience est la vitesse des mises à jour sur une longue séquence.

Pour la mémoire, c'est différent : le pic de mémoire allouée par PyTorch est passé d'environ 499 à 671 Mio. L'implémentation parcimonieuse crée des tenseurs intermédiaires lors de la collecte des voisins, tandis que l'attention complète utilise un noyau optimisé. Ces chiffres ne reflètent pas toute la mémoire vidéo occupée de l'ordinateur.

Ce qu'il faut vérifier avant de l'adopter dans son projet

Le Transformer de référence fonctionnait via PyTorch Efficient Attention. FlashAttention était absent de la build Windows utilisée. Le facteur d'accélération ne peut donc pas être transposé automatiquement à un autre système ni à une comparaison avec d'autres noyaux.

Il existe aussi une lacune plus importante : on a modifié simultanément l'ensemble des liens et la méthode de calcul de leurs poids. Pour comprendre la contribution des coordonnées, il faut le même graphe parcimonieux avec un produit scalaire ordinaire ou des poids fixes. Pour l'instant, l'avantage peut s'expliquer avant tout par la réduction du nombre d'interactions.

L'ingestion de 16 000 tokens ne prouve pas non plus que le modèle sait retrouver un fait précis au début du document. Pour cela, un test distinct est nécessaire. La vitesse de génération n'a pas été mesurée, le prototype n'a pas de cache KV, et un entraînement court sur un seul roman ne permet pas de juger de l'aptitude à des tâches de production.

La prochaine étape pratique consiste à reproduire la comparaison sur votre propre corpus et votre matériel, puis à ajouter un contrôle avec le même graphe sans coordonnées. Selon l'auteur, le code, les configurations et les mesures brutes sont disponibles dans le dépôt SpatialLM-Mini. Pour une équipe qui étudie l'entraînement économique à long contexte, c'est un point de départ utile pour sa propre expérience.