Differential Transformer V2

🖼️ 인포그래픽

Differential Transformer V2 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

DIFF V2는 같은 GQA 그룹의 두 쿼리 헤드 출력을 토큰·헤드별 계수로 차감해, 키·값 헤드와 표준 어텐션 커널을 그대로 유지하면서 디코딩 효율, 학습 안정성, 표현 자유도, 출력 투영의 매개변수 효율을 함께 개선한 구조다.

📌 핵심 요약

DIFF V2는 쿼리 헤드 수를 기준 트랜스포머의 두 배로 늘리되 키·값 헤드 수는 유지하고, 같은 GQA 그룹에서 키와 값을 공유하는 두 헤드의 출력을 차감한 뒤 다시 기존 출력 차원으로 줄인다.
DIFF V1의 전역적 계수와 헤드별 RMS 정규화를 없애고, 입력에서 투영한 토큰별·헤드별 계수에 시그모이드를 적용해 두 번째 어텐션 출력을 조절함으로써 표준 플래시 어텐션 커널을 그대로 사용할 수 있게 했다.
일반 소프트맥스 어텐션의 문맥 RMS는 값 벡터가 비상관이고 RMS가 1이라는 가정에서 1/√n 이상으로 제한되지만, DIFF V2의 차분 연산은 하한을 0까지 낮춰 어텐션 싱크를 완화하고 학습 안정성을 높일 수 있다.
진행 중인 대규모 사전학습 실험에서는 1조 토큰 시점에 기준 트랜스포머보다 언어 모델링 손실이 0.02~0.03 낮았으며, 큰 학습률에서 손실·그래디언트 급등과 활성값 이상치의 크기도 감소했다.
동일한 2h·d 쿼리 차원을 가진 트랜스포머와 비교하면 어텐션 커널 속도는 같지만 DIFF V2는 차분 연산을 출력 투영 전에 명시적으로 수행해 출력 투영 매개변수와 연산량을 줄이며, 현재 GQA 설정에서는 어텐션 모듈 매개변수의 약 25%를 절감할 수 있다고 설명한다.

🧩 주요 포인트

DIFF V2는 쿼리 헤드 수를 기준 트랜스포머의 두 배로 늘리되 키·값 헤드 수는 유지하고, 같은 GQA 그룹에서 키와 값을 공유하는 두 헤드의 출력을 차감한 뒤 다시 기존 출력 차원으로 줄인다.
DIFF V1의 전역적 계수와 헤드별 RMS 정규화를 없애고, 입력에서 투영한 토큰별·헤드별 계수에 시그모이드를 적용해 두 번째 어텐션 출력을 조절함으로써 표준 플래시 어텐션 커널을 그대로 사용할 수 있게 했다.
일반 소프트맥스 어텐션의 문맥 RMS는 값 벡터가 비상관이고 RMS가 1이라는 가정에서 1/√n 이상으로 제한되지만, DIFF V2의 차분 연산은 하한을 0까지 낮춰 어텐션 싱크를 완화하고 학습 안정성을 높일 수 있다.
진행 중인 대규모 사전학습 실험에서는 1조 토큰 시점에 기준 트랜스포머보다 언어 모델링 손실이 0.02~0.03 낮았으며, 큰 학습률에서 손실·그래디언트 급등과 활성값 이상치의 크기도 감소했다.
동일한 2h·d 쿼리 차원을 가진 트랜스포머와 비교하면 어텐션 커널 속도는 같지만 DIFF V2는 차분 연산을 출력 투영 전에 명시적으로 수행해 출력 투영 매개변수와 연산량을 줄이며, 현재 GQA 설정에서는 어텐션 모듈 매개변수의 약 25%를 절감할 수 있다고 설명한다.

🧠 상세 정리

1. DIFF V1에서 V2로 바뀐 핵심 구조

DIFF V1은 두 쌍의 쿼리와 키로 어텐션 출력을 각각 계산한 뒤, 학습 가능한 벡터들의 내적과 층별 초기값으로 구성한 전역적 계수를 사용해 두 번째 출력을 차감한다. 차감 후에는 헤드별 RMS 정규화를 적용하고 층별 초기값에 따른 배율을 다시 곱하며, 값 헤드의 차원도 일반적인 쿼리·키 헤드 차원과 다르게 두 배로 구성된다. DIFF V2는 하나의 표준 플래시 어텐션 호출로 2h개의 쿼리 헤드 출력을 얻은 다음 짝수·홀수 헤드로 나누고, 입력 X에서 투영한 토큰별·헤드별 계수에 시그모이드를 적용해 두 번째 출력을 차감한다. 이때 서로 차감되는 두 헤드는 같은 GQA 그룹에 속하므로 동일한 키와 값을 공유하며, 원문은 이 배치가 성능에 결정적으로 중요하다고 강조한다. 차분 연산 뒤 차원은 다시 h·d로 줄어들기 때문에 최종 출력 투영의 형태는 기준 트랜스포머와 동일하게 유지된다.

2. 빠른 디코딩과 표준 커널 호환성

