💬 NLP

T^2MLR: Transformer with Temporal Middle-Layer Recurrence

이 논문은 이전 토큰들의 캐싱된 중간 계층 표현을 현재 토큰의 더 이른 계층으로 융합하는 타겟팅된 접근 방식을 통해 잠재적 추론을 강화하는 트랜스포머 아키텍처인 T^2MLR을 소개하며, 이는 표준 베이스라인 및 전체 계층 순환 방식보다 우수한 성능을 보이면서도 기존의 사전 학습된 모델을 효율적으로 개조할 수 있게 한다.

원저자: Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora

게시일 2026-07-17

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 아주 똑똑한 로봇에게 까다로운 퍼즐을 푸는 법을 가르치려 한다고 상상해 보세요. 인공지능의 세계에서 오늘날 가장 인기 있는 로봇들은 '트랜스포머(Transformer)'라고 불리는 설계 구조를 기반으로 만들어집니다. 트랜스포머를 한 번에 한 단어씩 이야기를 읽어 내려가는 매우 빠르고 열정적인 독자로 생각해 보세요. 로봇이 새로운 단어를 읽을 때마다, 이전 단어에서 가졌던 깊고 복잡한 생각들을 거의 모두 잊어버리고, 그 풍부한 사고를 아주 작고 단순한 메모로 압축하여 새로 시작해야 합니다. 이것은 마치 수학 문제를 풀면서 다음 단계로 넘어가기 전에 매번 연습장을 지우고 최종 답안만 적어 넣는 것과 같습니다. 간단한 작업에는 효과적이지만, 로봇이 비밀 코드를 추적하거나 다단계 수수께끼를 푸는 것처럼 긴 논리의 사슬을 머릿속에 담아두어야 할 때는, 다음 단계로 넘어갈 때 '생각'을 계속 유지하지 못해 길을 잃곤 합니다.

과학자들은 로봇이 자신의 생각을 담은 '숨겨진 노트'를 가질 수 있게 함으로써 이 문제를 해결하려고 노력해 왔습니다. 하지만 대부분의 시도는 로봇이 전체 노트를 매번 다시 읽어야 하기 때문에 엄청나게 느려지거나, 로봇이 배우기 어려운 기묘하고 보이지 않는 언어로 생각을 적도록 강요하는 방식이었습니다. 큰 질문은 이것입니다. 로봇이 속도를 늦추거나 학습을 불가능하게 만들지 않으면서도, 한 단어에서 다음 단어로 넘어갈 때 그 깊고 추상적인 사고를 계속 유지할 수 있는 방법을 줄 수 있을까요? 이 논문은 T2MLR(Temporal Middle-Layer Recurrence를 가진 Transformers)이라는 영리한 새로운 기술을 도입하며, 그 답은 '예'라고 제안합니다. 다만, 우리가 로봇이 생각을 보관하는 '위치'를 바꾸어야 한다는 조건이 붙습니다.

연구진은 AI 모델이 사고 과정을 '기억'하는 새로운 방식을 제안합니다. 로봇에게 깊은 생각을 처리하는 과정의 맨 처음이나 맨 끝에서 아주 작은 메모로 적도록 강요하는 대신, 그들은 로봇의 뇌 중간 부분에서 다음 단어를 위한 뇌 중간 부분으로 직접 '비밀 메시지'를 전달하게 합니다. 이는 주자들이 바통을 오직 출발선이나 결승선에서만 주고받는 것이 아니라, 트랙의 중간 지점에서 서로에게 특별한 쪽지를 전달하는 릴레이 경주를 상상해 보세요. 이를 통해 AI의 '중간 레이어(middle layers)'에서 일어나는 복잡하고 추상적인 추론이, 로봇이 글을 읽는 동안 처음부터 다시 계산할 필요 없이 지속되고 진화할 수 있게 됩니다.

