¿Saltar una capa o ponerla en bucle? Aprendizaje de Programa de Capas en LLMs

Resumen

Los modelos de lenguaje de gran escala (LLMs) realizan inferencia siguiendo una profundidad y orden fijos, con una ejecución no recurrente de todas las capas. Revelamos la amplia existencia de programas de capas (PoLar) dinámicos, flexibles y sin entrenamiento, donde las capas preentrenadas pueden empaquetarse como módulos y luego omitirse o repetirse en bucle para formar un programa personalizado para cada entrada. Para la mayoría de las entradas, ejecuciones de programas sustancialmente más cortas pueden lograr la misma o mejor precisión, mientras que las predicciones incorrectas del LLM original pueden corregirse mediante programas alternativos con menos capas. Estas observaciones indican que la inferencia admite múltiples cómputos latentes válidos más allá del paso directo estándar. Para lograr PoLar de manera eficiente en la práctica, proponemos una red de predicción PoLar ligera, que aprende a generar programas de ejecución que omiten o repiten dinámicamente capas preentrenadas para cada entrada. Experimentos en puntos de referencia de razonamiento matemático demuestran que PoLar mejora consistentemente la precisión en comparación con la inferencia estándar y los métodos previos de profundidad dinámica, a menudo ejecutando menos capas, y que estas ganancias persisten bajo evaluación fuera de la distribución. Nuestros resultados sugieren que la ejecución de profundidad fija captura solo un subconjunto estrecho de la capacidad de razonamiento latente de 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.