DIFF V2는 기준 트랜스포머보다 쿼리 헤드를 추가하지만 키·값 헤드 수는 늘리지 않도록 설계됐다. 대규모 언어 모델의 디코딩은 통상 메모리 대역폭의 제약을 크게 받기 때문에, 키·값 캐시의 규모와 로드량을 늘리지 않는 이 구조는 표준 트랜스포머와 비슷한 디코딩 속도를 목표로 할 수 있다. 쿼리·키·값의 헤드 차원도 서로 정렬되어 있어 DIFF V1과 달리 별도의 맞춤형 어텐션 커널이 필요하지 않으며, 디코딩 중 어텐션 모듈의 산술 집약도를 높일 수도 있다. H 계열과 B 계열 GPU에서 최신 플래시 어텐션 커널을 사용한 사전학습에서는 DIFF V2가 추가하는 처리량 감소가 미미하다고 보고한다. 동일한 2h·d 쿼리 차원의 트랜스포머와 비교하면 어텐션 커널 속도는 같고, DIFF V2 쪽이 출력 투영의 매개변수와 부동소수점 연산량을 더 적게 사용한다.

3. 소프트맥스가 문맥 벡터 크기에 만드는 제약

표준 스케일드 내적 어텐션에서 문맥 벡터는 소프트맥스로 얻은 비음수 가중치로 값 벡터들을 합한 결과이며, 각 행의 가중치 합은 항상 1이다. 값 벡터들이 서로 비상관이고 각각의 RMS가 1이라고 가정하면, 문맥 벡터의 RMS는 어텐션 분포에 따라 1/√n 이상 1 미만의 범위에 놓인다. 한 토큰에 가중치가 집중되면 RMS는 1에 가까워지고, 모든 토큰에 1/n씩 균등하게 분산되면 RMS는 1/√n까지 낮아진다. DIFF V1의 헤드별 RMS 정규화는 이렇게 작아진 출력을 다시 RMS 1로 확대하므로, 길이 8,192의 균등 분포에서는 약 90.5배에 달하는 배율이 필요하다. 원문은 대규모 사전학습에서 이러한 약 100배 수준의 증폭이 거대한 그래디언트와 수치 불안정을 일으키며, 특히 큰 학습률을 사용할 때 후반부 그래디언트 노름의 증가와 분산이 기준 트랜스포머보다 커지는 현상을 관찰했다고 설명한다.

4. RMS 정규화 제거와 차분 계수의 역할

DIFF V2는 헤드별 RMS 정규화를 제거하고, 각 토큰과 각 헤드마다 입력에서 투영한 계수를 사용해 문맥 출력의 크기를 직접 조절한다. 구체적으로 두 소프트맥스 어텐션 분포 가운데 두 번째 분포에 시그모이드 계수를 곱한 후 첫 번째 분포에서 빼며, 이에 따라 차분된 문맥 RMS의 범위는 0 초과 √2 미만으로 확장된다. 저자들이 특히 중요하게 보는 변화는 상한의 확대가 아니라 문맥 RMS의 하한을 0까지 낮출 수 있게 된 점이며, 이는 어텐션 싱크 제거와 학습 안정성 향상에 도움이 된다고 보고한다. 분석 대상은 출력 투영 이전의 RMS로 한정되지만, 출력 투영 이후에 크기를 복원하거나 조절할 수 있다는 사실만으로는 소프트맥스 단계의 표현 자유도 부족이 해소되지 않는다고 지적한다. DIFF V2에서는 값 헤드 차원과 계수의 적용 방식이 V1에서 달라졌기 때문에 RMS 정규화를 제거해도 안전했고, 실제 그래디언트 노름 규모도 기준 트랜스포머와 비슷해지면서 급등 현상이 줄었다.

5. 소프트맥스 제약을 완화하는 다른 접근과의 비교

원문은 문맥 RMS의 하한을 0으로 낮추려는 관련 접근으로 ‘Attention Is Off By One’, gpt-oss의 헤드별 스칼라, 그리고 게이트 어텐션을 함께 제시한다. ‘Attention Is Off By One’은 소프트맥스 분모에 1을 더해 전체 어텐션 출력에 0과 1 사이의 축소 계수가 생기게 하며, 그 결과 문맥 RMS도 0 초과 1 미만의 범위를 가질 수 있다. gpt-oss는 각 헤드에 학습 가능한 스칼라를 두고 그 지수값을 분모에 추가해 비슷한 전역 축소 효과를 구현하며, 게이트 어텐션은 입력에서 투영한 원소별 시그모이드 게이트를 문맥 벡터에 곱한다. 이 접근들은 모두 표준 소프트맥스의 가중치 합이 항상 1이라는 제약을 완화해 출력 크기를 0에 가깝게 만들 수 있다는 공통점이 있다. DIFF V2는 하나의 어텐션 출력에 단순한 축소 계수를 곱하는 대신, 동일한 키와 값을 공유하는 두 소프트맥스 어텐션 출력을 토큰별·헤드별 계수로 차감한다는 점에서 구성 방식이 다르다.

6. 대규모 사전학습에서 관찰된 결과

