기술 요약: 임베딩 유사성 정규화를 통한 사전 학습 성능 향상 (SimReg)

1. 문제 제기

대형 언어 모델 (LLM) 은 근본적으로 다음 토큰 예측에 의존하며, 이는 이전 컨텍스트의 결합된 표현에 대한 분류 문제로 간주될 수 있습니다. 교차 엔트로피 손실을 이용한 사전 학습은 놀라운 발전을 이끌었지만, 저자들은 토큰 임베딩의 컨텍스트 의존적 성질이 높은 클래스 내 분산과 클래스 간 유사성을 초래한다는 중요한 한계를 식별했습니다.

특정 객체 (예: 고양이) 가 일관되게 안정적인 레이블과 임베딩에 매핑되는 이미지 분류와 달리, 언어 토큰은 다양한 컨텍스트 특징을 기반으로 예측됩니다. 결과적으로 동일한 토큰 (예: 다른 문장 속의 "walls") 을 예측하는 임베딩은 서로 다른 컨텍스트로 인해 크게 달라질 수 있습니다. 이 논문은 표준 교차 엔트로피 손실이 이러한 임베딩에 강력한 일관성을 부여하지 못한다고 주장합니다. 실증 분석에 따르면, 교차 엔트로피는 초기 학습 단계에서 대비적 유사성을 빠르게 제약하지만, 기본 분리 가능성 임계값에 도달하면 정체됩니다. 모델 크기가 증가하더라도 분류 마진 확대를 지속적으로 주도하거나 토큰 임베딩 간의 충분한 이질성을 강제하지는 않습니다. 이러한 한계는 표현 학습의 효율성을 저해하고 수렴 속도를 늦춥니다.

2. 방법론: SIMREG

이러한 문제를 해결하기 위해 저자들은 대규모 사전 학습 중 임베딩 유사성을 명시적으로 정규화하도록 설계된 손실 함수인 SIMREG(Similarity Regularization) 을 제안합니다.

핵심 메커니즘

시퀀스의 각 토큰에 대해, 이 방법은 기준 토큰의 실제 레이블 (ground-truth label) 을 기반으로 모든 토큰을 양 (+) 그룹과 음 (-) 그룹으로 분할합니다:

• 양 그룹 ($P_i$): 기준 토큰과 동일한 실제 레이블을 공유하는 토큰들.
• 음 그룹 ($N_i$): 다른 레이블을 가진 토큰들.

SIMREG 손실은 양 그룹 내 임베딩 간의 거리를 최소화하면서 음 그룹으로부터의 분리를 최대화하는 것을 목표로 합니다. $i$ 번째 토큰에 대한 손실은 다음과 같이 정의됩니다:
$$L_i = L_{ce}^i + \lambda \cdot \text{softplus}(L_{sr}^i)$$
여기서 $L_{sr}^i$는 유사성 정규화 항입니다:
$$L_{sr}^i \triangleq \log \sum_{j \in N_i} \phi_{i,j} - \log \sum_{j \in P_i} \phi_{i,j}$$
여기서 $\phi$는 온도 계수 $\tau$로 스케일링된 임베딩 간의 코사인 유사성을 나타냅니다. 안정성과 유효성을 보장하기 위해 이 방법은 다음을 포함합니다:

• 자기 유사성 (Self-Similarity): 시퀀스 내의 고유 토큰이라도 적어도 하나의 양 쌍이 존재하도록 기준 토큰을 자신의 양 그룹에 추가합니다.
• Softplus 스케일링: 정규화 손실이 음수가 되지 않도록 보장합니다.
• 그룹 수준 평균화: 토큰 간 기여도를 균형 있게 유지하고 긴 학습 기간 동안 안정성을 유지하기 위해 샘플 수준이 아닌 그룹 수준에서 손실을 계산합니다.

구현 세부 사항

• Chunk-wise 계산: 긴 시퀀스 내 유사성 계산의 $O(n^2)$ 복잡성을 처리하고 과도한 통신 오버헤드 없이 시퀀스 병렬성 (Sequence Parallelism, SP) 을 지원하기 위해 저자들은 시퀀스를 청크 (chunk) 로 분할합니다. 손실은 청크 내부에서 계산되어 복잡성을 $O(n^2/b)$로 줄입니다 (여기서 $b$는 청크 수).
• 배치: 실증 연구에 따르면 SIMREG 를 모델의 최종 레이어에 적용하는 것이 가장 효과적입니다. 중간 레이어는 얽힌 컨텍스트 정보를 인코딩하는 반면, 최종 레이어는 이러한 정보를 예측에 사용되는 고유한 의미 공간으로 투영하기 때문입니다.

