DAGFormer가 입력 토큰으로 Transformer 전층 연결 위상을 예측한다
DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections
무엇인가
이 논문은 표준 Transformer의 잔차 연결이 모든 층에 고정되어 있다는 문제를 다룬다. 각 어텐션 헤드와 MLP가 같은 running sum을 읽고 쓰기 때문에, 늦은 층은 이전 층의 결과를 모든 중간 층이 더해진 합 안에서만 본다. DenseFormer, hyper-connections, mHC, MUDDFormer 같은 기존 연구가 연결을 완화하지만, 동적 변형들은 대부분 층마다 별도 라우터를 두고 자기 층이 무엇을 읽을지 국소적으로 정한다. 그래서 전체 연결 위상은 층별로 탐욕적으로 조립되고, 입력으로부터 전역 계획을 세우거나 위상 자체를 하나의 객체로 검사하기 어렵다. DAGFormer는 백본이 토큰을 처리하기 전에 입력 토큰만 읽는 별도 예측기가 전체 층의 연결 위상을 한 번에 계획하도록 이 문제를 바꾼다.
어떻게 동작하나
방법의 핵심은 작은 topology predictor다. 예측기는 백본과 다른 자체 256차원 토큰·위치 임베딩을 쓰고, 4개 헤드를 가진 2개 causal Transformer 층으로 토큰을 인코딩한다. 각 출력은 GELU를 거치는 512차원 공유 trunk를 지나며, 라우팅되는 층마다 하나의 선형 헤드가 그 층의 adjacency matrix를 만든다. 이 행렬은 각 어텐션 헤드의 query, key, value 읽기와 residual 읽기에 대해 이전 출력 X0부터 Xl-1까지의 혼합 계수를 담는다. 계수는 softmax로 정규화되지 않은 부호 있는 실수라서 특정 소스를 증폭하거나 다른 소스에서 빼는 것도 가능하다. 선형 헤드의 가중치는 0에서 시작하고 bias의 각 행은 Xl-1만 선택하는 one-hot으로 초기화되므로, 학습 초기에는 정확히 기존 OLMo-2 Transformer처럼 동작한다. 예측기는 causal이어서 토큰의 그래프가 미래 토큰에 의존하지 않고, 모든 층이 같은 시퀀스 인코딩에서 계획되며 더 깊은 백본은 선형 헤드만 추가한다. 여기에 각 층이 자기 입력 Xl-1로부터 작은 네트워크로 계획된 가중치를 보정하는 local refinement가 붙는다. 이 보정은 W2를 0으로 초기화하고, 헤드와 읽기 단위로 예측기의 계획을 수정한다.
무엇과 다른가
백본은 이 그래프를 따라 계산한다. 라우팅되는 층 l에서 각 헤드의 query, key, value는 이전 출력들의 가중합으로 만들고, residual Rl도 이전 출력들의 가중합으로 만든다. 즉 qh는 sum over j<l of alpha_q,h,j 곱하기 Wq 곱하기 Xj이고, kh와 vh도 같은 방식이며, Rl은 alpha_r,j로 Xj를 섞는다. 어텐션의 Q/K 정규화, MLP, 층 출력 Xl=Rl+al+ml 같은 나머지 구조는 표준 Transformer와 같다. 300M와 1B에서는 혼합된 value에 RMSNorm을 추가한다. 예측기와 백본은 언어모델링 손실로 처음부터 공동 학습된다. 라우팅은 29M에서 37M 파라미터를 추가하며 그중 25.7M이 예측기의 토큰 임베딩이다. 학습 순서 기준 FLOPs는 75M에서 기존 Transformer의 1.18배, 1B에서 2.13배이고, prefill 시간은 1.5배에서 1.9배, 예측기 자체는 그중 1~3%다.
어떻게 쓰나
실험은 동일한 백본, 토크나이저, 데이터 순서, 토큰 예산, 옵티마이저를 맞춘 OLMo-2 쌍으로 한다. 크기는 75M 6층 D512 8헤드, 150M 8층 768 12헤드, 300M 12층 1024 16헤드, 1B 16층 2048 16헤드다. 데이터는 Dolma v1.7이고 시퀀스 길이는 1024, 스텝당 524K 토큰이다. 75M, 150M, 300M는 파라미터당 약 21토큰을 학습하고, 1B는 5B 토큰에서 10B 토큰으로 이어 학습한다. 비교 대상은 vanilla OLMo-2 Transformer, DenseFormer, 4스트림 dynamic hyper-connections, mHC, MUDDFormer다. 평가는 WikiText-2를 주 지표로 MathInstruct, GSM8K, 균형 Dolma 홀드아웃을 쓰고, 12개 lm-evaluation-harness 과제와 GSM8K를 합친 13개 과제, Alpaca와 Dolly 및 smol-smoltalk 지시 미세조정, AlpacaEval과 MT-Bench를 포함한다.
전제와 한계
주요 결과는 WikiText-2 퍼플렉시티가 모든 크기에서 vanilla Transformer보다 17~25% 낮다는 것이다. 손실로는 모든 크기에서 0.18~0.29 nats 낮고, 75M, 300M, 1B에서는 가장 강한 라우팅 기준선보다 0.1~0.2 nats 낮다. 150M에서는 MUDDFormer가 0.04 nats 더 낮다. 1B DAGFormer를 5B 토큰으로 학습한 모델이 약 같은 학습 연산의 10B 토큰 vanilla 1B를 WikiText-2, MathInstruct, GSM8K에서 앞선다. 13개 평가 평균 정확도는 300M와 1B에서 vanilla보다 2.0~2.5점, 가장 강한 라우팅 기준선보다 0.7~0.8점 높다. 수학에서는 mHC가 0.1점 앞선다. 지시 미세조정 후 홀드아웃 NLL은 vanilla보다 0.15~0.19 nats, 가장 강한 라우팅 기준선보다 0.04~0.06 nats 낮다. AlpacaEval에서 vanilla 대비 승률은 0.63~0.70이고 라우팅 기준선은 0.46~0.64다. MT-Bench 첫 턴에서는 0.56~0.69 승률이고 10개 중 7개에서 95% 구간이 0.5를 배제한다.
분석에 따르면 동일 연산에서 vanilla Transformer가 DAGFormer의 손실에 도달하려면 75M와 150M에서는 1.2~1.4배, 300M와 1B에서는 1.7~2.0배 학습 연산이 필요하다. 피팅된 지수는 파라미터 기준 0.56 대 0.48, 연산 기준 0.30 대 0.27이고, 외삽하면 10^20 FLOPs에서 1.9배, 10^21 FLOPs에서 2.3배가 되지만 두 지수 차이는 약 1 표준오차다. 추가 파라미터는 75M에서 격차의 70~89%를 설명하지만 150M에서는 4~42%만 설명한다. 라우팅 소스 절제에서 예측기만으로 약 90%의 이득을 회복하고, 같은 토큰을 읽는 독립 층별 예측기는 7M 더 많은 파라미터로 80% 미만을 회복한다. local refinement는 순전파 MACs의 0.5% 미만이다. 구조적 프루닝에서는 6개 설정 중 5개에서 DAGFormer의 성능 저하가 가장 작다. 헤드와 채널을 50% 제거하면 모든 크기에서 vanilla와 모든 라우팅 기준선보다 덜 나빠지고, 70% 제거에서는 150M와 300M에서 가장 덜 나빠진다. 75M에서는 vanilla, DenseFormer, hyper-connection 변형이 덜 나빠지지만 DAGFormer의 프루닝 손실 자체는 가장 낮다. MUDDFormer는 모든 크기와 희소도에서 가장 많이 나빠진다. 프루닝 중 예측기나 라우터를 고정해도 결과는 거의 바뀌지 않는다.
개발자 관점에서 이 논문은 층 간 연결을 학습 가능한 전역 위상으로 다루는 설계를 실험할 때 참고가 된다. 예측기 출력이 백본 실행 전에 고정되므로 그래프를 평균내거나 고정하거나 입력 간에 교체하는 개입이 가능하고, 이는 해석과 통제를 쉽게 만든다. 다만 이득의 대부분은 토큰별 재계획이 아니라 학습된 그래프 자체에서 나온다. 예측기 출력을 위치별 평균으로 바꿔도 거의 손해가 없고, 입력 의존성은 그래프를 학습하는 동안 주로 중요하며, 추론 시에는 refinement의 토큰별 보정에 의존한다. 따라서 실무에서는 파라미터 수만 맞추지 말고 라우팅 비용을 포함한 동일 연산 비교를 해야 하고, fused projection과 메모리 병목, 고정 위상으로 배포하려면 재학습이 필요한지 확인해야 한다.
저자들이 밝힌 한계도 분명하다. 가장 큰 모델이 1B이고 5B에서 10B 토큰만 학습해 compute-optimal 예산보다 낮다. 300M 쌍을 제외하면 사전학습은 단일 시드이고, 두 번째 시드에서 홀드아웃 NLL이 최대 0.08 nats per token 변했다. 이는 vanilla의 GSM8K solutions에서였고 WikiText-2에서는 최대 0.03이라 작은 차이는 노이즈다. 프루닝 셀도 단일 실행이다. 라우팅은 학습 순서에서 vanilla FLOPs의 1.18~2.13배를 쓰고, 학습된 모델은 refinement의 토큰별 보정에 의존해 고정 위상으로는 재학습 없이 실행할 수 없다. MUDDFormer는 150M에서 DAGFormer를 이기고 다른 크기에서는 지지만 각각 단일 시드라 손실만으로 결론이 나지 않는다. 1B 기준선은 5B 토큰만 학습했고, 파라미터 매칭 vanilla 통제는 75M와 150M에만 있으며, 학습률 스윕은 75M 쌍에서만 했다. 평가에서는 이 규모 모델이 어려운 벤치마크에서 우연 수준에 가깝고, 채팅과 AlpacaEval 비교를 단일 LLM 판정자가 채점했다. 해석 가능성도 그래프 전체를 집계해 읽는 수준이며, 입력 의존 부분을 사람이 이름 붙일 수 있는지와 1B 초과 규모에서 이득이 유지되는지는 후속 연구로 남겨졌다.