논문에 따르면 이 '중간 레이어' 방식은 놀라울 정도로 잘 작동합니다. 그들이 상태 추적이 필요한 작업(게임의 움직임 순서를 기억하는 것과 같은)이나 다단계 수학 문제를 테스트했을 때, 새로운 T2MLR 모델은 동일한 크기의 표준 모델들을 일관되게 앞질렀습니다. 흥미롭게도, 연구진은 로봇 전체가 모든 것을 기억하게 만들 필요는 없으며, 단지 모델의 중간 레이어 중 작은 부분(약 20%)만이 이러한 쪽지를 주고받게 하는 것이 로봇 전체가 자신의 생각을 반복(loop)하게 만드는 것보다 더 낫다는 것을 발견했습니다. 더욱 고무적인 것은 실제 세상에서도 이 기술을 사용하기 위해 로봇을 처음부터 새로 만들 필요가 없다는 점입니다. 그들은 이미 사전 학습된 로봇(17억 개의 파라미터를 가진 모델)을 가져와서, 이 새로운 중간 레이어 메모 시스템을 단순히 '이식'했습니다. 수학 문제에 대한 약간의 추가 학습만 거친 후, 이 업그레이드된 로봇은 수학 문제를 푸는 능력이 눈에 띄게 향려졌습니다. 한 테스트에서는 성공률이 35.8%에서 거의 40%로 뛰었고, 다른 테스트에서는 12.8%에서 18%로 상승했습니다.

저자들은 이것이 작동하는 이유가 트랜스포머의 '중간' 부분이 가장 추상적인 추론이 일어나는 곳이기 때문이라고 설명합니다. 즉, 그 특정 생각들이 지속되도록 함으로써 모델이 훨씬 더 나은 '잠재적 추론(latent reasoning, 모든 것을 글로 쓰지 않고도 생각하는 것)'을 할 수 있게 된다는 것입니다. 하지만 저자들은 이것이 트레이드오프(trade-off)를 동반한다는 점을 주의 깊게 언급합니다. 즉, 로봇이 실제로 질문에 답할 때는 속도가 크게 느려지지 않지만(단어당 약 8%의 추가 작업만 발생), 이 메모 쪽지들이 어떻게 작동해야 하는지 계산하기 위해 컴퓨터가 추가적인 수학 연산을 수행해야 하므로 '학습'하는 데는 시간이 더 오래 걸립니다. 이러한 학습 비용에도 불구하고, 결과는 AI의 효과적인 추론이 반드시 뇌의 모든 레이어를 반복할 필요는 없으며, 대신 목표가 명확한 중간 레이어 메모 시스템이 더 똑똑하고 지속적인 사고를 여는 열쇠가 될 수 있음을 시사합니다.

기술 요약: T2MLR: 시간적 중간 계층 재귀를 갖춘 트랜스포머 (Transformer with Temporal Middle-Layer Recurrence)

문제 정의

저자들은 표준 트랜스포머 아키텍처의 추론 능력에 관한 근본적인 한계를 지적한다. 트랜스포머는 인컨텍스트 검색(in-context retrieval)에는 뛰어나지만, 그들의 자기회귀적 디코딩 메커니즘은 정보 병목 현상을 생성한다. 각 디코딩 단계에서 풍부하고 고차원적인 잠재 표현(latent representations)은 다음 단계의 입력으로 전달되기 전, 희소한 토큰 공간(원-핫 벡터)으로 반복적으로 압축된다. 이 과정은 중간 단계의 추상적 추론 상태가 시간을 가로질러 지속되는 것을 방지하며, 모델이 어텐션이나 고도로 압축된 입력 임베딩을 통해 필요한 정보를 간접적으로 재구성하도록 강제한다.

기존의 잠재 추론(latent reasoning) 접근 방식은 추론 전체를 연속적인 공간(주로 마지막 레이어)에서 수행하거나, 토큰 임베딩의 선형 결합을 통해 불확실성을 전파함으로써 이를 완화하려고 시도한다. 그러나 이러한 방법들은 일반적으로 재귀적 신호를 입력 임베딩을 통해 다시 피드백함으로써, 추상적 추론이 일어나는 것으로 알려진 중간 계층을 우회한다. 대안적으로, "루프형(looped)" 트랜스포머 아키텍처는 모든 레이어를 반복하여 유효 깊이를 증가시키지만, 이는 루프 횟수에 따라 선형적으로 증가하는 추론 비용을 초래하여 배포하기에 비실용적이다.

방법론