3. 이론적 분석

이 논문은 분류 마진에 미치는 영향을 분석함으로써 SIMREG 의 효과에 대한 이론적 근거를 제공합니다.

• 마진 확대: 저자들은 교차 엔트로피 손실이 분류 마진 ($m_i$) 의 지수 함수에 의해 상한이 결정됨을 보여줍니다. SIMREG 손실을 최소화함으로써 토큰 임베딩의 양 중심은 동일 클래스 샘플 쪽으로 이동하고, 음 중심은 멀어집니다.
• 수렴 한계: 이론적 증명에 따르면, 이 과정은 유사성 손실로부터의 교란 $\epsilon^+$에 대해 분류 마진 ($m_k$) 을 적어도 $\delta \|\epsilon^+\| + \sqrt{2}L_P \gamma$만큼 증가시킵니다. 이러한 마진 확대는 교차 엔트로피 손실을 곱셈적으로 감소시켜 이론적으로 사전 학습 과정을 가속화합니다.

4. 주요 기여

1. SIMREG 제안: 대규모 사전 학습 중 클래스 내 일관성과 클래스 간 분리를 명시적으로 장려하는 새로운 임베딩 유사성 정규화 손실로, 교차 엔트로피 단독의 한계를 해결합니다.
2. 이론적 통찰: SIMREG 가 다중 분류 마진을 어떻게 확대하여 교차 엔트로피 손실의 상한을 낮추고 최적화 역학을 개선하는지에 대한 상세 분석.
3. 확장 가능한 구현: 시퀀스 병렬성 환경에서 유사성 정규화의 효율적인 적용을 가능하게 하는 청크 단위 계산 전략으로, 대규모 모델에 적용 가능하게 합니다.
4. 종합적 평가: 밀집형 (LLaMA-350M~7B) 및 전문가 혼합 (Mixtral-8×1B) 아키텍처에 걸친 광범위한 실험.

5. 실험 결과

저자들은 350M 에서 7B 파라미터까지의 모델과 Mixtral-8×1B MoE 모델을 사용하여 표준 벤치마크에서 SIMREG 를 평가했습니다.

• 학습 가속화: SIMREG 는 학습 수렴을 일관되게 가속화합니다. LLaMA 모델에서 기준선과 동일한 교차 엔트로피 손실 수준에 도달하는 데 30% 이상 (7B 모델의 경우 최대 38%) 의 속도 향상을 달성합니다.
• 하류 작업 성능: 이 방법은 표준 벤치마크에서 하류 작업의 평균 제로샷 성능을 1% 이상 향상시킵니다 (예: LLaMA-1.3B 에서 +1.14%, LLaMA-7B 에서 +1.22%).
• 임베딩 분리 가능성: 분석 결과, 교차 엔트로피 단독 사용 시 평균 클래스 간 각도가 약 61.3 도인 반면, SIMREG 를 추가하면 이 각도가 약 74 도로 증가하여 훨씬 더 우수한 특징 분리 가능성을 보여줍니다.
• 효율성: 전체 SIMREG 손실은 실행 시간을 2% 미만으로, GPU 메모리를 1% 미만으로만 증가시킵니다. 청크 단위 변형은 거의 자원 중립적입니다.
• 하이퍼파라미터 민감도: 이 방법은 광범위한 하이퍼파라미터 범위에서 견고합니다. 최적 온도 $\tau$는 0.01 에서 안정적이며, $\lambda$는 모델 임베딩 차원에 따라 효과적으로 스케일링됩니다 (약 $\sqrt{d}$에 비례).

6. 중요성 및 주장

이 논문은 SIMREG 가 LLM 사전 학습의 효율성과 일반화를 향상시키는 단순하면서도 광범위하게 적용 가능한 전략을 제공한다고 주장합니다. 명시적으로 임베딩 일관성을 강제함으로써 표준 교차 엔트로피 학습에서 관찰된 표현 학습의 정체성을 극복합니다. 저자들은 이 작업을 최소한의 추가 계산 비용으로 훈련 속도와 하류 작업 성능 모두에서 상당한 개선을 제공하는 실용적인 진전으로 위치시킵니다. 그들은 일관성 정규화가 주로 파인튜닝이나 소규모 설정에서 탐구되어 왔으나, 대규모 모델의 사전 학습 단계로 확장될 때 매우 효과적임을 강조합니다.