트랜스포머가 곱셈을 학습하지 못하는 이유는 무엇인가? 장거리 의존성 함정을 밝히는 리버스 엔지니어링

초록

언어 모델은 점점 더 강력해지고 있지만, 여전히 다중 자릿수 곱셈이라는 겉보기엔 단순한 작업에서 실패한다. 본 연구에서는 암묵적 사고 연쇄(implicit chain-of-thought)를 통해 곱셈을 성공적으로 학습하는 모델을 역공학적으로 분석하여 세 가지 발견을 보고한다: (1) 장거리 구조의 증거: 로짓 속성(logit attribution)과 선형 탐사(linear probe)는 모델이 다중 자릿수 곱셈에 필요한 장거리 의존성을 인코딩하고 있음을 나타낸다. (2) 메커니즘: 모델은 주의 메커니즘(attention)을 사용하여 방향성 비순환 그래프(directed acyclic graph)를 구성하고, 이를 통해 쌍별 부분 곱(pairwise partial product)을 "캐시"하고 "검색"함으로써 장거리 의존성을 인코딩한다. (3) 기하학: 모델은 주의 헤드(attention head)에서 부분 곱을 구현하기 위해 숫자 쌍 간의 민코프스키 합(Minkowski sum)을 형성하며, 숫자는 푸리에 기저(Fourier basis)를 사용하여 표현된다. 이 두 가지는 직관적이고 효율적인 표현 방식으로, 표준 미세 조정(fine-tuning) 모델에는 없는 특징이다. 이러한 통찰을 바탕으로, 표준 미세 조정의 학습 동역학을 재검토한 결과, 모델이 필요한 장거리 의존성이 없는 지역 최적점(local optimum)에 수렴함을 발견했다. 이를 검증하기 위해, 선형 회귀 탐사를 통해 "누적 합(running sum)"을 예측하는 보조 손실(auxiliary loss)을 도입했으며, 이는 모델이 다중 자릿수 곱셈을 성공적으로 학습할 수 있도록 하는 귀납적 편향(inductive bias)을 제공한다. 요약하면, 암묵적 사고 연쇄 모델의 메커니즘을 역공학적으로 분석함으로써 트랜스포머(Transformer)에서 장거리 의존성을 학습하는 데 있어 함정을 발견하고, 올바른 귀납적 편향이 이 문제를 해결할 수 있는 사례를 제시한다.

English

Language models are increasingly capable, yet still fail at a seemingly simple task of multi-digit multiplication. In this work, we study why, by reverse-engineering a model that successfully learns multiplication via implicit chain-of-thought, and report three findings: (1) Evidence of long-range structure: Logit attributions and linear probes indicate that the model encodes the necessary long-range dependencies for multi-digit multiplication. (2) Mechanism: the model encodes long-range dependencies using attention to construct a directed acyclic graph to ``cache'' and ``retrieve'' pairwise partial products. (3) Geometry: the model implements partial products in attention heads by forming Minkowski sums between pairs of digits, and digits are represented using a Fourier basis, both of which are intuitive and efficient representations that the standard fine-tuning model lacks. With these insights, we revisit the learning dynamics of standard fine-tuning and find that the model converges to a local optimum that lacks the required long-range dependencies. We further validate this understanding by introducing an auxiliary loss that predicts the ``running sum'' via a linear regression probe, which provides an inductive bias that enables the model to successfully learn multi-digit multiplication. In summary, by reverse-engineering the mechanisms of an implicit chain-of-thought model we uncover a pitfall for learning long-range dependencies in Transformers and provide an example of how the correct inductive bias can address this issue.

트랜스포머가 곱셈을 학습하지 못하는 이유는 무엇인가? 장거리 의존성 함정을 밝히는 리버스 엔지니어링

Why Can't Transformers Learn Multiplication? Reverse-Engineering Reveals Long-Range Dependency Pitfalls

초록

Support