Learning Linear Regression with Low-Rank Tasks in-Context

Titre : Learning Linear Regression with Low-Rank Tasks In-Context

Auteurs : Kaito Takanami, Takashi Takahashi, Yoshiyuki Kabashima (Université de Tokyo, RIKEN AIP)
Conférence : AISTATS 2026

1. Problématique et Contexte

L'apprentissage en contexte (In-Context Learning ou ICL) est une capacité fondamentale des grands modèles de langage (LLM) modernes, leur permettant d'effectuer de nouvelles tâches à partir de quelques exemples fournis dans le contexte, sans mise à jour des paramètres. Cependant, les mécanismes théoriques sous-jacents restent mal compris, en particulier dans des scénarios réalistes où les tâches partagent des structures communes (familles de tâches structurées).

La plupart des études théoriques précédentes se sont concentrées sur des tâches indépendantes ou des modèles "jouets" sans structure sous-jacente partagée. Ce travail vise à combler ce vide en analysant comment un modèle d'attention linéaire apprend à effectuer une régression linéaire lorsque les tâches de pré-entraînement partagent une structure de faible rang (low-rank structure). L'objectif est de caractériser précisément la distribution des prédictions et l'erreur de généralisation dans la limite des hautes dimensions.

2. Méthodologie

Les auteurs adoptent une approche théorique rigoureuse basée sur la physique statistique et la théorie des matrices aléatoires.

• Modèle : Ils utilisent un transformateur basé sur l'attention, spécifiquement un modèle d'attention linéaire à une seule couche. Ce modèle est choisi pour sa tractabilité analytique tout en capturant des phénomènes non triviaux de l'ICL.
• Configuration des Données :
  • Pré-entraînement : Le modèle est entraîné sur un ensemble de tâches $W$ générées à partir d'une structure latente de faible rang. Les vecteurs de tâche $w_\mu$ sont projetés dans un espace de dimension $D$ via une matrice de caractéristiques partagée $A$ de rang $r$ ($r \le D$).
  • Inférence : Le modèle est évalué sur de nouvelles instances (tâches vues, tâches de la même distribution, ou tâches hors distribution) avec un nombre limité d'exemples (few-shot).
• Régime Asymptotique : L'analyse est menée dans la limite des hautes dimensions où la dimension des entrées $D$, le nombre d'échantillons $M$, et le nombre de tâches $M_0$ tendent vers l'infini, tout en maintenant des ratios fixes (ex: $\alpha = L/D$, $\rho = r/D$, $\kappa = M_0/D$).
• Outils Mathématiques :
  • Méthode des Répliques (Replica Method) : Utilisée pour calculer les moments statistiques de la solution optimale du problème d'optimisation (minimisation de l'erreur quadratique moyenne).
  • Théorie des Matrices Aléatoires : Pour analyser le spectre de la matrice de covariance des tâches et comprendre les transitions de phase.
  • Décomposition de l'Erreur : La prédiction est décomposée en un signal algorithmique et des termes de bruit.

3. Contributions Clés et Résultats

L'article apporte plusieurs résultats théoriques majeurs :

A. Décomposition de la Prédiction en Contexte

Les auteurs montrent que la prédiction du modèle se décompose en trois composantes distinctes :

1. Signal Algorithmique ($\hat{y}_{algo}$) : Correspond à une procédure de régression linéaire apprise implicitement. Le modèle apprend à projeter l'estimateur naïf (filtre adapté) sur le sous-espace de faible rang appris durant le pré-entraînement.
2. Bruit de Mémorisation ($\hat{y}_{mem}$) : Un terme de bruit corrélé aux tâches spécifiques mémorisées durant le pré-entraînement. Ce bruit est supprimé lorsque la tâche actuelle correspond à une tâche vue, mais peut augmenter avec la longueur du contexte (sur-apprentissage aux motifs spécifiques).
3. Bruit Structurel ($\hat{y}_{struct}$) : Un bruit lié à la partie de l'estimation in-contexte qui est orthogonale à la structure apprise. Ce bruit domine lorsque la tâche test ne respecte pas la structure de faible rang (hors distribution).

