Mitigating Position Bias in Transformers via Layer-Specific Positional Embedding Scaling

기술 요약: 레이어별 위치 임베딩 스케일링을 통한 트랜스포머의 위치 편향 완화

문제 정의
Rotary Position Embeddings(RoPE)를 사용하는 것을 포함하여, 트랜스포머 기반의 대규모 언어 모델(LLM)은 "lost-in-the-middle(중간에서 길을 잃는)" 현상을 겪습니다. 컨텍스트 길이가 길어짐에 따라, 이러한 모델들은 입력 시퀀스의 시작과 끝 부분의 토큰에는 과도하게 주의(attention)를 기울이는 반면, 중간에 위치한 중요한 정보는 소홀히 하는 경향이 있습니다. Attention Buckets 및 MoICE와 같은 기존 솔루션들은 서로 다른 베이스(base) 또는 스케일링 계수를 가진 여러 개의 RoPE를 결합하여 이를 완화하려고 시합니다. 그러나 이러한 접근 방식은 일반적으로 수동으로 설계된 규칙에 의존하며, 추론 중에 여러 번의 포워드 패스(모델 레벨 또는 모듈 레벨 앙상블)를 필요로 합니다. 이는 상당한 계산 오버헤드, 지연 시간 증가, 높은 추론 비용을 초로하여 실질적인 적용을 제한합니다. 또한, 최적의 스케일링 계수를 결정하는 것은 조합 최적화 문제로서, 전수 조사를 위한 계산이 불가능할 정도로 복잡하며, 레이어별 할당의 이산적(discrete) 특성 때문에 경사 하강법(gradient-based methods)을 통해 해결하기 어렵습니다.

방법론
저자들은 모델 파라미터를 수정하거나 추론 지연 시간을 늘리지 않고도 트랜스포머의 각 레이어에 별도의 스케일링 계수를 할당하는 학습 불필요(training-free) 방식인 **레이어별 위치 임베딩 스케일링(Layer-Specific Positional Embedding Scaling, LPES)**을 제안합니다.

1. 레이어별 스케일링: 단일 글로벌 스케일링 계수를 적용하거나 어텐션 헤드당 여러 개의 베이스를 사용하는 방식과 달리, LPES는 각 레이어 $h$에 고유한 스케일링 계수 $s_h$를 할당합니다. 이를 통해 모델은 단 한 번의 포워드 패스 내에서 여러 개의 스케일링된 RoPE로부터 정보를 결합할 수 있습니다.
2. 베지에 곡선(Bézier Curve) 제약: $L$개의 독립적인 스케일링 계수를 찾는 것($M^L$ 크기의 탐색 공간)의 난해함을 해결하기 위해, 저자들은 베지에 곡선을 사용하여 레이어 깊이와 스케일링 계수 사이의 관계를 모델링합니다. 소수의 이산적인 제어점 $P_k$로 곡선을 정의함으로써, 탐색 공간을 $M^L$에서 $(M \times L)^C$ ($C$는 제어점의 수)로 획기적으로 줄입니다. 이 접근 방식은 매끄럽고 연속적인 스케일링 계수의 변화가 유익한 귀납적 편향(inductive bias)으로 작용하여 모델의 레이어별 표현 구조를 보존한다는 관찰 결과에 기반합니다.
3. 곡선 제약 유전 알고리즘: 저자들은 베지에 곡선의 제어점을 최적화하기 위해 특화된 유전 알고리즘을 사용합니다.
   • 적합도 함수(Fitness Function): 알고리즘은 관련 정보가 시작, 중간, 끝 부분에 나타나는 세 가지 서브 데이터셋을 기준으로 개체의 성능을 평가합니다.
   • 제약 조건: 변이(mutation) 및 교차(crossover) 연산자는 제어점 x-좌표의 단조성(monotonicity)을 강제하고, y-좌표 변화 범위를 제한하여 매끄러움을 유지하고 급격한 스케일링 변화를 방지합니다.
   • 효율성: 최적의 스케일링 구성은 표준 하드웨어(4개의 H100 GPU)에서 단 몇 백 개의 예시(예: 200개 인스턴스)만을 사용하여 3~4시간 내에 찾을 수 있습니다.

