Transformer RL: 단일 Layer만으로 Full-Parameter 성능 달성 가능성 분석

서론

최근 대규모 언어 모델(LLM)의 개발 트렌드는 단순히 규모를 키우는 것을 넘어, 강화학습(Reinforcement Learning, RL)을 활용한 후기 학습(Post-training)에 집중되고 있습니다. 이는 LLM이 단순한 다음 토큰 예측 기계에서 벗어나, 인간의 선호도나 복잡한 목표(예: 수학적 추론, 코드 생성, 에이전트 결정)를 충족하는 지능적인 시스템으로 진화할 수 있게 하는 핵심 동력입니다.

하지만 기존의 RL 기반 LLM 튜닝 접근법은 근본적인 가정을 깔고 있습니다. 바로 “모든 트랜스포머 레이어가 전체 성능 향상에 균일하게 기여한다"는 가정입니다. 즉, 모델 전체 파라미터를 한꺼번에 업데이트하면 모든 레이어가 비슷한 수준의 이득을 얻는다고 간주하는 것입니다. 그러나 실제 LLM 내부에서 RL 적응(adaptation)이 어떻게 분산되고 집중되는지에 대한 체계적인 이해는 여전히 부족했습니다.

본 연구는 이러한 가정을 정면으로 반박하며, 트랜스포머 레이어별 기여도(Layer Contribution)를 측정함으로써 RL 이득의 분포 패턴을 심층적으로 탐구합니다. 이 분석은 LLM 튜닝의 패러다임을 ‘전체 업데이트’에서 ‘정밀 타겟팅’으로 전환할 수 있는 중요한 단서를 제공하며, 모델 효율성 극대화와 계산 자원 최적화라는 실질적인 목표를 달성하는 데 결정적인 역할을 합니다.

본론: RL 이득의 분산 메커니즘 분석

1. Layer Contribution의 정의 및 원리

기존 연구들은 전체 파라미터 $\Theta$에 대한 강화학습 손실 함수 $L_{RL}(\Theta)$를 최소화하는 데 집중했습니다. 하지만 우리는 특정 레이어 $l$이 전체 성능 향상에 얼마나 기여했는지 정량적으로 측정하고자 합니다. 여기서 레이어 기여도는 단일 레이어 $l$만을 고립시켜 학습했을 때, 모델 전체 파라미터($\Theta$)를 모두 RL로 튜닝했을 때 얻을 수 있었던 총 성능 향상(Improvement) 중 얼마나 많은 부분을 회복했는지의 비율로 정의됩니다.

$$ \text{Layer Contribution}_l = \frac{\text{Performance}(\text{Model with } l \text{ trained in isolation}) - \text{Performance}(\text{Pre-trained Model})}{\text{Performance}(\text{Full RL Trained Model}) - \text{Performance}(\text{Pre-trained Model})} $$

이 지표를 통해 우리는 어떤 레이어가 ‘핵심적인 이득 창출자(Gain Driver)‘인지 식별할 수 있습니다.

2. 트랜스포머 내부의 정보 흐름과 RL 이득 집중 (Mermaid Diagram)

트랜스포머는 입력 임베딩에서 시작하여 여러 개의 인코더/디코더 블록을 거쳐 최종 출력까지 정보를 순차적으로 처리합니다. 각 레이어는 이전 레이어에서 추출된 추상화 수준의 특징(Abstract Feature)을 다음 단계로 전달하며, RL 이득 역시 이 정보 흐름에 따라 분산됩니다.

다음 다이어그램은 일반적인 트랜스포머 스택에서의 정보 흐름과 해당 레이어에 집중되는 RL 적응 이득($\Delta \text{RL}$)의 경향성을 보여줍니다.

  graph LR
    A[Input Embedding] --> B(Layer 1: Boundary)
    B --> C(Middle Layer k-2: Core Feature Extraction)
    C --> D(Middle Layer k-1: High Abstraction & Alignment)
    D --> E(Output Head / Final Token Prediction)

    subgraph RL Gain Distribution
        direction LR
        G[Low $\Delta$RL] --> B
        H[High $\Delta$RL] --> C
        I[Highest $\Delta$RL] --> D
        J[Moderate $\Delta$RL] --> E
    end

    B --> G
    C --> H
    D --> I
    E --> J

3. 핵심 발견: 단일 레이어의 충분성 (Table & Analysis)

본 연구는 이 메커니즘이 매우 강력하고 일관된 패턴을 따른다는 것을 밝혀냈습니다. 놀랍게도, **단 하나의 트랜스포머 레이어만 학습시켜도 전체 파라미터 RL 튜닝으로 얻은 성능 향상의 대부분(most of the gains)**을 회복하거나 심지어 초과 달성하는 경우가 관찰되었습니다.

이러한 현상은 특히 중간 레이어(Middle Layers)에 집중되어 나타납니다. 입력 경계층이나 출력 경계층은 모델의 초기 특징 추출 또는 최종 예측 조정 역할에 머무르는 반면, 중간 레이어는 복잡하고 고수준적인 추상적 표현을 생성하며, RL이 목표하는 ‘복잡한 행동’과 가장 밀접하게 연결되어 있기 때문입니다.

