L'attention latente : compresser le cache

Treizième article de la série Smallem. Précédent : Trois pièces modernes, et comment mesurer une idée.

L'article précédent s'est terminé sur un succès : RoPE améliore nettement le modèle. Celui-ci raconte un résultat négatif. L'attention latente de DeepSeek fait exactement ce qu'elle promet pour la mémoire, et elle dégrade notre modèle. Comprendre pourquoi a demandé trois entraînements de plus, et c'est la partie la plus instructive.

Le problème : le cache grossit avec le texte

L'article 11 a introduit le cache de génération. Pour ne pas recalculer tout le passé à chaque nouveau token, chaque couche garde la clé et la valeur de chaque position, pour chaque tête.

Notre modèle a 8 têtes de 64 nombres. Une position coûte donc, par couche, 512 nombres de clés et 512 de valeurs : 2 × C = 1 024 nombres. Avec 8 couches et un contexte plein de 512 tokens, le cache pèse 16 Mio (mébioctets, 2²⁰ octets).

C'est négligeable ici. Mais ce coût est proportionnel à tout : à la dimension du modèle, au nombre de couches, à la longueur du texte, et au nombre de conversations servies en même temps. Pour un modèle de dimension 7 168 à 61 couches qui lit 100 000 tokens, le même calcul donne plus d'une centaine de gigaoctets par conversation. À cette échelle, c'est la mémoire du cache, et non le calcul, qui limite le nombre d'utilisateurs servis par carte.

L'idée : garder un résumé

L'attention latente (MLA, Multi-head Latent Attention, introduite par DeepSeek-V2) part d'une observation. Les clés et les valeurs de toutes les têtes sont calculées à partir du même vecteur, celui de la position. Ces 1 024 nombres sont donc peut-être redondants. Au lieu de les garder tous, on pourrait garder un vecteur beaucoup plus petit, et recalculer à partir de lui les clés et les valeurs quand on en a besoin.

Chaque position est donc d'abord compressée en un vecteur latent, de taille d_c bien plus petite que C. Les clés et les valeurs en sont ensuite déduites par deux matrices :

c = rms(x · W_latent)        [d_c]      le latent, seul gardé dans le cache
k = c · W_clé                [H × D]    les clés de toutes les têtes
v = c · W_valeur             [H × D]    les valeurs de toutes les têtes

À la génération, on garde c. Quand un nouveau token arrive, les clés et les valeurs de tout le passé sont recalculées à partir des latents. On échange de la mémoire contre du calcul.

Les requêtes passent par leur propre latent, plus large. Elles ne sont jamais gardées : là, la compression ne sert qu'à économiser des paramètres.

Avec d_c = 128, le cache passe de 1 024 à 128 nombres par position.

L'obstacle : RoPE

RoPE fait tourner les clés selon leur position. Or le cache ne contient plus de clés, mais des latents, et la clé est c · W_clé.

Si l'on tourne la clé après l'avoir calculée, tout va bien pour le calcul, mais il faut alors garder la clé tournée, et le latent ne sert plus à rien. Peut-on tourner le latent à la place ? Non : une rotation et une multiplication par une matrice quelconque ne commutent pas. Tourner c puis multiplier par W_clé ne donne pas la clé tournée.

DeepSeek sort de l'impasse en donnant à la position son propre canal. Chaque clé et chaque requête reçoit une seconde partie, petite, qui porte seule la rotation :

k_r = rotation(x · W_clé_tournée)          [d_r]      une seule pour toutes les têtes
q_r = rotation(c_q · W_requête_tournée)    [H × d_r]  une par tête

La clé d'une tête est la concaténation [k ; k_r], sa requête [q ; q_r]. Le produit scalaire de deux vecteurs mis bout à bout est la somme des produits de leurs parties :

score = [q ; q_r] · [k ; k_r] / √(D + d_r)
      = (q · k  +  q_r · k_r) / √(D + d_r)

Le premier terme compare des contenus, le second des positions. C'est le RoPE découplé. Le cache garde c et k_r : avec d_c = 128 et d_r = 32, cela fait 160 nombres par position et par couche, contre 1 024.

Le code

La couche a dix matrices au lieu de quatre :

pub struct LatentAttention<B: Backend> {
    pub query_down: Projection<B>,   // C → d_q
    pub query_norm: RmsNorm<B>,
    pub query_up: Projection<B>,     // d_q → C
    pub query_rope: Projection<B>,   // d_q → H × d_r
    pub latent_down: Projection<B>,  // C → d_c
    pub latent_norm: RmsNorm<B>,
    pub key_up: Projection<B>,       // d_c → C
    pub value_up: Projection<B>,     // d_c → C
    pub key_rope: Projection<B>,     // C → d_r
    pub proj: Linear<B>,
}

