Training SAEs

Je me souviens assez clairement du moment où j'ai lu Scaling Monosemanticity par Templeton et al. en 2024. C'est en explorant Transformer Circuits que j'ai découvert la recherche en interpretabilité, et comme beaucoup d'autres j'ai trouvé la possibilité de piloter (steering) les LLMs particulièrement intriguante — et amusante, cf. Golden Gate Claude. Un peu plus d'un an après et quelques lectures en plus j'ai donc cherché à répliquer ces expériences, et entraîner moi-même un Sparse Autoencoder, ou SAE, pour faire de l'apprentissage par dictionnaire. Cet article détaille mes travaux dans ce but, les différentes optimisations que j'ai mises en place (kernels triton pour l'entraînement, inférence des activations plus rapide, ...) et quelques résultats de steering.

Cadre théorique

Transformers, caractéristiques et polysémie

Comme son nom l'indique, l'interprétabilité mécaniste cherche à comprendre le comportement d'un réseau de neurones en y explicitant des mécanismes internes qui lui font donner tel résultat. Une approche classique pour traiter ce type de questions est d'essayer de « reverse engineer » l'objet d'intérêt, en le décomposant en des parties plus petites et plus prédictibles, et en comprenant comment celles-ci intéragissent. Commençons donc par rappeler brièvement l'architecture des LLM récents. Dans les grandes lignes, elle est assez similaire entre ces derniers, et basée sur la partie décodeur du transformer^{[Ashish Vaswani, 2023]}Attention Is All You NeedAshish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N. Gomez and Lukasz Kaiser and Illia Polosukhin (2023)Source.

• un texte d'entrée est découpé en $N$ tokens, pour lesquels on donne un embedding $d_{\textrm{model}}$ qui dépend du token et de la position — on a donc un tenseur $x_0$ de forme $(N, d_{\textrm{model}})$ ;
• Pour chaque couche $0 \leq i < L$ du transformer, on lit $x_{i}$ et calcule $x_{i+1} = \ell_i(x_i)$, également de forme $(N, d_{\textrm{model}})$ ;
• en sortie de la dernière couche, on normalise et projette $x_L$ pour obtenir les logits utilisés pour prédire le token suivant.

Chaque couche $i$ est composée de deux blocs : un bloc d'attention et un bloc feed-forward (pour les modèles dense — c'est un peu différent pour les mixtures d'experts), chacun avec normalisation. Depuis GPT-2^{[Radford, 2019]}Language Models are Unsupervised Multitask LearnersRadford, Alec and Wu, Jeffrey and Child, Rewon and Luan, David and Amodei, Dario and Sutskever, Ilya (2019)OpenAI.Source, la normalisation par couche se fait généralement en entrée de chaque bloc, plutôt qu'en sortie comme dans le transformer original^{[Ashish Vaswani, 2023]}Attention Is All You NeedAshish Vaswani and Noam Shazeer and Niki Parmar and Jakob Uszkoreit and Llion Jones and Aidan N. Gomez and Lukasz Kaiser and Illia Polosukhin (2023)Source.

On remarquera que chaque bloc "met à jour" $(x_i)_i$, qu'on appele d'ailleurs le flux résiduel du modèle. Il faut également noter qu'ils interagissent linéairement^{[Elhage, 2021]}A Mathematical Framework for Transformer CircuitsElhage, Nelson and Nanda, Neel and Olsson, Catherine and Henighan, Tom and Joseph, Nicholas and Mann, Ben and Askell, Amanda and Bai, Yuntao and Chen, Anna and Conerly, Tom and DasSarma, Nova and Drain, Dawn and Ganguli, Deep and Hatfield-Dodds, Zac and Hernandez, Danny and Jones, Andy and Kernion, Jackson and Lovitt, Liane and Ndousse, Kamal and Amodei, Dario and Brown, Tom and Clark, Jack and Kaplan, Jared and McCandlish, Sam and Olah, Chris (2021)Transformer Circuits Thread.Source avec ce flux résiduel, au sens où :

