기술 요약: LaneRoPE

문제 정의

최적의 N 개 샘플링(best-of-N sampling)과 같은 병렬 테스트 시간 확장 기술은 동일한 입력 프롬프트에 조건부로 여러 시퀀스를 생성함으로써 대규모 언어 모델 (LLM) 의 정확도를 향상시킵니다. 이러한 방법들은 배치 추론의 계산 효율성을 활용하지만, 전통적으로는 조정되지 않은 병렬성으로 작동합니다. 표준 접근법에서는 배치 내의 각 시퀀스가 독립적으로 생성되므로, 다른 시퀀스들의 중간 생성물, 계산, 또는 관찰 결과를 재사용하지 않습니다. 이러한 독립성은 분해 가능한 문제 구조의 활용을 금지하고, 중간 출력의 공유를 제한하며, 다양하고 협력적인 응답 생성을 제약합니다.

기존의 "협력적" 접근법들은 중간 생성물에 대한 공유 가시성을 가능하게 함으로써 이를 해결하려 시도합니다. 그러나 많은 이전 방법들은 여러 시퀀스를 단일 긴 시퀀스로 가상으로 정렬하기 위해 위치를 재매개화 (reparameterizing) 하는 데 의존합니다. 저자들은 이 접근법이 본질적으로 토큰 차원과 시퀀스 (또는 "레인") 차원을 따라 위치를 정확하게 반영해야 하는 다중 시퀀스 생성 문제라고 주장합니다.

방법론: LaneRoPE

이 논문은 기본 LLM 아키텍처에 최소한의 변경을 가하면서 생성 시 $N > 1$개의 시퀀스 간 조정과 협력을 가능하게 하는 LaneRoPE를 제안합니다. 이 접근법은 두 가지 핵심 기술적 혁신에 의존합니다:

1. 시퀀스 간 어텐션 마스킹 (Inter-Sequence Attention Masking):
   LaneRoPE 는 표준 인과적 어텐션 메커니즘을 수정하여 시퀀스 샘플링이 서로 의존적이도록 합니다. 다음 토큰 $x^{(n)}_{i+1}$을 자신의 과거 $x^{(n)}_{1:i}$에만 조건부로 두는 대신, 모델은 모든 시퀀스의 과거 $\{x^{(m)}_{1:i}\}_{m=1}^N$에 조건을 둡니다. 이는 인과적 교차 시퀀스 어텐션 마스크를 통해 달성되어, 모든 병렬 레인에서 중간 추론 흔적이 가시화됩니다.

2. 시퀀스 인지 위치 인코딩 (LaneRoPE):
   표준 회전 위치 인코딩 (RoPE) 은 다른 시퀀스들의 동일한 시간 단계에 있는 토큰들을 구별할 수 없습니다. 왜냐하면 이들은 동일한 위치 인덱스를 공유하기 때문입니다. LaneRoPE 는 시퀀스 인덱스 $m$(즉, "레인") 을 인코딩하는 두 번째 블록 대각 회전 행렬 $R_{\Omega, m}$을 도입하여 RoPE 를 확장합니다.

   • 메커니즘: 쿼리 및 키 벡터는 토큰 인덱스 $i$와 레인 인덱스 $m$ 모두에 의해 회전됩니다: $f_{\{q,k\}}(x_i, i, m) = R_{\Omega, m} R_{\Theta, i} (W x_i + b)$.
   • 속성: 이 공식은 (토큰, 시퀀스) 결합 그리드를 2 차원 푸리에 기저로 취급합니다. 결정적으로 레인 회전 행렬이 직교하므로, 동일한 시퀀스 내 토큰 간의 어텐션 가중치는 변하지 않아 하위 호환성이 보장됩니다. $N=1$인 경우, LaneRoPE 는 표준 RoPE 로 단순화됩니다.
   • 초기화 전략: 저자들은 기존 행동을 모방하기 위해 사전 훈련된 모델에서 LaneRoPE 를 초기화하는 전략들을 제공합니다:
     • GroupThink: 레인 주파수를 $\Omega = K\Theta$(여기서 $K$는 시퀀스 간격) 로 설정함으로써, LaneRoPE 는 GroupThink 의 "가상 정렬"을 복원합니다.
     • 독립 샘플링: 프로젝션 레이어에 푸리에 편향을 추가함으로써, 모델은 교차 레인 어텐션을 억제하도록 초기화되어 효과적으로 독립 샘플링을 시뮬레이션할 수 있습니다.
     • NTK 인지 보정: GroupThink 초기화에서 발생하는 부정적인 가상 상대 위치 문제를 완화하기 위해, 저자들은 저주파 레인 회전을 감쇠시키는 스케일링 함수를 도입하여 모델의 절대 위치 인코딩 능력을 보존합니다.

