연산 학습을 위한 학습 친화적 토큰 순서 탐색

초록

본 논문은 Transformer 디코더가 출력 토큰을 어떤 순서로 생성하느냐가 학습 효율에 큰 영향을 미친다는 점에 주목한다. 저자는 다양한 순열을 섞어 학습한 뒤, 초기 학습 단계에서 손실이 가장 빠르게 감소하는 순열을 “학습 친화적”으로 정의하고, 이를 찾기 위한 두 단계 계층적 탐색 방법을 제안한다. 7개의 순서‑민감 연산 과제에 실험을 적용해 수십억 개의 후보 중 최적 순서를 발견했으며, 특히 곱셈 과제에서 기존 연구와 동일한 역순(least‑significant‑digit‑first) 순서를 자동으로 재발견했다.

상세 분석

이 논문은 Transformer 기반 자동 회귀 모델이 “체인 오브 사고(chain‑of‑thought, CoT)”를 통해 복합 연산을 수행할 때, 중간 단계 토큰의 배치 순서가 학습 난이도를 결정한다는 가설을 실험적으로 검증한다. 기존 연구는 CoT의 내용(어떤 중간 단계가 포함되는가)만을 설계했으며, 순서 자체는 직관에 의존하거나 고정된 형태(예: 왼쪽‑부터‑오른쪽)로 사용했다. 그러나 연산에서 캐리(carry)와 같은 비가역적 연산 흐름은 토큰이 생성되는 순서와 직접 연결된다. 예를 들어, 두 정수의 곱을 least‑significant‑digit‑first 로 출력하면 각 단계가 이전 단계의 캐리를 자연스럽게 활용할 수 있어 학습이 쉬워진다.

논문은 이 현상을 일반화하기 위해 학습 친화적 순열을 정의한다. 구체적으로, 동일한 데이터셋에 대해 서로 다른 순열을 적용한 여러 버전의 목표 시퀀스를 혼합 데이터셋으로 만든 뒤, 제한된 에폭(수천 배치) 동안 학습한다. 학습이 진행되는 동안 손실 감소 속도가 가장 빠른 순열을 “easy sample”로 간주하고, 이를 최종 후보로 선택한다. 이 접근법은 easy‑to‑hard 학습 다이내믹(Arpit et al., 2017)과 커리큘럼 학습의 원리를 출력 순서에 적용한 형태라 할 수 있다.

순열 탐색 자체는 factorial( L! ) 규모로 급증하므로, 저자는 두 단계 계층적 탐색을 제안한다.

1. 전역 단계: 전체 시퀀스를 일정 길이 블록으로 나누고, 블록 간 순서를 탐색한다. 블록은 사전 정의된 크기(예: 4~5 토큰)이며, 블록 순열은 전체 후보 공간을 크게 축소한다.
2. 지역 단계: 각 블록 내부에서 토큰 순서를 미세 조정한다. 전역 단계에서 선정된 블록 순서를 고정하고, 블록 내부 순열을 다시 혼합‑학습‑평가 과정을 통해 최적화한다.

이 두 단계는 탐색 효율을 크게 향상시키며, 실험에서는 L=13(≈10⁹ 후보)까지 무작위 초기화로, L=40(≈10⁴⁷ 후보)까지 구조화된 초기화로 성공적으로 최적 순서를 찾아냈다.

실험 설계는 7개의 순서‑민감 연산 과제(덧셈, 뺄셈, 곱셈, 나눗셈, ReLU 누적, 비가역 함수 연쇄 등)와 6개의 순서‑민감 인스턴스(입력·출력 길이 10~40)로 구성된다. 각 과제마다 “정방향”(most‑significant‑digit‑first)과 “역방향”(least‑significant‑digit‑first) 등 여러 후보 순열을 생성하고, 제안된 탐색으로 최적 순열을 도출한다. 결과는 다음과 같다.

• 곱셈 과제에서 역방향 순열이 최적임을 자동 재발견, 이는 기존 연구(Shen et al., 2023)와 일치.
• ReLU 누적 연산에서는 전방 순열이 우수, 이는 함수의 비가역성(음수 → 0) 때문에 앞쪽 토큰이 더 많은 정보를 제공함을 시사.
• 전체적으로 학습 친화적 순열을 사용했을 때 최종 성공률이 10 %→100 %까지 크게 향상됨.

또한 소프트 퍼뮤테이션(연속적인 확률 행렬) 기반 직접 최적화가 미래 토큰 누수를 야기해 손실이 인위적으로 급감하고, 실제 일반화 성능이 저하되는 문제를 확인한다. 따라서 제안된 “혼합‑학습‑평가” 방식이 현재 가장 실용적인 해결책으로 자리한다.

이 논문의 주요 기여는 (1) 출력 순서 최적화라는 새로운 연구 주제를 정의하고, (2) 학습 동적 기반 탐색이라는 효율적 방법론을 제시했으며, (3) 실제 연산 과제에 적용해 기존 휴리스틱을 자동 재발견함으로써, 향후 복합 연산, 프로그래밍 언어 모델, 심볼릭 추론 등에서 출력 순서 설계에 대한 체계적 접근을 가능하게 한다는 점이다.