• leur entrée est obtenue (à normalisation près) par des transformations linéaires du flux résiduel (matrices Q, K, V des têtes d'attention, première couche linéaire du FFN);
• leur sortie, juste avant addition au flux résiduel, est précédée par une transformation linéaire (concaténation des têtes d'attention, dernière couche linéaire du FFN).

On peut donc avoir l'intuition que les représentations que se fait le modèle se constituent progressivement au fil des couches, et par ailleurs qu'elles respectent aussi une certaine linéarité : si $x$ représente une voiture, $2x$ ne devrait pas représenter un chat. Les représentations sont donc des directions dans l'espace latent. Par ailleurs, la linéarité doit aussi impliquer une certaine additivité des représentations, et on peut par exemple espérer que [la représentation de] reine) soit proche de roi - homme + femme, ce type de régularités linguistiques étant déjà observées depuis 2013^{[Mikolov, 2013]}Linguistic Regularities in Continuous Space Word RepresentationsMikolov, Tomas and Yih, Scott Wen-tau and Zweig, Geoffrey (2013)Source.

Les représentations, qu'on appelle aussi caractéristiques (ou features) du modèle sont des directions dans l'espace latent $\mathbb{R}^{d_{\textrm{model}}}$. Or, il n'y a a priori pas de raisons de penser que les représentations les plus explicites aient comme directions celles de la base canonique, lesquelles correspondent aux sorties des neurones. Aussi, un même neurone peut très bien intervenir dans des contextes très différents (par exemple^{[Olah, 2017]}Feature VisualizationOlah, Chris and Mordvintsev, Alexander and Schubert, Ludwig (2017)Distill.DOI: 10.23915/distill.00007 un neurone d'un modèle de vision actif sur des visages d'animaux ou des voitures). Similairement, un même token $x_{i,n}$ peut traiter de plusieurs notions en même temps. Dans ce cas, on parlera parfois de de superposition^{[Olah, 2020]}Zoom In: An Introduction to CircuitsOlah, Chris and Cammarata, Nick and Schubert, Ludwig and Goh, Gabriel and Petrov, Michael and Carter, Shan (2020)Distill.DOI: 10.23915/distill.00024.001, et de caractéristiques polysémiques/polysémantiques.

Apprentisage de dictionnaire

Pour revenir à la question intiale, les neurones ne sont donc pas forcément facilement interprétables. On peut donc chercher à trouver des caractéristiques qui le soient, par exemple des caractéristiques monosémantiques. En quelque sorte, on cherche à "démêler" (disentanglement) les représentations. Pour cela, une approche consiste à faire de l'apprentisage de dictionnaire : pour un token $t_n = x_{i,n} \in \mathbb{R}^{d_{\text{model}}}$ en sortie de la couche $i-1$, on va chercher un vecteur $\alpha_n \in \mathbb{R}^h$ tel que

où $D \in \mathbb{R}^{d_{\text{model}}\times h}$ est le dictionnaire, $h \gg d_{\text{model}}$ et $\alpha_n$ est très parcimonieux/sparse (autrement dit $t$ est une combinaison linéaire de quelques colonnes de $D$). C'est une méthode utilisée en traitement du signal, puisqu'il permet de représenter un signal avec un nombre minimal de composantes (on pourrait faire une analogie avec la décomposition en série de Fourier, où $\alpha_n$ correspondrait ici aux amplitudes / phases de chaque harmonique).

On a donc approximativement le problème d'optimisation suivant

où ${||\cdot||}_0$ est la pseudo-norme 0 qui compte le nombre de valeurs non-nulles, et où $\lambda > 0$ arbitre donc entre la parcimonie/sparsité et la qualité de l'approximation, c.-à.-d. l'erreur entre $\widehat{t_n}$ et $t_n$.

Il existe plusieurs algorithmes permettant de résoudre ce problème. Cependant, celui-ci est non-convexe (à cause de la pseudo-norme 0) et NP-difficile^{[Tillmann, 2015]}On the Computational Intractability of Exact and Approximate Dictionary LearningTillmann, Andreas M. (2015)IEEE Signal Processing Letters, vol. 22(1), pp. 45-49.DOI: 10.1109/LSP.2014.2345761. En conséquence, le résoudre est potentiellement "trop fort" (voir les arguments de Tom Henighan et Chris Olah^{[Henighan, 2023]}Circuits Updates, May 2023, Dictionary Learning WorriesHenighan, Tom and Olah, Chris (2023)Source) : on risque de trouver des caractéristiques que le modèle n'utilise pas, car trop coûteuses à calculer (et donc trop coûteuses d'accès pour le modèle). Il y a donc un risque que notre dictionnaire overfit sur les données et n'extraie pas correctement les représentations que le modèle manipule. Ces méthodes ne sont également pas les plus simples pour passer à l'échelle^{[Bricken, 2023]}Towards Monosemanticity: Decomposing Language Models With Dictionary LearningBricken, Trenton and Templeton, Adly and Batson, Joshua and Chen, Brian and Jermyn, Adam and Conerly, Tom and Turner, Nick and Anil, Cem and Denison, Carson and Askell, Amanda and Lasenby, Robert and Wu, Yifan and Kravec, Shauna and Schiefer, Nicholas and Maxwell, Tim and Joseph, Nicholas and Hatfield-Dodds, Zac and Tamkin, Alex and Nguyen, Karina and McLean, Brayden and Burke, Josiah E and Hume, Tristan and Carter, Shan and Henighan, Tom and Olah, Christopher (2023)Transformer Circuits Thread.Source, or l'on souhaite idéalement traiter un corpus suffisament large (en milliards de tokens).

Sparse Autoencoders

Une autre approche consiste à entraîner un Sparse Autoencoder. Chaque token $t_n$, passe dans un encodeur $\alpha_n = \text{enc}(t_n) \in \mathbb{R}^{d_{\text{latent}}}$, puis dans un decodeur qui reprojette dans la dimension d'origine $\widehat{t_n} = \text{dec}(f_n) \in \mathbb{R}^{d_{\text{model}}}$. Classiquement, il s'agit d'un réseau de neurones qu'on peut entraîner par rétro-propagation :

On cherche à minimiser l'erreur entre $\widehat{t_n}$ et $t_n$ et à ce que $\alpha_n$ soit le plus sparse possible. On se ramène ainsi à une forme faible d'apprentissage de dictionnaire. Au biais du decoder près, $\widehat{t_n}$ a d'ailleurs la même forme, et le dictionnaire $D$ correspond aux poids du decodeur. On peut donc espérer que les neurones $f_1, f_2, \dots$ de la couche latente, autrement dit des dimensions / direction canonique de $\mathbb{R}^{d_{\text{latent}}}$, correspondent à des caractéristiques monosémantiques, et soient des représentations manipulées par le modèle. Notamment, l'architecture du bloc feed-forward des LLM est très similaire à celle du SAE que l'on vient de décrire.

Il y a plusieurs variations de SAE, et les premières itérations semblent dater de décembre 2022^{[Sharkey, 2022]}Interim Research Report: Taking Features Out of Superposition with Sparse AutoencodersSharkey, Lee and Braun, Dan and beren (2022)Source, notamment autour de la fonction d'activation en sortie de l'encodeur. On peut par exemple avec une $\text{ReLU}$ et imposer la sparsité au niveau de la loss (comme dans le problème d'optimisation de l'aprentissage de dictionnaire). Il s'agit de l'approche des chercheurs d'Anthropic documentée dans Scaling Monosemanticy^{[Templeton, 2024]}Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 SonnetTempleton, Adly and Conerly, Tom and Marcus, Jonathan and Lindsey, Jack and Bricken, Trenton and Chen, Brian and Pearce, Adam and Citro, Craig and Ameisen, Emmanuel and Jones, Andy and Cunningham, Hoagy and Turner, Nicholas L and McDougall, Callum and MacDiarmid, Monte and Freeman, C. Daniel and Sumers, Theodore R. and Rees, Edward and Batson, Joshua and Jermyn, Adam and Carter, Shan and Olah, Chris and Henighan, Tom (2024)Transformer Circuits Thread.Source (et^{[Bricken, 2023]}Towards Monosemanticity: Decomposing Language Models With Dictionary LearningBricken, Trenton and Templeton, Adly and Batson, Joshua and Chen, Brian and Jermyn, Adam and Conerly, Tom and Turner, Nick and Anil, Cem and Denison, Carson and Askell, Amanda and Lasenby, Robert and Wu, Yifan and Kravec, Shauna and Schiefer, Nicholas and Maxwell, Tim and Joseph, Nicholas and Hatfield-Dodds, Zac and Tamkin, Alex and Nguyen, Karina and McLean, Brayden and Burke, Josiah E and Hume, Tristan and Carter, Shan and Henighan, Tom and Olah, Christopher (2023)Transformer Circuits Thread.Source). Des chercheurs d'OpenAI ont introduit^{[Leo Gao, 2024]}Scaling and evaluating sparse autoencodersLeo Gao and Tom Dupré la Tour and Henk Tillman and Gabriel Goh and Rajan Troll and Alec Radford and Ilya Sutskever and Jan Leike and Jeffrey Wu (2024)Source une variante où la sparsité est imposée avec une fonction d'activation $\text{Top}_K$ (on ne conserve que les $K$ valeurs les plus élevées (et positives) de $\alpha_n$). D'autres variations ont également été étudiées (notamment GatedSAE^{[Rajamanoharan, 2024]}Improving Dictionary Learning with Gated Sparse AutoencodersRajamanoharan, Senthooran and Arthur Conmy and Lewis Smith and Tom Lieberum and Vikrant Varma and János Kramár and Rohin Shah and Neel Nanda (2024)Source et JumpReLU SAE^{[Rajamanoharan, 2024]}Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse AutoencodersRajamanoharan, Senthooran and Tom Lieberum and Nicolas Sonnerat and Arthur Conmy and Vikrant Varma and János Kramár and Neel Nanda (2024)Source).

Entrainement

La première partie ci-dessous détaille les conditions générales d'entrainement. Les deux suivantes donnent des détails sur les optimisations (kernels, ...) que j'ai implémenté pour l'entrainement.

Présentation générale

En premier lieu il convient de choisir le LLM qu'on souhaite étudier. J'ai personnellement choisi de travailler avec Mistral-Small-3.2-24B-Instruct-2506, pour plusieurs raisons.

• C'est un modèle (raisonnablement) petit : 24 milliards de paramètres, donc demande dans les 50 GO de VRAM en BF16 pour l'inférence (en pratique il faudra moins, voir sections optimisations plus bas).
• C'est un modèle dense. Comme mentionné plus haut, les mixtures d'experts ont une architecture un peu différente, et mettent potentiellement à mal les hypothèses formulées. D'autre part, il y a actuellement bien plus de littérature sur les SAE entrainés sur des modèles dense que sur des MoE.
• C'est un modèle "instruct". Une partie du jeu de données que j'ai utilisé était sous forme d'échanges, il était donc pratique de pouvoir formatter ainsi les entrées du modèle.

Il faut également choisir la couche $i$ où récupérer le flux résiduel. Une intuition est que les concepts les plus simples se forment dans les premières couches, et que ceux plus complexes se situent plus loin dans le modèle Un certain nombre d'observations empire vont dans ce sens, et de nombreux travaux placent le SAE au milieu du modèle. Mistral-Small-3 ayant 40 couches, j'ai décidé de placer le SAE en sortie de la 30è couche (on regarde donc $x_{i+1}$ avec $i = 29$).

Dans ma configuration, j'entraîne le SAE de manière "online" (terminologie approximative) : Le LLM et le SAE en entrainement sont donc en même temps sur GPU, et j'infère directement le LLM jusqu'à avoir assez d'activations pour faire un batch avec lequel entraîner le SAE. J'ai utilisé une NVIDIA RTX PRO 6000 Blackwell avec 96 GO de VRAM. mistral-small-3 ayant a une dimension $d_{\text{model}} = 5120$, j'ai choisi d'entraîner mon SAE avec

ce qui fait environ 3B paramètres. J'entraine nativement en BF16 avec AdamW (les poids, gradients et moments représentent donc dans les 24 GO de VRAM), et avec une batchsize de $2048$. Chaque batch est mélangé et normalisé par un facteur $\sqrt{\frac{\bdE{||t_n||}_2^2}{d_{\text{model}}}}$ (basé sur les travaux d'Anthropic). Le LLM est lui inféré sur des batch de $16$ séquences de $256$ tokens.

