TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Résumé technique : TAPIOCA – Pourquoi l'élagage conscient de la tâche améliore les capacités des modèles hors distribution

Énoncé du problème
Des travaux récents ont démontré que la suppression de couches dans les grands modèles de langage (LLM) lors de l'inférence — spécifiquement par le biais de méthodes d'élagage conscient de la tâche comme TALE (Task-Aware Layer Elimination) — peut améliorer les performances sur des tâches spécifiques sans réentraînement. Cependant, les mécanismes à l'origine de ces gains restent obscurs. Une question centrale est de savoir si ces améliorations découlent d'avantages généraux de compression, d'effets de régularisation, ou d'interactions spécifiques entre les représentations internes du modèle et la distribution des entrées. De plus, il est inconnu si ces gains sont universels ou spécifiques à certains types de décalages de distribution.

Méthodologie
Les auteurs investiguent ce phénomène en combinant des expériences synthétiques contrôlées et des évaluations à grande échelle sur des LLM affinés.

1. Cadre de régression contrôlé : L'étude commence avec un transformateur de 12 couches entraîné à partir de zéro sur une régression linéaire en contexte. La distribution d'entraînement est explicitement définie (coefficients échantillonnés selon $U(-1, 1)$). Les auteurs évaluent le modèle sur des données in-distribution (ID) (même distribution) et out-of-distribution (OOD) (coefficients échantillonnés selon des intervalles plus larges $U(-\sigma, \sigma)$ où $\sigma > 1$). Ils appliquent TALE, un algorithme d'élagage glouton qui supprime des couches pour maximiser la précision de validation sur une distribution spécifique.
2. Évaluation à grande échelle de LLM : Les résultats sont étendus à des LLM affinés (Llama 3.1 8B et GPT-OSS 120B) spécialisés dans le raisonnement mathématique et la génération de code. Ici, la distribution « adaptée » est la tâche d'affinage (par exemple, MATH500), et les tâches OOD incluent des benchmarks généraux comme MMLU et BoolQ.
3. Analyse géométrique et causale : Pour expliquer pourquoi l'élagage fonctionne, les auteurs analysent la géométrie des états cachés. Ils mesurent les normes représentatives par couche et les distances paires entre tokens. Ils emploient également des substituts linéaires locaux pour approximer les transformations de couche et mènent des interventions causales, telles que l'insertion de cartes de substitut inverses et le redimensionnement des mises à jour résiduelles, pour isoler l'effet de la distorsion géométrique.

Résultats clés

• Gains spécifiques à l'OOD : La découverte empirique principale est que l'élagage conscient de la tâche ne produit des gains de performance que sous un décalage de distribution. Lorsqu'évalué sur la distribution à laquelle le modèle est adapté (ID), TALE échoue soit à améliorer les performances, soit les dégrade activement en supprimant des couches nécessaires à une précision optimale. En revanche, sur des données OOD, TALE identifie et supprime systématiquement des couches qui améliorent la précision.
• Inadéquation de la géométrie des représentations : Les auteurs démontrent que les entrées OOD induisent une géométrie de représentation distordue par rapport aux entrées ID. Plus précisément, les entrées OOD font dévier systématiquement les normes d'états cachés par couche et les distances paires des profils observés sur les données ID. Dans le cadre de régression, les entrées OOD provoquent une inflation des distances à travers les couches ; dans les LLM, les trajectoires OOD divergent des trajectoires ID dans les couches intermédiaires à tardives.
• Mécanisme d'amplification : L'étude identifie que certaines couches agissent comme des « amplificateurs dépendants de la distribution ». Sur les entrées ID, ces couches effectuent des raffinements utiles. Cependant, sur les entrées OOD, ces mêmes couches exhibent une amplification à fort gain et faible rang, repoussant les représentations plus loin de la géométrie adaptée à la tâche. TALE améliore les performances en supprimant ces couches spécifiques, empêchant ainsi l'amplification de la distorsion géométrique.
• Alignement, pas seulement contraction : L'élagage ne supprime pas simplement les normes d'activation. Dans des expériences où la direction du décalage de distribution est inversée (entraînement sur des distributions larges, test sur des distributions étroites), l'élagage étend les distances de représentation pour les aligner sur le profil ID. Le mécanisme central est le réalignement des représentations OOD vers la géométrie apprise lors de l'adaptation, indépendamment du fait que cela nécessite une contraction ou une expansion.
• Preuve causale : Les interventions causales confirment que la distorsion géométrique est le moteur de la perte de performance. L'insertion d'une carte de substitut inverse après une couche élaguée (pour annuler l'effet géométrique de la couche) restaure la performance OOD presque au niveau du modèle élagué. De même, le redimensionnement continu de la mise à jour résiduelle d'une couche sélectionnée améliore la précision à mesure que la contribution de la couche est atténuée, écartant les artefacts topologiques ou la simple redondance comme cause unique.

Portée et affirmations
L'article prétend fournir une explication géométrique novatrice du succès de l'élagage lors de l'inférence. Il soutient que l'importance d'une couche n'est pas fixe mais conditionnelle à la fois à la tâche et à la distribution des entrées. Les auteurs postulent que les modèles pré-entraînés encodent plusieurs fonctions spécialisées par tâche, chacune associée à une géométrie de représentation locale. Les entrées ID s'ajustent à cette géométrie, tandis que les entrées OOD la violent.

La portée de ce travail réside dans :

1. Clarification des limites de l'élagage : Il établit que l'élagage conscient de la tâche n'est pas un booster universel de performance, mais un mécanisme correctif pour les décalages de distribution.
2. Interprétation géométrique : Il modifie la compréhension de la redondance des couches, passant des comptes de paramètres statiques aux distorsions géométriques dynamiques et dépendantes de la distribution.
3. Insight mécanistique : Il fournit des preuves causales que les gains de performance découlent de la correction de la trajectoire des représentations internes pour correspondre à la géométrie adaptée du modèle, plutôt que d'une augmentation de la capacité du modèle ou d'une réduction du bruit.

Les auteurs concluent que l'élagage améliore les performances OOD en atténuant les couches qui amplifient le décalage entre la distribution des entrées et la géométrie de tâche apprise du modèle, réalignant efficacement l'état interne du modèle avec les conditions dans lesquelles il a été optimisé.