다음 표는 다양한 태스크와 모델 가족(Qwen2.5, Qwen3)에서 관찰된 대표적인 레이어 기여도 패턴을 비교 정리한 것입니다. (기여도는 상대적 비율로 표시됨).

Layer Type	역할	평균 Layer Contribution (%)	주요 특징 및 관찰 내용
Input Boundary	초기 임베딩/단순 패턴 인식	8% ~ 12%	기본적인 토큰-특징 매핑에 기여. RL 이득은 낮음.
Middle Layers (k-3 to k-1)	복잡한 추상화, 목표 연관성 파악	55% ~ 70%	가장 높은 집중도. 수학적/코드 생성 등 고차원 태스크의 핵심 이득 창출.
Output Boundary	최종 예측 조정, 정렬(Alignment)	15% ~ 25%	RL 목표에 맞춘 미세한 출력 형태 조정에 기여.

4. 실무 적용을 위한 Layer-wise RL Fine-tuning 가이드 (Step-by-step & Code)

이러한 분석 결과는 LLM 개발자에게 ‘사후적 최적화’가 아닌 ‘선제적 설계’의 관점을 제공합니다. 이제 우리는 모든 레이어를 낭비적으로 업데이트할 필요 없이, 가장 기여도가 높은 중간 레이어에 자원을 집중 투입할 수 있습니다.

Step-by-step 가이드:

Baseline 측정: Pre-trained LLM을 준비하고, RL 목표 태스크(예: 코드 생성)에서 성능($P_{Full}$)을 측정합니다.
Layer Isolation & Training: 모델의 모든 레이어 $l$에 대해 독립적으로 학습을 수행합니다. (다른 레이어는 고정).
Contribution Calculation: 각 단일 레이어 학습 결과로 얻은 성능($P_l$)을 사용하여 Layer Contribution를 계산하고 순위를 매깁니다.
Targeted Fine-tuning: 기여도가 가장 높은 상위 $N$개 레이어를 선택하여, 해당 레이어들의 파라미터만 활성화(Unfreeze)시키고 RL 튜닝을 진행합니다.

개념 설명용 코드 예시 (PyTorch 기반): 다음은 Layer Contribution를 측정하기 위해 단일 레이어의 학습 이득을 계산하는 개념적인 PyTorch 함수입니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import torch
from transformers import AutoModelForCausalLM

def calculate_layer_contribution(model, layer_index, rl_optimizer, task_data):
    """
    주어진 모델에서 특정 레이어만 고정하고 RL 학습을 수행하여 기여도를 측정합니다.
    """
    # 1. 모든 파라미터 동결 (Freeze)
    for name, param in model.named_parameters():
        param.requires_grad = True # 기본적으로 모두 활성화 가정

    # 2. 특정 레이어만 해제 (Unfreeze) - 예: Transformer Block의 첫 번째 QKV 선형 레이어
    # 실제 구현에서는 layer_index에 따라 정확한 모듈을 지정해야 함
    for name, param in model.named_parameters():
        if f"model.h.{layer_index}.attn.q_proj" in name:
            param.requires_grad = True # 이 레이어는 활성화 유지
        else:
            # 다른 모든 레이어는 동결 (Freeze)
            param.requires_grad = False 

    # 3. RL 학습 루프 실행 (간단화)
    rl_optimizer.zero_grad()
    # ... (Forward pass, Reward calculation, Loss computation) ...
    loss = calculate_rl_loss(model, task_data)
    loss.backward()
    rl_optimizer.step()

    # 4. 성능 측정 및 반환 (이 예시에서는 손실 감소를 Proxy로 사용)
    return -loss.item() # Negative Loss = Performance Gain

결론

본 연구는 LLM 후기 학습에서 강화학습 이득의 분산 패턴에 대한 명확한 통찰을 제공했습니다. 핵심은 RL 적응이 모델 전체에 균등하게 퍼져 있지 않으며, 대부분 중간 트랜스포머 레이어에 고도로 집중되어 있다는 사실입니다. 단일 중간 레이어만으로도 Full-Parameter RL 튜닝의 성능을 회복하거나 능가할 수 있다는 발견은 LLM 개발 패러다임에 혁명적인 변화를 가져올 잠재력을 지니고 있습니다.

이러한 Layer Contribution 분석을 통해, 우리는 계산 비용과 시간 복잡도를 크게 줄이면서도 최적화된 결과를 얻는 ‘Surgical Fine-tuning’ 시대를 열 수 있게 되었습니다. 향후 연구에서는 특정 태스크(예: Agentic Decision Making)나 모델 아키텍처에 따라 기여도가 달라지는 미묘한 패턴을 더욱 정밀하게 분석하고, 동적으로 레이어 가중치를 조절하는 메커니즘을 개발할 것으로 기대됩니다.

이러한 Layer-wise 접근 방식은 LLM의 내부 작동 원리를 이해하는 것을 넘어, 모델 자체를 더 효율적이고 지능적으로 설계하는 데 필수적인 도구가 될 것입니다.

— 📚 참고 자료

arXiv:2607.01232v1 - Is One Layer Enough? Training A Single Transformer Layer Can Match Full-Parameter RL Training
http://arxiv.org/abs/2607.01232v1

출처: http://arxiv.org/abs/2607.01232v1