본 논문은 **T2MLR (Transformers with Temporal Middle-Layer Recurrence)**을 소개한다. 이는 추론 시의 깊이를 증가시키지 않으면서도 추상적인 중간 표현이 디코딩 단계 전반에 걸쳐 지속되고 진화할 수 있도록 설계된 새로운 아키텍처이다.

아키텍처 설계

T2MLR은 이전 토큰의 더 깊은 레이어와 현재 토کن의 더 이른 레이어를 연결하는 게이트형 재귀 경로를 도입하여 표준 $L$ -레이어 디코더 전용 트랜스포머를 수정한다.

재귀 경계 (Recurrence Boundaries): 아키텍처는 "중간" 블록의 레이어들을 정의하는 두 개의 레이어 인덱스 $\ell_{start}$ 와 $\ell_{end}$ 에 의해 매개변수화된다.
재귀 캐시 ( $R_t$ ): 각 디코딩 단계 $t$ 를 위해 상수 크기의 캐시 벡터가 유지된다. 이는 이전 단계의 레이어 $\ell_{end}$ 이후의 표현을 저장한다.
융합 모듈 ( $\Phi$ ): 현재 토큰의 표현( $h^{(\ell_{start}-1)}_t$ )을 처리하기 전, 모델은 이전 토큰으로부터 온 재귀 캐시( $R_{t-1}$ )와 현재 토큰을 융합한다. 이 융합은 학습 가능한 게이팅 메커니즘에 의해 제어된다:
$\Phi(h, R) = h + \tanh(\gamma_{cur}) \sigma(f_{cur}([h, R])) \odot h + \tanh(\gamma_{rec}) \sigma(f_{rec}([h, R])) \odot W_{rec}R$
이 설계는 모델이 현재 토큰의 스트림과 시간적으로 축적된 잠재 상태 중 어느 것에 의존할지를 학습할 수 있게 한다.
캐시 업데이트: 레이어 $\ell_{end}$ 를 통과한 후, 재귀 캐시는 잔차 연결(residual connection)과 RMSNorm을 통해 새로운 표현과 이전 캐시를 사용하여 업데이트된다.

훈련 전략: 시간적 병렬성 (Temporal Parallelism)

잠재 추론 아키텍처의 주요 과제는 시간 단계 사이에 의존성이 존재할 때 표준 시퀀스 병렬 훈련(토큰 병렬 방식)이 깨진다는 것이다. 이를 해결하기 위해 저자들은 Jacobi 고정점 반복(Jacobi fixed-point iterations)을 사용하는 근사된 시간적 병렬 훈련 방식을 제안한다.

순차적인 토큰 처리 대신, 모델은 병렬로 $d_{forward}$ 번의 정밀화(refinement) 반복을 수행한다.
표준 순방향 패스에서 유도된 초기 캐시로 시작하여, 이 캐시를 현재 토큰의 표현과 반복적으로 융합하고, 중간 블록을 정밀화하며, 캐시를 업데이트한다.
그래디언트는 절단된 깊이( $d_{backward}$ )를 통해 역전파되며, 이는 절단된 시간 역전파(Truncated Backpropagation Through Time, TBPTT)와 유사하다.
이 방식은 모델이 밀집된 티처 포싱(teacher forcing) 및 시퀀스 병렬성을 사용하여 훈련될 수 있게 하며, 저자들은 이것이 연속적인 체인 오브 쏘트(chain-of-thought) 변형 모델들 중 독보적인 기능이라고 주장한다.

주요 기여

아키텍처 혁신: T2MLR은 입력 임베딩이나 최종 출력에서가 아니라, 트랜스포머의 중간 계층에 특화된 재귀 경로를 도입한다. 이는 표준 자기회귀 인터페이스를 유지하면서 잠재 상태의 지속을 가능하게 한다.
확장 가능한 훈련: 저자들은 Jacobi 근사 단계를 사용하여 이러한 재귀적 의존성을 사용하여 모델을 훈련할 수 있음을 입증하였으며, 이를 통해 확장성을 희생하지 않고 대규모 데이터셋에서 사전 훈련 및 미세 조정을 가능하게 했다.
리트로피팅(Retrofitting) 능력: T2MLR은 처음부터 새로 훈련할 필요가 없음을 보여준다. 재귀 경로는 기존의 사전 훈련된 모델(예: SmolLM2-1.7B)에 삽입되어 상당한 추론 개선을 가져오도록 미세 조정될 수 있다.
효율성: 루프형 트랜스포머와 달리, T2MLR은 상수 크기의 캐시와 추론 시 단일 순방향 패스를 사용하기 때문에 무시할 수 있는 수준의 추론 오버헤드(토큰당 최대 ~8%)만을 추가한다.

