Universal Transformer는 어떻게 동작하는가
— 층을 쌓지 말고, 같은 층을 T번

Transformer는 서로 다른 층 6개를 쌓습니다. Universal Transformer(UT)는 층 하나를 만들어 6번 돌립니다. 바뀐 것은 이 한 줄뿐인데, 그 결과로 모델의 성격이 달라집니다. 층마다 가중치가 같으니 UT는 "깊이 방향으로 도는 RNN"이 되고, 몇 번 돌릴지를 학습이 끝난 뒤에도, 위치마다 따로 정할 수 있게 됩니다. 그리고 가장 실용적인 결과는 이것입니다 — 파라미터 수와 연산량이 서로 떨어집니다. 층을 공유하면 파라미터는 1/6로 줄지만 연산은 그대로이고, 거꾸로 층을 넓혀 파라미터를 base 수준으로 되돌리면 연산은 Transformer big과 똑같아집니다. 이 자료는 논문을 앞에서부터 따라가며 이 분리가 무엇을 사고 무엇을 내주는지를 브라우저에서 직접 계산해 봅니다.

how-ai-works 시리즈 · 원전: Dehghani, Gouws, Vinyals, Uszkoreit & Kaiser, Universal Transformers (ICLR 2019, arXiv:1807.03819)

§1 무엇이 바뀌었나 — 쌓기 대신 돌리기

2017년의 Transformer는 번역에서 순환 신경망을 이겼습니다. 그런데 이 논문은 첫 문단에서 불편한 사실을 지적합니다 — Transformer는 문자열 복사처럼 RNN이 쉽게 하는 일을, 학습 때보다 긴 입력에서 해내지 못한다는 것입니다(논문 초록). 저자들이 지목한 원인은 "서로 다른 층을 고정된 개수만큼 쌓는 구조"가 RNN이 가진 반복적(iterative)·재귀적 변환에 대한 귀납 편향을 버렸다는 점입니다(1절).

해법은 단순합니다. 층 N개를 쌓는 대신, 같은 가중치의 층 하나를 T번 적용합니다. 각 단계에서 모든 위치가 동시에(병렬로) 셀프 어텐션으로 서로의 정보를 주고받고, 이어서 위치별 전이 함수(transition function)를 통과합니다. 논문 그림 1의 캡션이 강조하는 한 줄은 이것입니다 — "Parameters are tied across positions and time steps."

네 모델을 번갈아 눌러 보면 차이가 선명해집니다.

• RNN은 위치를 따라 돕니다. 가중치는 하나지만 m개 위치를 순서대로 지나야 하므로 순차 틱이 m개입니다.
• Transformer는 모든 위치를 한 번에 계산하므로 순차 틱은 층 수 N뿐입니다. 대신 층마다 색이 다릅니다 — 가중치가 N벌입니다.
• UT는 Transformer처럼 한 틱에 모든 위치를 계산하지만 모든 칸이 같은 색입니다. 위치가 아니라 깊이를 따라 도는 RNN입니다. 그래서 T는 학습 때 정해 둔 상수가 아니라 돌리고 싶은 만큼 돌릴 수 있는 반복 횟수가 됩니다.
• UT + 동적 정지에서는 위치마다 멈추는 단계가 다릅니다. 멈춘 위치의 상태는 다음 단계로 그대로 복사되고(빗금 칸), 모든 위치가 멈추거나 최대 단계에 도달하면 끝납니다(2.2절).

§2 수식 두 줄 — 깊이 방향으로 도는 RNN

인코더는 입력 m개 기호의 임베딩으로 만든 행렬 H⁰ ∈ ℝ^m×d에서 시작해, 단계 t = 1, …, T마다 다음 두 줄을 적용합니다(논문 식 4·5).

Transformer 블록과 나란히 놓으면 차이는 세 군데입니다.

디코더도 같은 구조를 반복합니다. 셀프 어텐션 뒤에 인코더의 최종 표현 H^T를 읽는 교차 어텐션이 붙고, 셀프 어텐션에는 인과 마스크가 걸립니다. 학습은 teacher forcing, 생성은 한 기호씩 자기회귀로 — 여기까지는 Transformer와 같습니다(2.1절). 즉 UT의 "순환"은 생성 방식과는 관계가 없습니다. 토큰 방향으로는 여전히 병렬이고, 순환은 오직 깊이 방향에만 있습니다.

§3 파라미터와 연산이 떨어진다 — "같은 파라미터"의 진짜 의미