저자들은 밀집 모델과 30A3 혼합전문가 모델을 포함한 실제 운용 규모의 언어 모델을 수조 개 토큰으로 사전학습하며 DIFF V2를 평가하고 있다. 실험에 사용한 학습률은 6×10⁻⁴에서 1×10⁻³ 범위로 크며, 원문 작성 시점에도 학습은 계속 진행 중이라고 명시한다. 현재까지 가장 두드러진 결과는 1조 개 학습 토큰 시점에서 기준 트랜스포머보다 언어 모델링 손실이 0.02~0.03 낮게 나타난 것이다. 기준 트랜스포머가 불안정해지는 큰 학습률 조건에서 손실과 그래디언트의 급등도 줄었고, 활성값 이상치의 크기 역시 감소했다고 보고한다. 다만 중간학습과 후속학습의 학습 효율, 그리고 문맥 부패 완화를 포함한 장문맥 하위 벤치마크 성능은 향후 학습 단계에서 살펴볼 대상으로 남겨 두었으므로, 원문에는 이에 대한 확정적인 결과가 제시되지 않았다.

7. 차분 연산을 명시적으로 구성하는 이유

이론적으로 2h개의 어텐션 헤드를 가진 표준 트랜스포머도 같은 GQA 그룹의 인접한 두 헤드에 대응하는 출력 투영 가중치를 서로 정확한 음수로 학습하면 차분 연산을 구현할 수 있다. 그러나 저자들의 첫 번째 가정은 두 개의 독립된 매개변수 집합이 정확히 반대값으로 수렴해야 하므로, 실제 최적화 과정에서 이러한 해를 학습하기 어렵다는 것이다. 두 번째 가정은 차분 연산 자체가 모델에 유용하고 모델이 학습 과정에서 그것을 선택한다는 것이며, 이 가정 아래에서는 연산을 출력 투영 전에 구조적으로 명시하는 편이 더 효율적이다. DIFF V2는 차분을 먼저 수행해 출력 투영에 들어가는 헤드 차원을 절반으로 줄이므로 출력 투영 매개변수의 절반을 절약할 수 있다. 현재 GQA 구성에서는 어텐션 모듈의 매개변수가 주로 쿼리 투영과 출력 투영에 집중되어 있어, 이 절감량은 전체 어텐션 모듈 매개변수의 약 25%에 해당한다고 설명한다.

8. 현재 결론과 적용 범위의 경계

DIFF V2의 설계상 핵심은 쿼리 계산을 늘리면서도 디코딩 병목에 직접 영향을 주는 키·값 헤드와 캐시는 늘리지 않고, 표준 플래시 어텐션 인터페이스를 유지하는 데 있다. 장문 시퀀스의 프리필 자체를 선형 복잡도로 바꾸는 구조라고 주장하지는 않으며, 이 구간에는 이미 시퀀스 길이에 대해 선형 복잡도를 제공하는 YOCO 같은 기법과 결합할 것을 권장한다. 사전학습 처리량 감소가 미미하다는 관찰도 H 계열과 B 계열 GPU에서 최신 플래시 어텐션 커널을 사용한 조건을 전제로 제시되어 있다. 학습 손실, 그래디언트 급등, 활성값 이상치에 관한 초기 결과는 긍정적이지만 실험이 아직 진행 중이고, 중간학습·후속학습 및 장문맥 벤치마크 결과는 공개된 결론이 아니라 추후 평가 항목이다. 따라서 원문이 현재 뒷받침하는 범위는 구조적 효율성과 사전학습 중간 관찰까지이며, 이후 학습 단계와 하위 과제의 최종 성능은 아직 확인되지 않았다.

🧾 핵심 주장 / 시사점

DIFF V2의 디코딩 효율을 좌우하는 핵심은 쿼리 헤드 증가 자체보다 키·값 헤드와 캐시를 늘리지 않는 데 있으며, 같은 GQA 그룹의 두 헤드가 키와 값을 공유하도록 배치하는 것이 필수적이다.
문맥 RMS의 상한을 크게 만드는 것보다 하한을 0까지 낮추는 자유도가 중요하다는 분석은, 소프트맥스 가중치 합이 1로 고정되는 제약과 어텐션 싱크 및 학습 불안정을 하나의 관점에서 연결한다.
차분 연산을 모델이 출력 투영에서 우연히 학습하도록 두지 않고 구조적으로 명시하면, 정확한 음수 가중치 쌍을 찾아야 하는 최적화 부담을 줄이는 동시에 출력 투영 매개변수와 연산량도 절감할 수 있다.

✅ 액션 아이템

DIFF V2의 두 쿼리 헤드 차분 후 출력 차원 복원 경로를 기준 트랜스포머와 동일 조건에서 비교해 파라미터 절감 폭을 계산한다.
DIFF V1의 전역 계수·헤드별 RMS 정규화를 제거한 토큰별·헤드별 시그모이드 계수 방식이 표준 플래시 어텐션 커널을 유지하는지 호환성 검증을 수행한다.
1조 토큰 학습 실험에서 LM 손실 0.02~0.03 개선과 큰 학습률 구간의 그래디언트 급등·활성값 이상치 감소 동시 발생 여부를 실험 설계와 지표로 재확인한다.