Is One Layer Enough? A Single Transformer Layer Matches Full-Parameter RL Train (arxiv.org)

Researchers have found that training a single transformer layer can achieve comparable, and sometimes superior, performance to full-parameter reinforcement learning (RL) training in large language models (LLMs). Traditional methods assume each layer contributes uniformly to the model's improvement post-training, but this study challenges that perspective by introducing the concept of "layer contribution." This metric quantifies how much improvement can be attributed to training individual layers in isolation, revealing that a small subset—mostly centered in the middle of the transformer architecture—drives most of the RL gains. This discovery is significant for the AI/ML community as it offers a more efficient approach to fine-tuning models, suggesting that resources can be focused on specific layers rather than updating all parameters uniformly. The findings were consistent across seven models and three different RL algorithms, indicating a robust pattern in the importance of particular transformer layers across various tasks, including mathematical reasoning and code generation. This understanding could streamline RL training processes, reduce computational costs, and enhance model performance by targeting crucial layers more effectively.