Ces choix sont principalement motivés par la contrainte en VRAM (s'agissant d'un projet personnel je n'ai pas pu me permettre de meilleure configuration). J'ai éxécuté plusieurs tests afin d'arbitrer et choisir au mieux ces valeurs (cf. benchmarks.md). J'ai cherché à me rapprocher le plus proche possible de Scaling Monosemanticity, aussi la fonction d'activation est une $\text{ReLU}$ et la sparsité est obtenue au niveau de la loss : pour un token $t_n$, on calcule

où $W_{\text{enc}}, W_{\text{dec}}, b_{\text{enc}}, b_{\text{dec}}$ sont les poids et biais de l'encodeur et du decodeur. La loss est

$\lambda > 0$ arbitre entre reconstruction et sparsité. La normalisation décrite plus haut m'a permis de reprendre $\lambda = 5$ des travaux d'Anthropic. $W_{\text{dec},j}$ correspond à la colonne $j$ du decodeur (soit la représention de la caractéristique $f_j$). Il est important de l'intégrer à la pénalité sans quoi le modèle peut baisser artificiellement les valeurs de $\alpha_n$ en augmentant les valeurs du decodeur. Une autre approche^{[Leo Gao, 2024]}Scaling and evaluating sparse autoencodersLeo Gao and Tom Dupré la Tour and Henk Tillman and Gabriel Goh and Rajan Troll and Alec Radford and Ilya Sutskever and Jan Leike and Jeffrey Wu (2024)Source consiste à normaliser les $W_{\text{dec},j}$ à chaque step.

$W_{\text{enc}}$ est initialisé comme la transposée de $W_{\text{dec}}$, lui-même initialisé selon He/Kaiming uniforme ; les biais $b_{\text{enc}}$ et $b_{\text{dec}}$ sont initialisés à 0. AdamW est configuré avec $\beta_1 = 0.9$, $\beta_2 = 0.999$, pas de weight decay et $\varepsilon = 10^{-15}$. J'utilise du gradient clipping, un "warmup" de $\lambda$ de $0$ à $5$ sur les 5 premiers % de l'entrainement, et un learning rate de $5 \times 10^{-5}$ constant puis qui décroit sur les derniers 20 % de l'entrainement.

J'ai effectué des entraînements sur 600 millions, 1.2 milliards et 1.6 milliards de tokens. Pour surveiller le bon déroulement, j'ai calculé et log périodiquement:

• le nombre de valeurs non nulles (nnz) par token, avec des statistiques de répartition (moyenne, médiane, p90, p99) pour caractériser la sparsité effective ;
• les fréquences d'activation des latents, moyenne, maximale, quantiles (p50, p90, p99), ainsi que la proportion de latents dans différentes plages de fréquence (0.1-1%, 1-10%, >10%) ;
• les amplitudes des activations conditionnelles à l'activité, moyenne, écart-type et p99, pour caractériser la distribution des contributions ;
• des métriques sur le résidu de reconstruction, notamment la variance expliquée, la corrélation avec l'entrée et la norme du résidu.

Optimisation de l'inférence (LLM)

Il existe de nombreux travaux pour accélérer l'inférence des modèles de langages. On peut notamment citer PagedAttention du projet vLLM^{[Kwon, 2023]}Efficient Memory Management for Large Language Model Serving with PagedAttentionKwon, Woosuk and Zhuohan Li and Siyuan Zhuang and Ying Sheng and Lianmin Zheng and Cody Hao Yu and Joseph E. Gonzalez and Hao Zhang and Ion Stoica (2023)Source ou FlashAttention de Tri Dao^{[Dao, 2023]}FlashAttention-2: Faster Attention with Better Parallelism and Work PartitioningDao, Tri (2023)Source. Si les optimisations du premier concernent l'utilisation autorégressive et concurrente du modèle (ce qui n'est pas notre cas), celles de FlashAttention nous sont tout à fait utiles, augmentent le débit de tokens par seconde et réduisent la mémoire utilisée.

La manière classique de récupérer les activations/états intermédiaires d'un modèle est d'utiliser des hook ou des bibliothèques comme TransformerLens (qui elle aussi utilise des hook). Dans notre cas on peut aussi utiliser output_hidden_states=True dans le forward d'un modèle chargé avec transformers. Cette approche présente plusieurs défauts.

• Elle infère toutes les couches du modèles, y compris les $i' > i$ situées après le SAE, dont nous n'avons pas besoin.
• Elle charge en VRAM la partie vision du modèle (vision tower et projecteur multimodal) dont nous n'avons pas besoin.
• Elle est peu ou pas compatible avec torch.compile().

On peut plutôt modifier l'implémentation du modèle dans transformers pour retirer la partie multimodale ("langage seul"), et les couches non utiles et normalisation finale ("tranché").

Comme dit plus haut, un autre avantage et qu'on peut compiler le module obtenu. On teste la performance pour différentes méthodes de compilation possibles (avec et sans FlashAttention), puis pour différentes tailles de batch et de séquence. Il faut noter que la compilation demande idéalement d'avoir des tenseurs de forme fixe. On ajoute des tokens de padding et masque d'attention pour les séquences trop courtes, et on tronque les séquences trop longue. Il faut noter que les tokens de padding doivent être retirés après inférence avant d'être donnés au SAE, ce qui crée un surcoût (croissant en la longueur de séquence souhaitée). Au final, on parvient atteint plus du double en débit par rapport à l'implémentation de base.

Optimisation de l'entraînement (SAE)

Le SAE a globalement une structure très simple. Les optimisations à mettre en place tirent principalement profit de son caractère parcimonieux/sparse : pour un token, seulement quelques caractéristiques parmi $f_1, f_2, \dots{}, f_{d_{\text{latent}}}$ sont actives. On peut donc trouver une meilleure manière de représenter la sortie $\alpha_n$ de l’encodeur, et implémenter de manière moins naïve l’encodeur et le decodeur. Commençons par l'encodeur : on a

En pratique, on traite un batch de $B$ tokens. Dans la suite, on notera $\cdot_{ABC}$ un tenseur de forme $(A, B, C)$ ; notamment on notera $W_{HD} = W_{\text{enc}}$ les poids et $b_H = b_{\text{enc}}$ biais de l'encodeur. Pour le forward, notre encodeur effectue

On va donc chercher à représenter de manière sparse $a_{BH}$.

Pour le backward, on cherche à calculer $\frac{\partial\bcL}{\partial W_{HD}}$ le gradient des poids, $\frac{\partial\bcL}{\partial b_{H}}$ celui des bias et potentiellement $\frac{\partial\bcL}{\partial x_{BD}}$ celui des entrées. Pour rappel, on a $a_{BH} = \text{ReLU}(z_{BH})$ donc

où $\odot$ est le produit élément par élément (Hadamard).

Ensuite, $z_{BH} = x_{BD}(W_{HD})^\top + 1_{B}b_H^\top$ nous livre

Puisqu'on a une représentation sparse de $a_{BH}$, il en va de même de son gradient $\frac{\partial \bcL}{\partial a_{BH}}$ et donc du gradient $\frac{\partial \bcL}{\partial z_{BH}}$. Dans notre cas, on n'a pas besoin des gradients sur les entrées $x_{BD}$. On cherche donc à faire la somme d'une matrice sparse le long d'une dimension pour le gradient de $b_H$, et la multuplication de la transposée d'une matrice sparse par une matrice dense pour celui de $W_{HD}$. Pour le premier, un scatter_add fait l'affaire. Pour le second, il va s'agir d'implémenter un kernel SpMM dédié.

Commençons par le cas $\text{TopK}$. Le gradient $G_z = \frac{\partial \bcL}{\partial z_{BH}}$ est une matrice sparse, représenté de la même manière que $a_{BH}$ : pour chaque ligne $0 \leq b < B$ (soit chaque token du batch), on a exactement $K$ valeurs non-nulles, situées aux colonnes d'indices $i_{BK}[b, :]$. On veut calculer le gradient $G_W = \frac{\partial \bcL}{\partial W_{HD}} = G_z^\top x_{BD}$. On peut donc initialiser $G_W$ à zéro, puis ajouter la contribution de chaque token $b$ du batch. On multiplie la ligne la ligne $b$ de $x_{BD}$ par les $K$ valeurs de la colonne $b$ de $G^\top$, situées aux lignes $i_{BK}[b, :]$. On ajoute ces $K$ lignes à $G_W$ aux positions $i_{BK}[b, :]$.

On démarre $B$ programmes, chacun s'occupant d'un $b$. On peut même subdiviser manuellement le travail selon $D$, en démarrant une grille $B \times \bbB\bbS_D$ où chaque programme s'occupe d'un $b$ et de $k\bbB\bbS_D \leq d < (k+1)\bbB\bbS_D$. Un des avantages de cette approche est que les accès mémoire sont contigus : on lit $x_{BD}$ ligne par ligne et $(G_z)^\top$ colonne par colonne. Le défaut est cependant la présence d'additions atomiques à $G_W$ : une même caractéristique $h$ peut être active dans plusieurs tokens donc plusieurs programmes peuvent ajouter aux mêmes positions. On a ici un kernel "scatter".