훈련 파이프라인

레인 간 상호작용을 유도하기 위해, 저자들은 두 가지 파인튜닝 전략을 제안합니다:

1. 지도 파인튜닝 (SFT): LLM 이 $N$명의 서로 다른 어시스턴트 역할을 하여 문제 해결에 협력하는 데이터셋이 합성됩니다. 어시스턴트들은 서로의 부분적 사고를 보고 중복 작업을 피하도록 프롬프트됩니다. 최종 데이터셋은 이러한 대화의 개별 흔적으로 구성되며, 모델로 하여금 시퀀스 간 의존성을 학습하도록 강제합니다.
2. 카네만 - 트버스키 최적화 (KTO): 모델은 기본 모델에 의해 독립적으로 생성된 추론 흔적들의 배치 (일부는 정확하고 일부는 부정확함) 로 훈련됩니다. KTO 손실은 정확한 완성의 가능성을 최대화하고 부정확한 것의 가능성을 최소화하여, 병렬 레인들이 서로를 교정하거나 보완할 수 있는 비자명한 결합 분포를 학습하도록 장려합니다.

결과

저자들은 DeepSeek-R1-Distill-Qwen 모델 (1.5B 및 7B 파라미터) 을 사용하여 수학적 추론 벤치마크 (MATH500, AMC 23, AIME 24/25) 에서 LaneRoPE 를 평가했습니다.

• 성능: LaneRoPE 는 순차적 베이스라인 및 독립 샘플링에 비해 신뢰할 수 있는 정확도 향상을 보여줍니다. 7B 모델은 특히 NTK 인지 초기화 및 학습 가능한 주파수를 사용한 KTO 훈련 시 일관된 이득을 보입니다.
• 비교: LaneRoPE 는 7B 모델에서 다른 상호 의존적 생성 방법인 "Bridge" 베이스라인을 일반적으로 능가합니다. 이는 LaneRoPE 의 더 높은 정보 공유 능력과 더 효율적인 아키텍처에 기인합니다.
• 소규모 모델: 1.5B 모델은 혼합된 결과를 보여주며, 훈련되지 않은 GroupThink 초기화로부터 덜 이득을 보고 레인 증가에 따른 확장성이 약합니다. 이는 아키텍처가 아닌 기본 모델 크기의 한계를 시사합니다.
• 효율성: LaneRoPE 는 무시할 수 있는 오버헤드를 도입합니다. 추론 시간 측정 결과, 표준 트랜스포머 대비 약 6% 증가에 불과하며, 이는 주로 어텐션 메커니즘 자체가 아닌 전처리/후처리 때문입니다. 반면, Bridge 베이스라인은 추가 레어로 인해 추론 시간이 최대 25% 증가합니다.

중요성 및 주장

이 논문은 LaneRoPE 가 시퀀스 내부 및 외부의 토큰 상대 위치를 근본적으로 활용함으로써 병렬 협력적 추론을 가능하게 하는 유망한 단계라고 주장합니다.

• 아키텍처 효율성: 이전 연구들이 전용 런타임 엔진이나 복잡한 분기 메커니즘을 요구하는 것과 달리, LaneRoPE 는 최소한의 아키텍처 변경을 요구하며 무시할 수 있는 추론 오버헤드를 도입하여 기존 LLM 파이프라인에 신속하게 통합하기에 적합합니다.
• 유연성: 이 공식은 기존 기술 (독립 샘플링 및 GroupThink 등) 을 포괄할 만큼 표현력이 풍부하면서도 새로운 협력 전략을 발견하기 위한 추가 파인튜닝을 허용합니다.
• 데이터 효율성: 이 방법은 간단한 SFT 또는 합성 데이터를 사용한 KTO 를 통해 기본 모델 대비 매우 적은 추가 파라미터 (< 0.5%) 로 잠금 해제될 수 있습니다.

저자들은 한계를 인정하며, 아직 여러 레인에서 출력물을 병합하기 위한 전용 훈련 체계를 사용하지 않으며, 현재 실험은 작은 배치 크기 ($N \le 4$) 에 초점을 맞추고 있음을 지적합니다. 그들은 향후 연구가 협력적 행동을 더욱 강화하기 위해 검증 가능한 피드백을 통한 강화 학습 (RLVF) 을 탐구할 수 있음을 제안합니다.