DepthGraft가 계층적 KV 재구성으로 깊은 층 정보 희석을 줄인다

무엇인가

이 논문은 Pre-LN 트랜스포머에서 층을 깊게 쌓아도 잔차 스트림이 고정된 덧셈 업데이트로 모든 이전 계산을 하나의 은닉 상태 궤적에 모으기 때문에, 초기 층 특징이 깊어질수록 희석되어 후속 층이 꺼내 쓰기 어려워지는 문제를 다룬다. 깊이 확장은 LLM 성능의 핵심 레버지만 이 희석 때문에 추가 층의 효율이 표현 용량만큼 나오지 않는다는 것이 출발점이다. 기존 교차 층 연결은 직접 이전 표현을 읽게 해 완화하지만 계산·메모리 접근·시스템 복잡도를 늘린다.

어떻게 동작하나

저자들은 교차 층 KV 공유를 단순 캐시 압축이 아니라 교차 층 그래디언트 정규화로 본다. 깊은 층이 재사용한 KV로 역전파하면 소스 층의 K/V 투영에 직접 그래디언트가 흐르고, 하나의 소스 KV가 여러 쿼리 공간의 학습 신호를 받는다. 이때 K는 softmax 경쟁 때문에 공유 판별 방향으로 더 날카로워지고, V는 서로 다른 어텐션 패턴의 출력 그래디언트를 선형으로 받아 더 넓은 콘텐츠 부분공간을 채우는 비대칭 전문화가 강화된다고 분석한다. 실제로 학습 중 K의 effective rank는 낮아지고 V는 높아지는 기준선 경향이 관찰되며, 공유는 이를 증폭한다.

무엇과 다른가

공유 효과는 연결 위상에 달려 있다. 인접 층 공유(CLA/MLKV)는 깊이상 가까운 쿼리 공간이 겹쳐 K 랭크를 기준선 아래로 충분히 낮추지 못하고, 여러 층이 같은 메모리를 보면서 유효 깊이를 줄일 수 있다. 전역 앵커 공유(YOCO/FusedKV)는 소수 앵커에 모든 깊은 층의 그래디언트가 몰려 과부하된 절충 표현을 만들고, 앵커가 아닌 층에는 직접 정규화가 잘 퍼지지 않는다. 저자들은 충분한 교차 층 거리와 소수 재사용 층으로 분산된 그래디언트, 그리고 대부분 층의 참여가 필요하다고 정리한다.

어떻게 쓰나

DepthGraft는 L개 층을 대략 L/3씩 얕은·중간·깊은 단계로 나눈다. 얕은 단계와 중간 단계는 KV를 정상 계산하고, 깊은 단계는 독립 KV를 만들지 않고 같은 상대 깊이 위치의 얕은 소스와 중간 소스를 짝지어 재구성한다. 깊은 층 l의 소스는 s(l)=l-2L/3, m(l)=l-L/3이며, K_l=a_l1⊙K_s(l)+a_l2⊙K_m(l), V_l=b_l1⊙V_s(l)+b_l2⊙V_m(l)처럼 KV 그룹별 채널 단위 학습 계수로 두 소스를 섞는다. 키 계수는 RoPE 회전 쌍에 묶여 위치 인코딩과 재구성이 교환 가능하고 표준 KV 캐시 저장 논리를 유지한다. 이 배치는 비국소 그래디언트를 여러 소스 층에 분산하면서 프리필 조기 완료와 PP/VPP 친화성을 보존한다.

전제와 한계

실험은 dense 0.5B·1B·1.5B와 DeepSeekMoE 스타일 3B·9B·16B MoE, GQA와 MLA를 포함한다. dense 평균은 각각 +0.75, +0.85, +2.42점 올랐고, 1.5B에서 MMLU +3.33, MMLU-Pro +1.43, SIQA +2.86, GSM8K +0.91, MBPP +5.20, C3 +9.46이 제시된다. MoE GQA 평균은 3B +1.05, 9B +2.29, 16B +3.61이며, 16B에서 MMLU-Pro +7.00, CMMLU +6.61, C3 +6.41, GSM8K +5.84다. 9B MLA도 평균 +3.68로 어텐션 변형에 묶이지 않음을 보인다.

효율 측면에서 깊은 단계의 독립 KV가 사라져 KV 캐시와 프리필 비용이 약 1/3 줄고, 128K 문맥에서 TTFT가 9B MLA 26.8%, 16B GQA 28.8% 감소하며 TPOT 오버헤드는 전 구간 1% 미만이다. 32K까지 확장한 RULER 평균은 4K에서 94.8 대 86.8, 32K에서 82.4 대 76.2이고, 32K 변수 추적은 37.8에서 76.6, 빈출 단어 추출은 64.0에서 74.7로 오른다. 절제실험에서는 재사용 KV의 그래디언트를 끊은 shallow-copy(detach)가 기준선과 비슷한 42.86(기준 42.38)인 반면 역방향을 살리면 45.30, 계층 재구성은 47.56으로 올라 교차 층 그래디언트가 핵심임을 보인다. 2단계 단일 소스 변형은 43.58이고, 4단계 변형들은 3단계 DepthGraft보다 낮았다.

개발자에게 이 논문은 KV 캐시를 줄이는 추론 최적화가 아니라, 사전학습 단계에서 깊은 층의 메모리 경로를 정규화하는 아키텍처 선택으로 읽힌다. 깊은 LLM 스택에서 KV 캐시와 프리필 지연이 병목이고 GQA/MLA를 이미 쓰는 팀이라면, 깊은 단계에서 KV를 재사용하되 소스 층과의 거리 및 그래디언트 분산을 설계 변수로 삼아야 한다는 시사점이 있다. 다만 추론 시 체크포인트에 붙이는 플러그인이 아니라 모델 구조와 사전학습 레시피를 바꾸는 방식이므로, 재학습 비용과 PP/VPP 배치, RoPE 계수 묶음, 실제 커널 최적화 여부를 함께 검토해야 한다.

저자들이 밝힌 한계는 이론 분석이 명시적 가정 아래 이루어졌고, 더 엄밀한 이론 이해와 더 나은 공유 위상, 더 큰 사전학습 규모로의 확장이 향후 과제라는 점이다. 실험은 최대 16B MoE와 1.5B dense까지이며, 효율 수치는 고정 2K 생성 길이와 특정 9B MLA·16B GQA 설정에서 보고된다. 또한 TTFT 이득은 문맥이 길수록 커지지만 TPOT는 거의 그대로이고, 전용 커널 최적화를 적용하면 이론적 속도 향상과의 격차가 달라질 수 있다고 언급한다.