Sauter une couche ou la boucler ? Apprentissage de programme de couches dans les LLM

Résumé

Les grands modèles de langage (LLMs) réalisent l'inférence en suivant une exécution non récurrente de toutes les couches, selon une profondeur et un ordre fixes. Nous révélons l'existence répandue de programme de couches (PoLar) flexible, dynamique et sans entraînement, où les couches pré-entraînées peuvent être regroupées en modules puis sautées ou bouclées pour former un programme personnalisé pour chaque entrée. Pour la plupart des entrées, des exécutions de programmes sensiblement plus courtes peuvent atteindre la même précision ou une meilleure, tandis que les prédictions incorrectes du LLM original peuvent être corrigées par des programmes alternatifs avec moins de couches. Ces observations indiquent que l'inférence admet plusieurs calculs latents valides au-delà du passage avant standard. Pour atteindre efficacement PoLar en pratique, nous proposons un réseau de prédiction PoLar léger, qui apprend à générer des programmes d'exécution qui sautent ou répètent dynamiquement des couches pré-entraînées pour chaque entrée. Les expériences sur des références de raisonnement mathématique montrent que PoLar améliore constamment la précision par rapport à l'inférence standard et aux méthodes dynamiques antérieures, souvent tout en exécutant moins de couches, et que ces gains persistent lors d'une évaluation hors distribution. Nos résultats suggèrent que l'exécution à profondeur fixe ne capture qu'un sous-ensemble étroit de la capacité de raisonnement latent d'un LLM.

English

Large language models (LLMs) perform inference by following a fixed depth and order, non-recurrent execution of all layers. We reveal the wide existence of training-free, flexible, dynamic program-of-layers (PoLar), where pretrained layers can be packed as modules and then skipped or looped to form a customized program for each input. For most inputs, substantially shorter program executions can achieve the same or better accuracy, while incorrect predictions of the original LLM can be corrected by alternative programs with fewer layers. These observations indicate that inference admits multiple valid latent computations beyond the standard forward pass. To efficiently achieve PoLar in practice, we propose a lightweight PoLar prediction network, which learns to generate execution programs that dynamically skip or repeat pretrained layers for each input. Experiments on mathematical reasoning benchmarks demonstrate that PoLar consistently improves accuracy over standard inference and prior dynamic-depth methods, often while executing fewer layers, and that these gains persist under out-of-distribution evaluation. Our results suggest that fixed-depth execution captures only a narrow subset of an LLM's latent reasoning capacity.