층을 공유하면 무엇이 줄어들까요? 파라미터는 줄지만 연산은 줄지 않습니다. 한 번 계산한 층을 재사용하는 것이 아니라, 같은 가중치로 다른 입력을 T번 다시 계산하기 때문입니다. Transformer에서 "층 수"는 파라미터와 연산을 함께 정했는데, UT에서는 두 손잡이가 분리됩니다 — 파라미터는 d와 d_ff가, 연산은 d·d_ff와 T가 정합니다.

논문 표 7에는 이런 각주가 붙어 있습니다 — "All base results have the same number of parameters." 층 하나로 6층짜리와 파라미터를 맞추려면 층을 넓혀야 합니다. 공개된 구현(tensor2tensor)의 universal_transformer_base는 실제로 hidden_size = 1024, filter_size = 4096, num_heads = 16으로 설정되어 있고, 주석에 이유가 적혀 있습니다 — "To have a similar capacity to the transformer_base with 6 layers, we need to increase the size of the UT's layer." 그 결과가 무엇인지 계산해 봅시다.

네 개의 프리셋을 차례로 눌러 보면 이런 표가 나옵니다.

세 가지를 읽을 수 있습니다.

1. UT small은 base와 연산이 같고 파라미터는 42%입니다. 결과는 1.2 BLEU 낮습니다 — 공유는 공짜가 아닙니다. 같은 연산이라면 층마다 다른 가중치를 쓰는 편이 표현력이 큽니다.
2. UT base는 파라미터가 base와 거의 같지만(67.28M, +6.7%) 연산은 정확히 4배입니다. d를 2배로 늘리면 행렬 하나의 크기가 4배가 되는데, 층 수가 6에서 1로 줄어 파라미터는 4/6배, 연산은 6단계를 그대로 도니 4배입니다. 참고로 d_ff = 4d를 유지하며 파라미터를 base(63.08M)에 정확히 맞추는 d를 풀면 979가 나옵니다 — 1024는 그 근처의 2의 거듭제곱입니다.
3. UT base와 Transformer big은 토큰당 연산이 한 자리까지 같습니다. 블록 모양(1024 / 4096)과 깊이(6)가 같으니 당연합니다. 차이는 파라미터가 1/3.2라는 것뿐이고, BLEU는 UT base가 0.5 높습니다.

§4 좌표 임베딩 — 공유된 층에게 "지금 몇 번째인가" 알려 주기

가중치를 공유하면 새로운 문제가 생깁니다. 층 하나가 1단계에서도 6단계에서도 같은 함수이므로, 입력만 보고는 지금이 몇 번째 반복인지 알 방법이 없습니다. Transformer의 층들은 각자 다른 가중치로 "나는 아래쪽 층이니 국소적인 것을 본다" 같은 역할 분담을 할 수 있었지만, UT의 층은 그럴 수 없습니다.

논문의 해법은 위치 인코딩과 똑같은 사인파를 단계 번호 t에도 한 번 더 계산해 위치 쪽에 더하는 것입니다(식 6·7).

그리고 이 P^t를 입력에 한 번이 아니라 매 단계 H^t−1에 다시 더합니다(식 5). 공유된 층이 받는 입력이 단계마다 조금씩 달라지는 것이 곧 "지금 몇 번째인가"라는 신호입니다.

히트맵에서 단계 부분만을 누르면 모든 행이 같은 줄무늬입니다 — 단계 신호는 위치와 무관한 하나의 벡터라는 뜻입니다. 오른쪽 곡선은 조금 더 흥미롭습니다. sin·cos 짝마다 전개하면

그래서 곡선은 t를 옮기면 모양 그대로 평행이동하고, 봉우리는 언제나 i = t에서 √(2d) = 32(d = 512)입니다. 위치 인코딩 하나의 길이가 항상 √(d/2) = 16이었던 것(How Transformer Works §3)과 비교하면, 두 사인파를 더하면서 길이가 16과 32 사이에서 흔들리게 된 것입니다. i − t가 커져도 저주파 차원들은 cos ≈ 1로 남아 있어서 √d ≈ 22.6까지 금방 내려가지는 않습니다(i − t = 997에서도 24.7). 이것이 학습에 문제가 되지는 않습니다 — 임베딩 쪽 크기와 LayerNorm이 흡수하기 때문입니다. 다만 "위치와 단계를 그냥 더했다"는 선택이 벡터 공간에서 어떤 모양이 되는지는 이 식이 보여 줍니다.

§5 고정 깊이의 한계 — 길이가 늘면 단계도 늘어야 한다