실험 결과

저자들은 합성 벤치마크, 자연어 사전 훈련, 그리고 다운스트림 추론 작업에 대해 T2MLR을 평가한다.

합성 추론 (S5-Retrieval): T2MLR은 상태 추적과 인컨텍스트 검색을 동시에 요구하는 작업에서 표준 트랜스포머와 LSTM 모두를 크게 능가한다. 트랜스포머는 깊이 제약으로 인해 긴 시퀀스 길이에서 실패하고, LSTM은 고정된 상태 압축으로 인해 검색에서 실패하는 반면, T2MLR은 높은 정확도를 유지한다.
사전 훈련 (FineWeb-Edu): 135M 파라미터 모델에서 T2MLR 변형 모델들은 제로샷 NLP 벤치마크(ARC, HellaSwag 등)에서 파라미터 수가 동일한 트랜스포머 베이스라인과 일치하거나 이를 능가한다. 특히, 중간 계층 재귀(예: 레이어 13–18과 같이 약 20%의 레이어를 루핑)가 전체 계층 재귀보다 성능이 우수한 경우가 많았다.
다운스트림 추론: GSM-Aug (수학), HotPotQA (멀티홉 QA), Variable Assignment와 같은 작업에서 T2MLR은 베이스라인 대비 일관된 이득을 보여준다. 중간 계층 재귀 변형들이 다시 한번 전체 계층 재귀보다 우수한 성능을 보였으며, 이는 타겟팅된 재귀가 글로벌 루핑보다 더 효과적임을 시사한다.
스케일링 (Scaling): 성능 향상은 361M 및 1B 파라미터로 규모를 키우고 사전 훈련을 50B 토큰까지 확장할 때도 지속되며 더욱 커진다.
리트로피팅: 사전 훈련된 SmolLM2-1.7B-Instruct 모델에 T2MLR을 적용하고 수학 데이터로 미세 조정한 결과, GSM8K 정확도가 35.8%에서 39.9%로, MATH500 정확도가 12.8%에서 18.0%로 향상되었으며, 이는 전면적인 사전 훈련 없이도 실질적인 적용 가능성을 입증한다.

의의 및 주장

본 논문은 트랜스포머의 효과적인 잠재 추론을 위해 모든 레이어를 루핑할 필요는 없으며(이는 추론 비용을 증가시킨다), 재귀를 입력 임베딩에 국한할 필요도 없다(이는 추상적 계산을 우회한다)고 주장한다. 대신, 저자들은 타겟팅된 중간 계층 재귀만으로도 상당한 추론 능력을 끌어올리는 데 충분하다고 주장한다.

T2MLR의 의의는 다음과 같다:

추론 깊이와 추론 비용의 분리: 루프형 아키텍처와 관련된 선형적인 추론 페널티 없이도 시간적으로 지속되는 잠재 경로를 도입한다.
도입 장벽 완화: T2MLR이 기존 사전 훈련된 모델에 리트로핏될 수 있음을 보여줌으로써, 저자들은 현재 모델에 전면적인 사전 훈련 비용의 극히 일부만으로 고급 추론 능력을 추가할 수 있다고 주장한다.
귀납적 편향(Inductive Biases)의 검증: 결과는 트랜스포머에서 재귀의 "스위트 스폿(sweet spot)"이 입력이나 출력 인터페이스가 아니라 추상적 계산이 일어나는 중간 계층임을 시사한다.

저자들은 Jacobi 근사 단계로 인해 T2MLR이 훈련 시간 오버헤드(재귀 블록 크기에 따라 벽시계 시간 기준 2~4배)를 초환한다는 점을 인정한다. 그들은 이를 하나의 트레이드오프(trade-off)로 규정한다: T2MLR은 훈련 컴퓨팅 자원을 희생하여 추론 작업에 대해 낮은 오버헤드와 높은 성능을 가진 추론 아키텍처를 얻는다.