Une Projection est une couche linéaire sans biais. Les RMSNorm sur les latents ont un rôle précis : un latent est le produit de l'entrée par une matrice tirée avec un petit écart-type, et la clé le produit de ce latent par une seconde matrice du même genre. Sans normalisation entre les deux, les clés partiraient plusieurs fois trop petites.

La passe avant est celle de l'attention classique, à deux détails près. Le premier est le cache : ce qui en sort n'est pas utilisable tel quel.

if let Some(cache) = cache {
    let all = cache.latents.extend(tape, new);
    k_rope = cache.rotated_keys.extend(tape, k_rope);
    total = tape.shape(all)[2];
    latent = tape.reshape(all, &[1, total, latent_dim]);
}
// Les clés et les valeurs de toutes les têtes, pour les `T` positions.
let k = split(self.key_up.forward(tape, latent), total, head_dim);
let v = split(self.value_up.forward(tape, latent), total, head_dim);

Comme à l'article 11, il n'y a qu'une seule passe avant. L'entraînement est le cas sans cache. La génération, le cas où une seule position est nouvelle.

Le second détail est le calcul des scores, qui a demandé une mesure.

Deux produits ou une concaténation ?

La formule du score offre deux écritures, mathématiquement égales.

• Deux produits : calculer q · kᵀ, puis q_r · k_rᵀ, et additionner les deux tableaux de scores. Aucune opération nouvelle : le produit matriciel sait déjà étirer la clé tournée, commune à toutes les têtes.
• Une concaténation : mettre q_r à la suite de q, k_r à la suite de k, et faire un seul produit. Cela demande une opération que le projet n'avait pas.

La première était la plus simple, et elle a été écrite d'abord. Mais un petit banc d'essai a montré qu'elle coûtait 1,4 ms de plus par couche qu'un produit unique sur des vecteurs de 96 nombres. La raison : les tableaux de scores [B, H, T, T] sont les plus gros tenseurs de la couche, 33 millions de nombres, et les additionner coûte cher. Les vecteurs à concaténer sont cinq fois plus petits.

L'opération Concat a donc été écrite, avec ses deux kernels :

pub fn concat_last(a: &[f32], b: &[f32], width_a: u32, width_b: u32,
                   mut out: DisjointSlice<f32>) {
    let idx = thread::index_1d();
    let flat = idx.get();
    if let Some(out_elem) = out.get_mut(idx) {
        let (width_a, width_b) = (width_a as usize, width_b as usize);
        let row = flat / (width_a + width_b);
        let column = flat % (width_a + width_b);
        *out_elem = if column < width_a {
            a[row * width_a + column]
        } else {
            b[row * width_b + column - width_a]
        };
    }
}

Chaque thread écrit un élément de la sortie et va le chercher dans la bonne entrée. La rétropropagation fait le chemin inverse : elle découpe le gradient en deux. Sa définition de référence, écrite avec les primitives, amène le dernier axe en premier, utilise la concaténation de lignes qui existait déjà, puis le remet en place.

Les deux écritures ont ensuite été mesurées sur un vrai pas d'entraînement : 189,3 ms avec deux produits, 180,1 ms avec la concaténation. Elle est restée.

Vérifier une couche compliquée

Dix matrices, deux normalisations, une rotation, une concaténation : les occasions de se tromper ne manquent pas, et une erreur ici ne ferait pas planter le programme. Elle donnerait un modèle un peu moins bon, ce qui est précisément le résultat qu'on va observer. Il faut donc pouvoir écarter le bug.

Le test principal réécrit la couche à la main, d'après les formules du rapport, sans aucun tenseur : des boucles sur des vecteurs, les clés réellement concaténées, le softmax écrit terme à terme.

for head in 0..HEADS {
    let mut query = q[head * D..(head + 1) * D].to_vec();
    query.extend(rotate(&q_rope[head * rope_dim..(head + 1) * rope_dim], position));
    let mut key = k[head * D..(head + 1) * D].to_vec();
    key.extend(&k_rope);
    // …
}

Les deux calculs s'accordent à 1e-11 près. S'y ajoutent les contrôles de toujours : le test du futur, la séquence traitée d'un coup ou morceau par morceau à travers le cache, les scores inchangés quand toutes les positions sont décalées, le gradient de tous les paramètres par différences finies, le sur-apprentissage d'un lot, et l'accord du CPU et du GPU.

Les résultats

Même protocole qu'à l'article précédent : 49 millions de tokens, une graine et la référence modern.

Le cache est bien 6,4 fois plus petit. Mesuré sur le modèle entraîné : 2,5 Mio au lieu de 16,0 pour un contexte plein.

Le modèle est nettement moins bon. L'écart vaut douze fois le seuil de bruit. Il rend plus de la moitié du gain de RoPE.