논문 부록 B의 논증은 짧습니다. Transformer의 전체 연산량은 입력 길이에 따라 늘지만, 순차적인 연산 단계 수는 층 수 N으로 고정되어 입력 길이와 무관합니다. 유한 정밀도를 가정하면 이것만으로 Transformer는 계산 보편적(computationally universal)일 수 없습니다. 반면 UT는 단계 수를 입력 길이의 함수로 정할 수 있고, 이는 가중치를 단계 간에 공유하기 때문에 학습이 끝난 뒤에도 가능합니다.

"for any given choice of depth T, one can construct an input sequence of length N > T that cannot be processed correctly by a standard Transformer." — 부록 B

이것을 눈으로 보기 위해 누적 패리티(prefix parity) 문제를 골랐습니다 — 비트열 x₁…x_n이 주어지면 각 위치 i에 x₁ ⊕ … ⊕ x_i를 출력합니다. 아래 데모의 "층"은 학습된 것이 아니라 손으로 정한 규칙입니다 — 단계 t에서 위치 i는 어텐션으로 정확히 한 위치 i − 2^t−1을 읽고 자기 값에 XOR합니다(Hillis–Steele 누적 스캔). 읽을 거리 2^t−1이 단계마다 다르다는 점에 주목하세요 — 가중치가 같은 층이 단계마다 다른 일을 하려면 바로 §4의 단계 신호가 필요합니다.

단계 t가 끝나면 위치 i는 [i − 2^t + 1, i] 구간을 본 상태가 되므로, i ≤ 2^t인 위치만 정답이 확정됩니다. 층 수를 N으로 고정하면 2^N보다 긴 입력의 뒷부분은 절대 맞힐 수 없습니다 — N = 6이면 64까지. 논문의 알고리즘 과제가 정확히 이 상황입니다. 길이 40으로 학습하고 길이 400으로 평가합니다(3.4절). 이 규칙으로는 40에 ⌈log₂ 40⌉ = 6단계, 400에 9단계가 필요합니다. Transformer는 층 수를 바꾸면 새 층을 학습해야 하지만, UT는 같은 층을 세 번 더 돌리면 됩니다.

또 하나 볼 것이 있습니다. UT 모드에서 앞쪽 위치일수록 일찍 초록이 됩니다 — 위치 i에 필요한 단계는 ⌈log₂ i⌉로, 위치마다 다릅니다. 이미 답이 확정된 위치를 계속 계산하는 것은 낭비입니다. 위치마다 멈출 때를 스스로 정하게 하자 — 이것이 다음 절의 동적 정지입니다.

§6 동적 정지(ACT) — 위치마다 다른 깊이

문장 속 기호는 모호함이 제각각입니다. 논문은 Graves(2016)의 ACT(Adaptive Computation Time)를 가져와 위치마다 붙입니다(2.2절). 각 단계에서 위치별로 정지 확률 p = σ(w·h + b)를 내고, 이를 누적해 임계값 1 − ε을 넘으면 그 위치는 멈춥니다. 멈춘 위치의 상태는 그대로 다음 단계로 복사되지만, 다른 위치의 어텐션 대상으로는 계속 남습니다.

논문 부록 C의 Listing 2가 한 단계의 계산을 그대로 보여 줍니다. 아래 실행기는 그 코드를 한 줄씩 옮긴 것입니다.

still_running = (halting_probability < 1.0)
new_halted    = (halting_probability + p*still_running > threshold) * still_running
still_running = (halting_probability + p*still_running <= threshold) * still_running
halting_probability += p * still_running
remainders          += new_halted * (1 - halting_probability)
halting_probability += new_halted * remainders
n_updates           += still_running + new_halted
update_weights = p*still_running + new_halted*remainders
new_state = transformed_state*update_weights + previous_state*(1 - update_weights)

실행해 보면 확인할 수 있는 것들입니다.

