DiScoFormer réduit de 37 fois l’erreur du KDE

Ai2 présente DiScoFormer, un transformer conçu pour estimer simultanément la densité et le score d’une distribution. L’annonce a été publiée le 29 juin 2026 par Ai2. Le modèle reçoit un ensemble de points et produit ces deux estimations en un seul passage, sans nouvel entraînement pour chaque distribution. La densité indique les zones où les observations se concentrent. Le score, gradient du logarithme de la densité, indique la direction vers les régions plus probables. Ai2 relie cette tâche à plusieurs usages, notamment les modèles génératifs fondés sur la diffusion, l’échantillonnage bayésien et les simulations de particules utilisées pour modéliser des systèmes comme les plasmas. Lire l’annonce d’Ai2.

Deux méthodes s’opposent aujourd’hui sur cette estimation. L’estimation par noyaux, ou KDE, fonctionne sans entraînement et s’applique à toute distribution. Elle devient toutefois nettement moins précise lorsque la dimension augmente. Les modèles neuronaux entraînés par appariement de scores restent précis en grande dimension, mais chacun doit apprendre une distribution particulière. Ils doivent ensuite être réentraînés pour une autre distribution. DiScoFormer cherche à réunir ces avantages dans un modèle unique. Son architecture évalue la densité et le score à n’importe quel point, y compris hors des observations fournies. Elle utilise l’attention croisée pour relier l’échantillon complet à chaque requête.

DiScoFormer couple densité et score

Une architecture partagée forme le cœur de DiScoFormer. Le modèle utilise des blocs de transformers empilés, avec un socle commun et deux têtes de sortie. La première estime la densité. La seconde estime le score. Ai2 exploite leur relation mathématique, puisque le score correspond au gradient du logarithme de la densité. L’écart entre les deux sorties fournit ainsi une perte de cohérence sans étiquette supplémentaire. À l’inférence, le modèle conserve le contexte, puis effectue quelques étapes de gradient sur cette perte. Ai2 indique que cette procédure l’adapte immédiatement à une entrée hors distribution, sans connaître la densité ou le score de référence.

Le KDE reste présent dans la conception plutôt que d’être remplacé par une boîte noire. Ai2 montre analytiquement que les poids d’une tête d’attention ressemblent presque à un noyau gaussien appliqué aux données. Un bloc d’attention croisée peut donc déjà reproduire la densité et le score du KDE. DiScoFormer apprend ensuite plusieurs échelles et les adapte aux données. Pour l’entraînement, Ai2 génère un nouveau modèle de mélange gaussien à chaque lot. Ces modèles peuvent approximer pratiquement toute distribution lisse avec suffisamment de composantes. Leurs densités et leurs scores disposent aussi d’une forme exacte, utilisée comme cible d’apprentissage. Consulter le rapport technique.

Les gains augmentent en grande dimension

Les résultats atteignent leur niveau le plus marqué en 100 dimensions. Selon Ai2, DiScoFormer réduit alors l’erreur de score d’environ 6,5 fois face au meilleur KDE réglé manuellement. Son erreur de densité diminue de plus de 37 fois dans la même comparaison. Le modèle continue aussi de progresser lorsque le nombre d’échantillons augmente. Ai2 rapporte parallèlement que le KDE finit par manquer de mémoire dans cette situation. DiScoFormer reste précis sur des mélanges comportant davantage de modes que ceux vus pendant l’entraînement. Il conserve également cette précision sur des formes non gaussiennes, notamment les distributions de Laplace et de Student-t.