주요 기여

• LPES 프레임워크: 모델 가중치를 미세 조정하거나 추가적인 추론 지연 시간을 발생시키지 않고도, 레이어별 스케일링 계수를 할당하여 균형 잡힌 어텐션 분포를 달려 균형 잡힌 어텐션 분포를 달성하는 새로운 방법입니다.
• 효율적인 최적화: 베지에 곡선에 의해 제약된 유전 알고리즘을 도입하여, 최적의 스케일링 계수를 찾는 조합 최적화 문제를 효율적으로 해결하고 탐색 공간을 줄이며 빠른 수렴을 가능하게 했습니다.
• 실증적 검증: 다양한 벤치마크와 모델 아키텍처에 걸쳐 일반적인 능력을 보존하면서도 롱 컨텍스트 성능을 크게 향상시킨다는 것을 입증하는 포괄적인 실험을 수행했습니다.

결과
Vicuna-7B, StableBeluga-7B, Qwen2.5-7B 모델을 대상으로 MDQA, Key-Value Retrieval, ZeroSCROLLS, L-Eval과 같은 벤치마크에서 광범위한 실험을 진행했습니다.

• 성능 향상: LPES는 (Positional Interpolation, Attention Buckets, Ms-PoE, MoICE를 포함한) 베이스라인들을 일관되게 능가했습니다. 특히, Vicuna-7B의 Key-Value Retrieval 데이터셋에서 최대 11.2%의 정확도 향상을 달ach했습니다.
• 위치 편향 완화: 이 방법은 컨텍스트 중간(50% 위치)에 위치한 정보에 대한 정확도를 크게 개선하여, 기존 방법들보다 "lost-in-the-middle" 문제를 더 효과적으로 해결했습니다.
• 추론 효율성: LPSE는 단 한 번의 포워드 패스만을 요구합니다. MoICE보다 2.42배, Ms-PoE보다 1.45배 빠르며, 베이스라인과 동일한 추론 시간을 유지하는 것으로 보고되었습니다.
• 일반화: 한 데이터셋(MDQA)에서 학습된 스케일링 계수는 MMLU 및 C-Eval 벤치마크에서의 일반적 능력을 저하시키지 않으면서, 다른 롱 텍스트 작업(ZeroSCROLLS, L-Eval) 및 다른 모델로 효과적으로 전이되었습니다.
• 귀납적 편향: 표현 구조 편차 분석 결과, 매끄러운 베지에 기반 스케일링은 노이즈가 있거나 섞인(shuffled) 스케일링 전략에 비해 모델의 내부 레이어별 표현에 미치는 영향이 최소화됨을 보여주었습니다.

의의
본 논문은 LPES가 롱 컨텍스트 LLM의 위치 편향 문제에 대해 폭넓게 적용 가능하고 효율적인 솔루션을 제공한다고 주장합니다. 모델 전체의 미세 조정에 따른 계산 비용과 멀티 패스 추론에 따른 지연 시간 없이도, LPS는 이미 학습된 모델에 롱 컨텍스트 능력을 배포할 수 있게 합니다. 저자들은 베지에 곡선을 통해 어려운 조합 탐색 문제를 다룰 수 있는 최적화 문제로 전환함으로써, 이전의 앙상블 기반 방법들이 가졌던 지연 시간 페널티 없이 균형 잡힌 어텐션 분포를 구현하는 실질적인 경로를 제공한다고 강조합니다. 이 연구는 매끄러운 레이어별 스케일링이 강력한 귀납적 편향으로서 작론하며, 고성능 구성으로서 자연스럽게 나타난다는 점을 부각합니다.