• 나머지(remainder) R · 멈추는 단계에서는 p가 아니라 1 − (그때까지의 누적)이 가중치가 됩니다. 그래서 멈춘 위치의 정지 확률 합은 항상 정확히 1이 됩니다.
• 초기화 직후 · tensor2tensor는 정지 유닛의 편향을 act_halting_bias_init = 1.0으로 둡니다. 가중치 항을 무시하면 p = σ(1) = 0.731이고, 0.731 + 0.731 > 0.99이므로 모든 위치가 2단계에서 멈춥니다(R = 0.269). 학습은 "대부분 2단계"에서 출발해, 필요한 위치만 더 오래 돌도록 바뀌어 갑니다. 최대 단계 기본값은 2 × num_hidden_layers = 12입니다.
• 한 번에 멈추는 위치 · 예시의 the처럼 첫 단계에 p > 0.99를 내면 누적 0에서 바로 멈추고 R = 1을 받습니다 — 단계 1의 상태가 그대로 출력이 됩니다.
• 최대 단계에 걸린 위치 · "전부 망설임"(p = 0.05)을 고르면 12단계 동안 누적이 0.60에 그쳐, 0.99에 닿기 전에 최대 단계에서 루프가 끝납니다. 이 위치들은 R을 받지 못해 가중치 합이 1이 되지 않습니다.
• 숨은 차이 — 섞는 방식 · Graves의 ACT는 최종 출력을 Σ_t w_t·s_t로 섞습니다. 그런데 Listing 2는 new_state = s·w + prev·(1 − w)로 매 단계 보간하고, 누적 상태가 0에서 시작합니다. 그래서 k번째 상태의 실효 가중치는 w_k·Π_j>k(1 − w_j)가 되고, 합은 1 − Π(1 − w_j)로 1보다 작아질 수 있습니다 — 모자란 만큼은 초기값 0의 몫입니다. 초기화 직후 예에서는 0.731 × 0.731 + 0.269 = 0.803, 즉 출력의 약 20%가 0 벡터 쪽으로 끌려갑니다(뒤에 오는 LayerNorm이 크기를 되돌리긴 합니다). tensor2tensor는 Graves 방식도 act_type = "accumulated"로 따로 제공합니다.

마지막으로 비용입니다. 정지 유닛은 "오래 생각할수록 좋다"는 방향으로만 학습되면 늘 최대 단계까지 돌 것이므로, Graves는 위치별 숙고 비용(ponder cost) ρ = N + R을 손실에 더합니다. tensor2tensor는 이것을 act_loss_weight = 0.01 × 평균(N + R)으로 인코더와 디코더에 각각 더합니다. 위 실행기의 ρ 값이 그 항입니다.

§7 ACT는 무엇을 배웠나 — bAbI와 LAMBADA

ACT가 정말로 "어려운 곳에서 오래 생각하는지"는 두 실험에서 드러납니다.

bAbI — 근거 사실이 늘수록 더 오래 생각한다

bAbI는 짧은 이야기와 질문을 주고 답하게 하는 20개 합성 과제입니다. 답하는 데 필요한 근거 사실(supporting fact) 개수가 과제마다 1·2·3개로 다릅니다. 논문은 테스트 데이터의 모든 위치에서 평균 숙고 시간(ponder time)을 쟀습니다(3.1절).

평균만 늘어난 것이 아닙니다. 논문은 근거 사실이 3개인 과제에서 많은 위치가 1~2단계에서 이미 멈추고 몇몇 위치만 오래 변환된다고 보고합니다 — 표준편차가 0.8 → 2.2로 커진 것이 그 흔적입니다. 이 과제들은 이야기가 훨씬 길고 관련 없는 사실이 많은데, 모델이 그런 문장에 계산을 쓰지 않는 법을 배운 것으로 보입니다. 부록 F의 어텐션 시각화도 같은 이야기를 합니다 — 첫 단계의 어텐션은 거의 균일하다가 단계가 진행될수록 정답에 필요한 사실 쪽으로 날카로워집니다. 단계가 곧 다단계 추론의 "홉(hop)" 역할을 하는 셈입니다.

LAMBADA — "그냥 더 많이 돌아서"가 아니다

LAMBADA는 앞의 4–5문장을 다 읽어야만 맞힐 수 있는 마지막 단어 예측 과제입니다. ACT를 붙인 UT는 테스트에서 평균 8.2 ± 2.1단계를 돌았는데, 고정 6단계 모델보다 좋았습니다. 그러면 자연스러운 의문 — 단지 단계를 더 써서 이긴 것 아닌가? 저자들은 고정 8단계·9단계 모델을 따로 학습해 확인했습니다(표 3).

세 가지가 눈에 띕니다.

