상하이자오퉁대 연구진이 JTok과 xHC로 Transformer 확장의 새 축을 제안했다

雷锋网2026-11-04

상하이자오퉁대학교 인공지능대학의 Yan Junchi 교수가 CNCC 2026 발표에서 Transformer 구조를 크게 바꾸지 않으면서 모델 용량을 늘리는 두 가지 방법 JTok과 xHC를 소개했다. 두 방법에 공통으로 깔린 질문은 하나다. 새로 붙인 용량이 적은 추가 연산으로 계속 학습 이득을 주는가.

JTok은 토큰 번호를 인덱스로 쓰는 학습 가능한 파라미터 테이블을 Transformer 백본 옆에 붙인다. 토큰마다 고유 번호가 있으니 그 번호로 벡터를 조회하고, 조회한 벡터를 정규화한 뒤 차원별 스케일을 곱해 변조 신호를 만든다. 이 신호를 전방향 네트워크(FFN)의 잔차 증분에 게이트로 걸어 다시 백본에 더하는 구조다. 핵심은 늘어난 파라미터가 대규모 행렬곱을 타지 않는다는 점이다. 조회와 원소별 연산만 하면 되고 백본과 병렬로 계산할 수 있어 추론 구현 부담도 작다. 여기에 라우팅을 더해 일부 분기만 동적으로 활성화하는 JTok-M 변형도 함께 제안됐다.

발표에 따르면 JTok의 손실 곡선은 여러 규모의 dense·MoE 베이스라인에서 계속 아래에 머물렀고, 학습 후반부에도 격차가 사라지지 않았다. 하류 과제 평가에서도 손실 감소만으로 설명되지 않는 개선이 나타났다는 설명이다. IsoFLOPs 분석에서는 같은 연산 예산에서 손실이 약 2.2% 낮았고, 같은 손실에 도달하는 데 필요한 학습 연산은 베이스라인의 약 65% 수준이었다. 다만 구현상 추론 지연이 약 7% 늘었고 파라미터 테이블의 저장과 대역폭 부담도 남는다고 연구팀은 밝혔다.

xHC(Expanded Hyper-Connections)는 잔차 스트림을 넓히는 쪽이다. 기존 Hyper-Connections가 잔차 경로를 여러 갈래로 늘리고 학습 가능한 매핑을 넣었다면, 연구팀은 왜 늘 네 갈래인지에 주목했다. 스트림을 늘릴 때 생기는 학습 불안정은 매핑 행렬을 행·열 합이 1인 이중 확률 행렬로 제약하고 Sinkhorn 정규화를 적용해 완화한다. 이렇게 만든 mHC도 4갈래에서 8갈래로 늘리면 이득이 줄고 16·32갈래는 비용이 급증한다. 원인은 두 가지로 분석됐다. 추가된 경로의 갱신이 서로 강하게 상관돼 실효 차원이 늘지 않는 정보 중복, 그리고 스트림 수에 따라 세제곱으로 커지는 매핑 생성 비용이다.

xHC는 MLP 출력 뒤에 다중 스케일 인과 컨볼루션을 넣어 각 갱신이 더 많은 정보를 조합하도록 만들고, 스트림 수를 16개까지 늘리되 '촘촘히 읽고 드물게 쓰는' 방식을 택했다. 모든 층이 16개 스트림을 전부 읽지만 실제 갱신에 참여하는 것은 4개뿐이다. 2개는 고정, 2개는 라우터가 고른다. 나머지는 그대로 다음 층으로 흘러간다. 읽기를 조밀하게 유지하는 이유는 방금 갱신된 정보가 다음 층에서 읽히지 않으면 층 간 전달이 끊기기 때문이다. 매핑 비용은 활성화된 스트림 수에만 비례하므로 전체 스트림을 늘려도 세제곱 비용을 피할 수 있다.

발표된 수치로는 18B MoE 모델의 평균 점수가 mHC 44.8에서 xHC 48.8로, 28B 모델은 50.5에서 53.6으로 올랐다. 같은 손실 기준으로 표준 Transformer 대비 학습 연산을 약 33%, mHC 대비 약 16% 절감한다는 결과도 제시됐다. 16갈래 스트림의 메모리 접근량은 여전히 4갈래 mHC보다 많지만, 연산자 융합 등을 거친 뒤 처리량 오버헤드는 약 1.3% 수준이라고 한다.

두 연구의 방향은 대학 연구실의 현실과 맞닿아 있다. 파라미터를 늘리고 전문가 수를 확장하고 데이터를 더 투입하는 정공법은 산업계 규모의 연산을 전제로 한다. 하드웨어와 드라이버, 연산자부터 모델 소프트웨어까지 생태계가 Transformer에 맞춰져 있어 구조를 크게 뜯어고치는 것도 쉽지 않다. 그래서 백본은 두고 용량을 붙이는 축을 찾는 쪽을 택했다.

실무에서 눈여겨볼 지점은 학습비와 추론비의 교환이다. 두 방법 모두 학습 연산을 줄이는 대신 파라미터 조회와 메모리 접근이라는 추론 측 비용을 남긴다. JTok의 7% 지연, xHC의 16갈래 스트림 접근량이 그 예다. 저자도 소규모 실험의 이득이 더 큰 규모에서 유지되는지, 연산 절감이 추론 지연·메모리 비용과 어떻게 상쇄되는지가 남은 과제라고 밝혔다. 규모의 법칙 역시 아직 경험식 수준이라는 단서도 붙였다. 두 방법은 서로 다른 축을 건드려 직교성이 있고, 결합이 다음 방향으로 언급됐다.