Pular uma Camada ou Criar um Loop? Aprendendo Programa de Camadas em LLMs

Resumo

Modelos de linguagem de grande escala (LLMs) realizam inferência seguindo uma profundidade e ordem fixas, executando todas as camadas de forma não recorrente. Revelamos a ampla existência de programas-de-camadas (PoLar) flexíveis, dinâmicos e que não exigem treinamento, nos quais camadas pré-treinadas podem ser agrupadas como módulos e depois puladas ou repetidas para formar um programa personalizado para cada entrada. Para a maioria das entradas, execuções de programas substancialmente mais curtas podem alcançar a mesma ou melhor acurácia, enquanto predições incorretas do LLM original podem ser corrigidas por programas alternativos com menos camadas. Essas observações indicam que a inferência admite múltiplos cálculos latentes válidos além do passe direto padrão. Para alcançar PoLar de forma eficiente na prática, propomos uma rede de predição PoLar leve, que aprende a gerar programas de execução que dinamicamente pulam ou repetem camadas pré-treinadas para cada entrada. Experimentos em benchmarks de raciocínio matemático demonstram que o PoLar melhora consistentemente a acurácia em relação à inferência padrão e a métodos dinâmicos de profundidade anteriores, frequentemente executando menos camadas, e que esses ganhos persistem sob avaliação fora da distribuição. Nossos resultados sugerem que a execução de profundidade fixa captura apenas um subconjunto estreito da capacidade latente de raciocínio de um 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.