• LSTM과 Transformer는 dev/test 퍼플렉시티가 수천이고 정확도가 0.0입니다. 반면 control 세트(걸러내지 않은 일반 문장)에서는 130~140대로 UT와 큰 차이가 없습니다 — 평범한 언어 모델링은 되는데 넓은 문맥이 필요한 예측만 못 하는 것입니다.
• 고정 단계를 6 → 8로 늘리면 test 퍼플렉시티가 319 → 202로 좋아지지만, 9로 늘리면 오히려 239로 나빠집니다. 그리고 어떤 고정 단계도 ACT의 142에 닿지 못합니다.
• 저자들의 해석은 ACT가 정규화(regularizer)로 작용한다는 것입니다 — 어떤 기호에는 단계를 적게 쓰도록 유도하고 다른 기호에는 더 많이 허용하면서, 모든 위치를 같은 깊이로 돌리는 것보다 일반화가 좋아진다는 것입니다.

§8 실험 결과 — 알고리즘, 언어, 번역

과제별로 읽을 거리를 정리하면 이렇습니다.

• 주어–동사 일치 · "The keys to the cabinet ___"에서 are/is를 맞히는 과제입니다. 주어와 동사 사이에 수가 반대인 명사(유인자, attractor)가 끼어들수록 어려워집니다. 전체 평균은 모두 0.99대로 비슷하지만, 유인자가 늘수록 UT + ACT와 이전 최고 모델의 격차가 벌어집니다(유인자 1개에서는 −0.008로 뒤지지만 3·4·5개에서 +0.009·+0.016·+0.027). 계층 구조를 여러 단계에 걸쳐 풀어야 하는 경우일수록 반복이 도움이 된다는 신호입니다. Transformer는 유인자 0개에서도 0.973으로 가장 낮습니다.
• 알고리즘(복사·뒤집기·덧셈) · 길이 40으로 학습하고 400으로 평가합니다. UT는 복사·뒤집기에서 문자 정확도 0.91·0.96으로 Transformer(0.53·0.13)를 크게 앞서지만, 덧셈의 문장 전체 정답률은 0.02입니다. 커리큘럼으로 학습한 Neural GPU만 1.0입니다 — 표현할 수 있다는 것과 학습으로 찾는 것은 다른 문제입니다.
• 프로그램 실행(LTE) · 암기 과제(copy·double·reverse)에서 UT는 전부 1.0입니다. 프로그램 평가 과제에서도 LSTM·Transformer보다 모두 높습니다.
• bAbI · 표준 Transformer는 하이퍼파라미터와 크기를 바꿔 봐도 항상 과적합했다고 각주 2에 적혀 있습니다(10K 단일 과제 평균 오차 15.2%). UT는 0.23%, ACT를 붙이면 0.21%로 20개 과제 모두 통과합니다. 데이터가 1K로 줄어든 설정에서 격차가 특히 큽니다 — 가중치 공유가 데이터 효율에 도움이 된다는 해석과 맞닿아 있습니다.
• 번역 · §3에서 본 대로입니다. 연산이 같은 UT small은 base보다 낮고, 파라미터가 같은 UT base는 base보다 0.9 높으며, 그 UT base의 연산은 big과 같습니다.

§9 그 뒤의 이야기 — 공유의 계보

UT가 제기한 두 가지 아이디어 — 층 사이 가중치 공유와 입력마다 다른 깊이 — 는 이후 따로따로 이어졌습니다.

• ALBERT(Lan et al., 2019)는 BERT의 모든 층 파라미터를 공유해 BERT-base 약 108M을 ALBERT-base 약 12M으로 줄였습니다. §3의 교훈이 그대로 적용됩니다 — 파라미터는 줄었지만 추론 연산은 줄지 않습니다. 메모리와 통신은 아끼지만 속도는 그대로입니다.
• 깊이의 조건부 계산 · 위치마다 층을 건너뛰게 하거나(Mixture-of-Depths 계열), 예측이 확실하면 일찍 빠져나가는(early exit) 방법들은 ACT의 "위치별 정지"를 다른 방식으로 구현한 것입니다.
• 반복(looped) Transformer · 같은 블록을 반복 적용해 학습 때보다 어려운 문제에 반복 횟수를 늘려 대응하려는 연구가 이어지고 있습니다. §5의 "길이가 늘면 단계도 늘려라"가 그 출발점입니다.

반면 대형 언어 모델의 주류는 여전히 층마다 다른 가중치를 쌓는 쪽입니다. §3의 표가 그 이유를 보여 줍니다 — 연산 예산이 제약일 때, 같은 연산으로 더 많은 파라미터를 쓰는 쪽이 유리한 경우가 많기 때문입니다. UT는 "파라미터가 제약이고 연산은 여유가 있을 때", 그리고 "입력마다 필요한 계산량이 크게 다를 때" 빛나는 설계입니다.

§10 정리 — 한 장 요약

§11 이해도 퀴즈