Cette décomposition révèle que l'ICL agit comme un mécanisme de réduction de bruit dépendant du contexte.

B. Régularisation Implicite par les Fluctuations Statistiques

Un résultat contre-intuitif est mis en évidence : l'apprentissage de tâches de faible rang est stabilisé par les fluctuations statistiques des données de pré-entraînement finies.

• Dans la limite idéale de données infinies (sans bruit d'échantillonnage), le problème d'optimisation devient mal posé (ill-posed) pour les tâches de faible rang, conduisant à une instabilité et une mauvaise généralisation.
• Les fluctuations finies introduisent une régularisation implicite (un terme de régularisation effectif $\hat{\bar{q}} > 0$) qui brise les directions plates du paysage d'optimisation, rendant l'apprentissage stable et robuste. Cela suggère que le "bruit" des données est une fonctionnalité cruciale, et non un bug.

C. Transition de Phase et Compromis Spécialisation/Robustesse

L'analyse spectrale de la matrice de covariance des tâches révèle une transition de phase aiguë gouvernée par la relation entre la complexité des tâches ($\rho = r/D$) et la diversité des tâches ($\kappa = M_0/D$).

• Régime de Tâches Rares ($\rho > \kappa$) : La diversité des tâches est le goulot d'étranglement. La performance est plafonnée et ne s'améliore pas même si les tâches deviennent plus simples.
• Régime de Tâches Riche ($\rho < \kappa$) : Le modèle peut pleinement exploiter la structure de faible rang.
  • Compromis fondamental : Dans ce régime, une spécialisation parfaite aux tâches de la distribution (faible erreur TM et IDG) s'accompagne d'une dégradation drastique de la robustesse aux tâches hors distribution (ODG).
  • La transition se produit à $\rho = \kappa$, marquant un changement qualitatif dans la capacité du modèle à généraliser.

4. Validation Expérimentale

• Simulations Numériques : Les prédictions théoriques (erreurs de généralisation, paramètres d'ordre) sont validées par des expériences numériques sur des modèles linéaires, montrant une excellente concordance.
• Attention Non-Linéaire : Les auteurs étendent l'analyse à un modèle d'attention softmax non linéaire. Bien que la transition de phase ne soit pas aussi aiguë (en raison de la complexité du modèle), les tendances qualitatives (amélioration avec la simplicité des tâches, plateau lié à la diversité) sont reproduites, confirmant la pertinence des résultats théoriques au-delà du cadre linéaire.

5. Signification et Implications

Ce travail fournit un cadre théorique unifié pour comprendre comment les transformateurs "apprennent à apprendre" (meta-learning) :

1. Mécanisme d'ICL : L'ICL n'est pas seulement une mémorisation, mais un processus dynamique de réduction de bruit où le contexte permet de sélectionner la bonne structure algorithmique.
2. Rôle des Données : L'étude remet en question l'idée que des données parfaites sont toujours préférables, montrant que les imperfections statistiques des ensembles de données finis sont essentielles pour la stabilité de l'apprentissage de structures complexes.
3. Stratégie de Pré-entraînement : Les résultats suggèrent qu'il existe un "point idéal" (sweet spot) dans la conception des ensembles de données de pré-entraînement. Pour obtenir un modèle robuste, il ne faut ni trop de diversité (qui empêche l'apprentissage de structures complexes) ni trop de simplicité (qui nuit à la robustesse hors distribution). Une diversité modérée proche de la transition de phase ($\rho \approx \kappa$) semble optimale pour équilibrer spécialisation et généralisation.

En résumé, cet article offre une compréhension profonde des mécanismes mathématiques derrière l'efficacité de l'ICL sur des tâches structurées, reliant la physique statistique des systèmes désordonnés aux performances des modèles de langage modernes.