Tout est plus lent. L'entraînement de 5 % : dix projections et deux normalisations par couche au lieu de quatre projections. La génération de 19 %, 130 tokens par seconde contre 161 : à chaque token, les clés et les valeurs de toutes les positions sont recalculées.

Chercher la cause

Un résultat négatif ne vaut rien tant qu'on ne sait pas d'où il vient. Trois explications étaient possibles :

1. Le goulot. 128 nombres doivent porter à la fois les clés et les valeurs de 8 têtes, là où l'attention classique dispose de deux matrices 512 × 512 indépendantes.
2. RoPE affaibli. Dans la référence, chaque tête tourne ses 64 composantes : tout le contenu d'une clé est modulé par la position. Ici, seules 32 composantes tournent, et elles sont communes aux 8 têtes. Or RoPE était le seul gain réel de l'article précédent.
3. L'optimisation. Chaque projection est devenue un produit de deux matrices, entraîné avec les réglages de la référence.

Chacune se teste par un entraînement.

Le goulot explique les deux tiers de l'écart. Doubler le latent le fait tomber de 0,025 à 0,008. À 128 nombres pour 8 têtes, le latent était trop étroit pour ce modèle.

RoPE n'y est pour rien. Doubler la partie tournée ne change rien : 0,9666 contre 0,9670, un écart cinq fois plus petit que le bruit. Trente-deux composantes communes aux têtes suffisent à porter la position. C'était pourtant l'hypothèse la plus séduisante, et sans l'entraînement qui la teste, elle aurait fini dans cet article comme une explication.

Il reste un écart que la compression n'explique pas. Avec un latent de 512, aussi large que le modèle, plus rien n'est compressé du côté des clés et des valeurs. L'écart ne disparaît pas : il vaut encore 0,012, avec 14 % de paramètres de plus que la référence. Il est même un peu plus grand qu'avec 256, de deux fois le seuil de bruit, ce qui est trop peu pour en tirer une conclusion.

Ce qui reste ne peut venir que de la forme de la couche : des projections écrites comme un produit de deux matrices séparées par une normalisation, des requêtes qui passent elles aussi par un latent, et le tout entraîné avec le taux d'apprentissage d'une couche qui n'avait pas cette forme. Départager ces causes demanderait d'autres entraînements, à commencer par une recherche du taux d'apprentissage.

Pourquoi DeepSeek n'observe pas cette perte

Leur rapport annonce que MLA fait au moins aussi bien que l'attention classique, avec une compression bien plus forte que la nôtre : un latent de 512 pour une dimension de 7 168.

L'explication courante est qu'un grand modèle a des représentations redondantes, qu'on peut compresser sans rien perdre, alors qu'un modèle de 30 millions de paramètres se sert de toutes ses dimensions. Nos mesures sont compatibles avec cette idée mais elles ne la prouvent pas : il faudrait répéter l'expérience à plusieurs tailles.

Ce qui est établi est plus modeste. À notre échelle, MLA échange de la qualité et de la vitesse contre une mémoire dont nous ne manquons pas. Seize mébioctets de cache ne gênent personne. La technique répond à un problème que ce modèle n'a pas, et la référence des articles suivants reste l'attention classique.

Ce que DeepSeek fait de plus

Notre couche recalcule les clés et les valeurs de tout le passé à chaque token. DeepSeek évite ce travail par une réécriture. Le score d'une requête face à une clé est q · (c · W_clé), ce qui s'écrit aussi (q · W_cléᵀ) · c : on peut projeter la requête dans l'espace des latents une fois, puis la comparer directement aux latents du cache. De même, la somme pondérée des valeurs devient une somme pondérée des latents, projetée une seule fois par W_valeur.

L'attention se fait alors sur les latents eux-mêmes, et la génération retrouve la vitesse d'un cache classique. Cette réécriture n'a pas été faite ici : elle demande une seconde passe avant, propre à la génération, alors que tout le projet tient à n'en avoir qu'une, testée dans ses deux usages.

Ce qu'il faut retenir

• Le cache de génération coûte 2 × C nombres par position et par couche. À grande échelle, c'est lui qui limite.
• MLA garde un latent dont clés et valeurs sont recalculées. La mémoire baisse, le calcul monte.
• Une rotation ne commute pas avec une matrice : RoPE doit passer par une partie séparée de la clé, et le score devient une somme « contenu + position ».
• Une couche compliquée se vérifie contre une réécriture naïve de ses formules.
• Un résultat négatif se dissèque : un entraînement par hypothèse.
• Une technique conçue pour un problème qu'on n'a pas peut coûter sans rien rapporter.

La suite s'attaque à l'autre moitié du bloc, le feed-forward : Un mélange d'experts.