Looped Transformers are Better at Learning Learning Algorithms: 반복 transformer로 in-context 학습 알고리즘 흉내 내기

개요

transformer는 프롬프트에 (x₁, y₁, …, x_k, y_k, x_test)를 넣어 주면 가중치를 바꾸지 않고 그 자리에서 y_test를 맞히는 **in-context learning(문맥 내 학습: 예시만 보고 새 함수를 배우는 능력)**을 보인다. Garg et al.(2022)은 처음부터 학습한 12층 transformer가 선형회귀를 최소제곱법만큼 잘 푼다는 것을 보였다. 그런데 전통 기계학습은 이런 문제를 경사하강법처럼 같은 갱신을 반복하는 알고리즘으로 푼다. 반면 일반 transformer는 층마다 다른 가중치를 가진 일회성 계산이다. 저자들의 출발점은 이것이다. 최소제곱 해 (XᵀX)⁻¹Xᵀy는 역행렬이 들어가서 transformer가 배우기 어렵지만, 경사하강 한 걸음 Xᵀ(Xw−y)는 행렬곱뿐이다. 다만 여러 번 반복해야 한다. 그렇다면 한 층짜리 블록을 반복하는 looped transformer가 반복 알고리즘을 더 적은 파라미터로 흉내 낼 수 있지 않을까. 핵심은 학습법에 있다. 매 반복마다 원래 프롬프트를 더해 주는 **input injection(입력 주입: Y_{t+1}=M(Y_t+P))**이 없으면 학습한 반복 횟수를 넘기자마자 오차가 폭발했다. 또 b번 반복한 결과 중 마지막 T번에만 손실을 거는 truncated loss로 메모리를 아끼면서, b를 충분히 크게 잡아 학습 범위 밖에서도 출력이 **고정점(fixed point)**에 머물게 했다. 결과적으로 0.79M 파라미터의 1층 looped transformer가 9.48M의 12층 transformer와 선형회귀·결정트리·2층 ReLU 신경망에서 비슷한 오차를 냈다. sparse linear regression에서는 더 적은 예시로 오차가 줄어 transformer와 Lasso보다 빨리 내려갔다(예시 10개 미만에서 Lasso보다 오차가 작음). 대가도 있다. 입력을 3배로 키운 분포 밖 시험에서는 looped transformer가 일반 transformer보다 나빴고, 저자들은 모델이 진짜 알고리즘이 아니라 학습 분포 안에서만 작동하는 고정점 반복을 배웠다고 인정한다.

Implication: 반복 구조는 “더 단순한 해”를 선호하는 귀납 편향을 준다. 이 편향은 sparse 문제나 적은 예시에서는 이득이고, 입력 스케일이 바뀌는 분포 이동에서는 손해다. 이후 looped transformer 연구가 다루는 입력 주입, 반복 수 외삽, 고정점 안정성 문제가 실험적으로 처음 정리된 자리다.

📎 https://arxiv.org/abs/2311.12424

배경과 문제

왜 어렵고 왜 중요한가. LLM의 in-context learning 메커니즘은 아직 불분명하다(Section 1, author-stated). 이를 이해하려는 흐름은 transformer를 작은 함수 클래스 학습 문제로 처음부터 학습시키는 실험(Garg et al. 2022), 이를 학습 알고리즘의 암묵적 구현으로 해석한 연구(Akyürek et al. 2022), 선형 self-attention만 있는 transformer가 경사하강 한 걸음을 배운다는 실험(von Oswald et al. 2022)으로 이어진다. 저자들은 von Oswald et al.의 단순화된 구조와 실제 decoder transformer 사이에 큰 간극이 있고, 표준 decoder transformer를 작은 구조 변경만으로 학습해 학습 알고리즘을 재현하는 문제가 열려 있다고 본다(Section 1).

반복 구조가 필요한 이유 (Section 1, author-stated). 선형회귀의 닫힌 해는 역행렬이 필요하고, 역행렬은 transformer가 배우기 더 어렵다(더 많은 층·헤드가 필요, Charton 2021; Giannou et al. 2023). 경사하강은 행렬곱만 쓰지만 반복이 필요하다. looped transformer는 문제를 단순한 부분 과제로 쪼개 반복하므로 파라미터를 크게 줄일 수 있다고 본다. 연구 질문은 “looped transformer가 일반 transformer보다 반복 학습 알고리즘을 더 효율적으로 흉내 낼 수 있는가”다.

학술적 계보 (원문 Section 2, author-stated).

• 가중치 공유 모델: RNN·LSTM, Transformer-XL 등 순환형 transformer, ALBERT·Perceiver·Universal Transformer(층 사이 가중치 공유, 제한된 층 수 또는 halting 기준). 극단에는 함수를 무한히 반복하는 implicit model(Bai et al. 2019)이 있다.
• Deep implicit model: 고정점을 black-box solver로 찾는 계열(Bai et al. 2018–2022)과 Jacobian 정규화(Bai et al. 2021). 저자들은 이 접근이 하이퍼파라미터 튜닝을 많이 요구하고 여전히 불안정하다고 지적한다.
• 수학·추론 과제: Bansal et al.(2022) 등 반복 모델로 반복 과정을 푸는 연구. 이들은 유한 상태·선택지 과제이고, 회귀는 연속 출력이라는 점이 다르다.
• in-context learning 이론: 선형 attention이 경사하강을 구현한다는 구성·이론(von Oswald, Ahn, Zhang, Mahankali 등), 다층 transformer의 고차 최적화(Fu et al. 2023; Giannou et al. 2024), Bayesian 추정기 관점 등.
• looped transformer의 표현력: Giannou et al.(2023)은 encoder looped transformer가 범용 컴퓨터로 동작할 수 있음을 보였고, Bai et al.(2019)의 Theorem 3은 임의의 L층 네트워크를 weight-tied·input-injected 네트워크로 표현할 수 있음을 보였다(Section 6).

이 논문의 차이 (author-stated). Garg et al.(2022)과 같은 구조(softmax attention이 있는 GPT-2 decoder)를 쓰고, 층 수만 1로 줄여 반복한다. 표현력 증명이 아니라 looped decoder transformer를 실제로 학습시키는 방법과 경험적 비교에 집중한다(Section 2, 6).

지표

지표	무엇을 세나	단위·범위	방향	분모·정규화와 측정 조건	원문 위치
Squared error	질의 x의 예측과 정답 f(x)의 차이 제곱 (M(P) − f(x_test))²	0 이상 실수. 그래프는 선형 또는 로그 축	낮을수록 좋음	문제 차원과 무관하게 Garg et al.(2022) 방식으로 정규화. 1280개 프롬프트 평가, 1000회 bootstrap으로 90% 신뢰구간	Section 3, 5.1
Squared error vs loop iterations	시험 시 반복 횟수 t에 따른 위 오차	같음	낮고 평평할수록 좋음(고정점 안정)	주로 in-context 예시 k=40의 마지막 예측 기준으로 보이나, 어떤 예시 수에서 잰 것인지 그림마다 명시되지 않음(unknown)	Figure 2 왼쪽, 3, 4, 7, 9
Squared error vs in-context examples	예시 수에 따른 위 오차	같음	낮을수록, 빨리 떨어질수록 좋음	baseline 0.0(완벽)과 1.0(점선, 0을 예측하는 수준으로 해석됨, analyst-inference)	Figure 2 오른쪽, 6, 10, 11
Probe MSE	각 층·반복의 표현에서 2층 MLP probe가 XᵀY 또는 w_OLS를 복원할 때의 오차	실수	낮을수록 그 정보가 표현에 담겨 있음	문제 차원 d로 나눔, 여러 문맥 길이 평균. 점선은 w를 1로 고정한 control task의 최소 오차	Figure 8, Appendix F
Test accuracy (OpenML)	보지 못한 데이터셋에서 in-context 이진 분류 정확도	0–1	높을수록 좋음	10개 중 9개로 학습, 남은 1개로 시험	Appendix G Table 2 (부록)

방법과 벤치마크 구성

1. 문제 설정 (Section 3). 함수 f를 분포 D_F에서, 입력 x₁…x_k와 x_test를 D_X에서 뽑아 프롬프트 P=(x₁, f(x₁), …, x_k, f(x_k), x_test)를 만들고, f(x_test)를 맞힌다. 모델은 GPT-2 decoder. 일반 transformer는 L=12(Garg et al. 설정), looped transformer는 L=1.

2. 반복 구조와 input injection (Section 4, 4.1). t번 반복한 출력은 Y_t(Pⁱ|θ) = M_θ(M_θ(⋯M_θ(Y₀ⁱ + Pⁱ) + Pⁱ)⋯ + Pⁱ), Y₀는 0 텐서. 즉 매 반복 입력 Pⁱ를 더한다. 이를 빼고 Y_{t+1}=M(Y_t), Y₀=P로 두는 weight-tying 형태(ALBERT, Universal Transformer 방식)는 t가 커지면 초기 입력의 영향이 사라져 해가 무작위해진다고 저자들은 설명하며, Bai et al.(2019)과 Bansal et al.(2022)을 근거로 든다(author-stated). 주입 방법이 덧셈에 한정되지는 않는다고 덧붙인다.

3. Truncated loss (Section 4, 식 (1)). 손실은 반복 b₀부터 b까지(b₀ = max(b−T, 0))의 출력에 대해, 모든 프롬프트 접두사(예시 0개~k개)의 예측 오차를 평균한 것이다. truncated BPTT에서 착안해 계산을 아낀다. b는 학습 시 최대 반복 수(순전파 시간 결정), T는 손실 창 크기(메모리 결정)다(Section 6).

4. b와 T의 선택 (Section 4.2, Figure 4). 선형회귀에서 b ∈ {12, 20, 30, 40, 50}, T ∈ {5, 10, 15}를 비교했다. b가 작으면 학습 반복 이후 발산하고, b가 크면 발산하지 않는 고정점을 찾지만 일정 값을 넘으면 수렴 속도가 오히려 줄어든다. T가 작으면 성능이 떨어지고, 크면 학습이 불안정해진다(Jaeger 2005 인용). 기준은 “발산하지 않는 최소 b”와 “메모리를 아끼는 작은 T”이고, 선형회귀는 b=20, T=15를 택했다.

5. 학습 설정 (Section 4, 5.1). D=256, h=8. 파라미터 수는 일반 9.48M, looped 0.79M. Adam, 학습률 0.0001, weight decay·gradient clipping·데이터 증강 없음(Garg et al.과 공정 비교를 위해 정규화를 뺐다고 Section 6에서 밝힘). Garg et al.의 d, k curriculum을 따르고, b도 학습 중 점점 늘리는 schedule을 쓴다. 저자들은 이 schedule이 대부분 결과에 큰 영향이 없고 결정트리에서만 더 나았다고 한다(Section 5.1, Appendix E).

6. 과제 (Section 5.1, 5.4, Appendix A).

• 선형 함수: d=20, k=40, w와 x 모두 N(0, I_d).
• sparse 선형 함수: d=20, k=100(Appendix B는 k=40으로 적음), 0이 아닌 성분 s=3.
• 결정트리: 깊이 4, d=20. 각 내부 노드에서 무작위 좌표의 부호로 분기, 잎 값은 N(0, I)에서 초기화.
• 2층 ReLU 신경망: 은닉 100, d=20.
• 과제별 b, T: sparse linear b=20·T=10, 결정트리 b=70·T=15, 2층 NN b=12·T=5(Section 5.4). 구조 변형과 sparsity 실험은 b=30·T=15.
• 비교 대상: 12층 transformer, 최소제곱, 평균(Averaging), Lasso(α=0.01 선택), 3-NN, greedy tree learning, XGBoost, 경사하강으로 학습한 2층 NN.
• 분석 실험: 학습 중 본 서로 다른 프롬프트 수(d=10, k=20, curriculum 끔), 분포 밖 입력(공분산 왜곡, 출력 잡음, 입력 3배 스케일), L·D·h 변화, 층별 probing.

주요 결과

주장	근거 위치	관찰/수치	조건·주의	출처 라벨
1층 looped transformer가 12층 transformer와 선형회귀에서 대등	Figure 2 오른쪽, Section 5.2	두 곡선과 최소제곱이 거의 겹침. 파라미터 0.79M vs 9.48M	학습 분포 안(in-distribution) 평가. 본문은 b=20·T=15, Figure 2 캡션은 “30 loop iterations로 학습”이라고 다르게 적음	author-stated
학습한 반복 수를 넘어도 고정점에 머문다	Figure 2 왼쪽	약 20회 반복에서 약 10⁻³으로 떨어진 뒤 100회까지 평평(그래프 판독)	오차 바닥(error floor)이 있음(Section 6)	author-stated + analyst-inference
input injection이 없으면 학습 반복 이후 발산	Figure 3, Section 4.1	weight-tying은 학습 반복 20 근처에서 최저 후 100회에 약 10⁰까지 상승, input injection은 약 10⁻³ 유지(그래프 판독)	선형회귀	author-stated
b가 작으면 발산, 크면 수렴하지만 너무 크면 수렴이 느려짐	Figure 4, Section 4.2	T=15에서 b=40, 50의 수렴 속도는 비슷. b=12는 effective depth가 12층 transformer와 같지만 성능을 재현하지 못함	선형회귀만의 sweep	author-stated
looped transformer는 더 적은 서로 다른 프롬프트로 학습한다	Figure 5, Appendix C	looped 곡선이 transformer보다 먼저 0 근처로 떨어짐(약 10k vs 약 40k, 그래프 판독)	d=10, curriculum 없음. 저자는 원인을 적은 파라미터로 돌림. 본문은 k=20, 캡션은 k=21	author-stated + analyst-inference
분포 밖에서 “더 단순한 해” 쪽 귀납 편향	Figure 6, Section 5.2	(a) 공분산 왜곡에서 더 나음, (b) 잡음 선형회귀에서 double descent 정점이 더 낮음, (c) x=3z 스케일에서 더 나쁨	본문은 (c)를 “서로 다른 orthant”라고 소개하지만 그림의 (c)는 스케일 실험이다(부록 Figure 15에 orthant 실험이 있음)	author-stated
looped transformer는 같은 성능에 층 수보다 많은 반복이 필요	Figure 7 왼쪽, Section 5.3	1층 looped는 약 20회 반복으로 8층 transformer 성능에 도달	저자는 두 모델이 층(반복)마다 다른 표현을 배운다고 해석	author-stated
D·h의 영향	Figure 7 오른쪽, Section 5.3	D를 키우면 수렴이 좋아지다 D=256에서 포화. h는 2→8로 갈 때 오차 감소, 8→32에서 증가(헤드당 차원 D/h 감소와의 trade-off)	L=1 기준	author-stated
looped transformer는 반복마다 표현을 꾸준히 다듬는다	Figure 8, Appendix F	transformer는 약 10번째 층에서 probe 오차가 최저이고 이후 증가, looped는 반복할수록 계속 감소	control task 점선보다 모두 낮아 probe가 in-context 정보를 쓰고 있음	author-stated
필요한 b, T가 과제 난이도를 반영한다	Figure 9, Section 5.4, 6	T=10에서 선형회귀는 b=20에서 발산, sparse linear는 수렴. 2층 NN은 오히려 더 작은 b, T로 충분	저자 제안: b, T를 transformer 관점의 과제 난이도 지표로 쓸 수 있다	author-stated
sparse linear에서 transformer와 Lasso보다 빠르게 오차가 준다	Figure 10(a), Section 5.4	예시 40개에서 Lasso 1.32e-04, 최소제곱 8.21e-14, looped 6.12e-04, transformer 0.0017. 예시 10개 미만에서는 looped가 Lasso보다 작음	”outperform”은 최종 정밀도가 아니라 예시 수에 따른 추세를 뜻한다고 저자가 명시	author-stated
결정트리·2층 NN에서는 transformer와 대등	Figure 10(b)(c)	두 곡선이 비슷하게 겹침(그래프 판독)	파라미터 1/12	author-stated
모든 sparsity 수준에서 looped가 같거나 낫다	Figure 11	s=1, 3, 6, 9, 12 모두, 특히 예시가 적을 때 차이	b=30, T=15	author-stated
실제 표 데이터(OpenML)에서도 대등하거나 일부 더 나음	Appendix G Table 2	예: 847번 0.668 → 0.757, 819번 0.483 → 0.568, 871번 0.532 → 0.51	부록 실험, 10개 데이터셋 leave-one-out	author-stated

결과는 “학습법 → 대등성 → 편향” 순으로 이어진다. Section 4의 그림들(Figure 3, 4)은 input injection과 충분히 큰 b가 있어야 학습 범위 밖에서도 발산하지 않는 고정점을 얻는다는 것을 보인다. 이 학습법으로 Section 5는 1/12 파라미터로 표준 transformer와 대등한 in-context 학습을 보인다. 마지막으로 Figure 6, 10, 11은 looped 모델이 단순·희소한 해를 선호한다는 공통된 해석으로 묶이며, 이 편향이 sparse 문제에서는 이득, 입력 스케일 이동에서는 손해로 나타난다(analyst-inference: “단순성 편향”은 저자의 추정이며 직접 측정한 양은 아니다. 저자도 “We conjecture”라고 쓴다).

Take-home

핵심 메시지. 1층 transformer 블록을 입력 주입과 함께 반복하고 마지막 T회 반복에 손실을 거는 학습법이면, 학습 범위 밖 반복에서도 안정적인 고정점 해를 얻으며 12층 transformer의 약 1/12 파라미터로 같은 in-context 학습 성능을 낸다.

성립 범위.

• 합성 함수 클래스(선형, sparse 선형, 깊이 4 결정트리, 2층 ReLU NN), d=20 수준, 입력이 N(0, I)인 학습 분포 안에서 확인했다(Section 5). OpenML 실험은 부록에 있다.
• 비교는 Garg et al.(2022) 설정의 12층 GPT-2(D=256, h=8)와의 비교다.
• b, T는 과제마다 하이퍼파라미터 탐색으로 정했다(Section 5.4, Appendix B).

저자가 인정한 한계 (author-stated, Section 6).

• 고정점 해에 고유한 오차 바닥이 있고 분포 밖 프롬프트에서 성능이 제한적이다. 학습한 모델은 임의 입력 분포에 일반화하는 진짜 알고리즘이 아니라, 학습한 입력 분포에서만 알고리즘을 흉내 낸다.
• b, T 선택에 튜닝이 필요하고, T가 크면 학습이 불안정해질 수 있다. 공정 비교를 위해 mixed precision, gradient clipping, weight decay 같은 안정화 기법을 쓰지 않았다.
• looped decoder transformer의 표현력 분석은 미래 과제로 남긴다.

작성자가 보는 우려 (analyst-inference).

• “Better at learning learning algorithms”라는 제목에 비해, 주된 증거는 “1/12 파라미터로 대등”이다. 명확한 우위는 sparse linear(추세)와 적은 프롬프트 수의 학습 효율에서 나오고, 분포 밖 스케일 이동에서는 오히려 열세다.
• 계산량은 줄지 않는다. b=20 이상 반복하므로 순전파 계산은 12층 transformer보다 많을 수 있다. 원문은 파라미터 수만 비교하고 실행 시간·FLOPs는 보고하지 않았다.
• “고정점”은 오차 곡선이 평평해지는 관찰이며, 은닉 상태가 실제로 수렴하는지(‖Y_{t+1} − Y_t‖)는 본문에서 직접 측정하지 않았다.
• 원문 내부 불일치: Figure 2 캡션(30회 학습)과 본문(b=20), Figure 5 캡션(k=21)과 본문(k=20), Section 5.2 본문의 Figure 6(c) 설명(orthant)과 실제 그림(스케일 3), sparse linear의 k(Section 5.4는 100, Appendix B는 40).

재현성. 코드 공개(https://github.com/Leiay/looped_transformer, 각주 1). 평가는 1280 프롬프트·1000회 bootstrap 90% 신뢰구간. 학습 seed 수(독립 학습 반복)는 본문과 부록에서 보고되지 않았다. 하드웨어와 학습 시간도 보고되지 않았다.

본문 도표

Figure 1 — 문제 설정과 이 논문의 접근

• 무엇을 보여주는가: 선형최소제곱 문제를 decoder-only transformer로 푸는 설정. 위는 프롬프트 (x₁, y₁, …, x_k, y_k, x_test)와 출력 ŷ, 아래는 “학습된 TF 해법”과 “전통 해법”의 대응이다. 일반 transformer는 닫힌 해 ŵ=(XᵀX)⁻¹Xᵀy에, 이 논문의 looped transformer는 반복 알고리즘(경사하강 w_{t+1} ← w_t − η(XᵀXw_t − Xᵀy))에 대응시킨다.
• 읽는 법: looped transformer 그림에서 빨간 블록 하나가 자기 출력을 다시 받는 화살표와 ⊕(입력 더하기) 기호가 있다.
• 핵심 관찰: 개념도이며 수치는 없다.
• 의미와 주의: 저자 목표는 표준 transformer와 같은 성능을 더 적은 파라미터로 내는 것(캡션, author-stated). 그림은 대응 관계를 제안할 뿐, 학습된 looped transformer가 실제로 경사하강을 구현한다고 증명하지는 않는다(analyst-inference). Figure 8의 probing이 이에 대한 간접 증거다.
• 확인 상태: viewed

Figure 2 — 선형회귀에서의 예비 결과

• 무엇을 보여주는가: 왼쪽은 시험 반복 횟수에 따른 오차, 오른쪽은 in-context 예시 수에 따른 오차.
• 읽는 법: 왼쪽 x축 loop iterations(0–100), y축 squared error(로그, 10⁰–10⁻³). 오른쪽 x축 in-context examples(0–40), y축 squared error(0–1). 회색 Least squares, 빨강 Looped transformer, 파란 점선 Transformer.
• 핵심 관찰: 왼쪽에서 오차가 약 20회 반복까지 급감한 뒤 약 10⁻³에서 100회까지 평평하다. 오른쪽에서 세 곡선이 거의 겹치며 예시 약 20개에서 0 근처에 이른다(그래프 판독, analyst-inference). d=20이므로 예시 수가 차원과 같아질 때 해가 결정되는 것과 맞는다(analyst-inference).
• 의미와 주의: 캡션은 1/12 파라미터로 12층 transformer와 대등하고 최소제곱에 근접하며, 학습 반복을 넘어서도 안정적인 고정점 해를 얻는다고 말한다(author-stated). 캡션의 “30 loop iterations로 학습”은 Section 5.2의 b=20과 다르다.
• 확인 상태: viewed

Figure 3 — input injection 유무

• 무엇을 보여주는가: 선형회귀에서 input injection을 쓴 looped transformer와 쓰지 않은 weight-tying 모델의 시험 오차를 반복 횟수에 따라 비교.
• 읽는 법: x축 loop iterations(0–100), y축 squared error(로그). 세로 점선은 학습 반복 수(20). 빨강 Looped Transformer, 파랑 Weight Tying.
• 핵심 관찰: 두 모델 모두 20회 근처까지 오차가 줄어든다. 그 뒤 weight-tying은 오차가 다시 커져 100회에 약 10⁰에 가까워지고, input injection 모델은 약 10⁻³에서 평평하다(그래프 판독).
• 의미와 주의: 저자 해석은 input injection이 없으면 학습 반복 이후 빠르게 발산한다는 것(Section 4.1, author-stated). 같은 학습 조건(b, T)에서의 단일 비교이며 seed 반복 여부는 보고되지 않았다.
• 확인 상태: viewed

Figure 4 — b와 T의 영향 (선형회귀)

• 무엇을 보여주는가: 학습 시 b(최대 반복)와 T(손실 창)를 바꿨을 때 시험 반복 수에 따른 오차.
• 읽는 법: 왼쪽부터 T=5, 10, 15. 색은 b=12, 20, 30, 40, 50. 실선은 시험 오차 곡선, 같은 색 세로 점선은 그 모델의 b. x축 0–200회, y축 로그.
• 핵심 관찰: 세 패널 모두 b=12(진한 파랑)는 학습 반복 뒤 오차가 다시 올라 약 10⁻¹ 근처까지 간다. T=10에서는 b=20(연한 파랑)도 학습 반복 뒤 크게 올라간다. b가 30 이상이면 200회까지 평평하다. b가 클수록 오차가 최저에 도달하는 반복 수가 뒤로 밀린다(그래프 판독, analyst-inference).
• 의미와 주의: 저자 결론은 “발산하지 않는 최소 b, 작은 T”이며 선형회귀는 b=20, T=15를 택했다(author-stated). 또 looped 모델은 학습 b 이전에 포화하는 고정점을 찾는다고 해석한다. T=10에서 b=20이 발산한다는 점은 b와 T가 서로 얽혀 있음을 보여 준다(analyst-inference).
• 확인 상태: viewed

Figure 5 — 학습 중 본 서로 다른 프롬프트 수

• 무엇을 보여주는가: 미리 만든 고정 학습 집합의 크기(서로 다른 프롬프트/함수 수)에 따른 시험 오차.
• 읽는 법: x축 distinct prompts/functions seen during training(1k–80k, 로그 간격), y축 squared error. 파랑 Transformer, 빨강 Looped Transformer. 점선 1.0.
• 핵심 관찰: 두 모델 모두 1k에서 약 1.1이다. looped transformer는 약 10k에서 0 근처로 떨어지고, transformer는 약 40k에서야 0 근처에 도달한다(그래프 판독).
• 의미와 주의: 저자는 looped transformer의 파라미터가 적어 더 적은 프롬프트로 학습한다고 해석한다(author-stated). 다른 차원(d=5, 20) 결과는 Appendix C. 캡션은 k=21, 본문은 k=20으로 적혀 있다.
• 확인 상태: viewed

Figure 6 — 분포 밖 평가

• 무엇을 보여주는가: N(0, I) 입력의 선형회귀로 학습한 모델을 (a) 공분산이 왜곡된 입력, (b) 잡음 섞인 선형 함수, (c) 3배로 스케일한 입력에서 평가.
• 읽는 법: x축 in-context examples(0–40), y축 squared error. 진한 파랑 Transformer L=12, 연한 파랑~빨강 계열 Looped transformer b=20, 30, 50, 100, 진한 빨강 Least squares. 점선 1.0.
• 핵심 관찰: (a) 모든 모델이 낮지만 looped 곡선이 transformer보다 약간 아래에 있다. (b) 예시 수 약 20(=d)에서 오차가 튀는 double descent 정점이 looped에서 더 낮다. (c) 최소제곱은 예시 약 20개에서 0이 되지만, transformer는 약 4–5에서 머물고 looped transformer는 약 7.5–10으로 더 높다(그래프 판독, analyst-inference).
• 의미와 주의: 저자 해석은 looped transformer가 진짜 알고리즘을 익히기보다 “더 단순한 해”를 선호하는 귀납 편향을 보이며, 이것이 (a)(b)에서는 이득, (c)에서는 손해라는 것(author-stated). 본문은 (c)를 “서로 다른 orthant”라고 소개하지만 그림은 스케일 실험이다.
• 확인 상태: viewed

Figure 7 — 층 수, 헤드 수, 임베딩 차원

• 무엇을 보여주는가: 왼쪽은 L층 일반 transformer(반복 없음, 수평선)와 L층 looped transformer(반복에 따른 곡선), 오른쪽은 looped transformer의 L, D, h 조합.
• 읽는 법: x축 loop iterations(0–100), y축 squared error(로그). 왼쪽 범례 unlooped L=1, 2, 4, 8, 12, 20과 looped L=1, 2, 4. 오른쪽 범례는 (L=1, D=256에서 h=2, 4, 16, 32), (L=1, h=8에서 D=64, 128, 384, 512), (D=256, h=8에서 L=1, 2, 4).
• 핵심 관찰: 왼쪽에서 unlooped L=1은 약 10⁰, L=2는 약 10⁻¹, L=4는 약 10⁻² 수준의 수평선이고, L=8 이상은 약 10⁻³ 근처다. looped L=1, 2, 4는 모두 약 20–30회 반복 뒤 약 10⁻³ 근처에 도달하며 L이 클수록 빨리 내려간다. 오른쪽에서 D=64는 약 10⁻¹에서 멈춰 가장 나쁘다(그래프 판독, analyst-inference).
• 의미와 주의: 저자는 looped transformer가 L층 transformer 성능에 L보다 많은 반복이 필요하다는 점(1층이 8층 성능에 약 20회)을 근거로 두 모델이 층(반복)마다 다른 표현을 배운다고 해석한다. D는 256에서 포화, h는 8이 최적이다(author-stated).
• 확인 상태: viewed

Figure 8 — 층·반복별 probing

• 무엇을 보여주는가: 각 층(transformer) 또는 반복(looped)의 출력에서 2층 MLP probe가 XᵀY(경사 성분)와 w_OLS(최소제곱 해)를 얼마나 잘 복원하는지.
• 읽는 법: x축 Layer / Loop Iteration(0–20), y축 MSE. 파랑 막대 Transformer, 빨강 막대 Looped Transformer. 위쪽 점선은 w를 1로 고정한 control task의 최소 probe 오차.
• 핵심 관찰: 두 probe 모두 transformer는 약 10번째 층에서 오차가 가장 낮고 그 뒤 다시 커진다. looped transformer는 반복이 늘수록 오차가 꾸준히 줄어 20회에서 transformer보다 낮다. 모든 막대가 control 점선(약 0.04 이상)보다 낮다(그래프 판독).
• 의미와 주의: 저자 해석은 transformer는 처음에 XᵀY와 w_OLS 관련 표현을 담았다가 다른 통계로 초점을 옮기고, looped transformer는 반복마다 목표 관련 표현을 계속 다듬는다는 것(author-stated). probe 오차 차이가 0.01 수준이고, probing은 정보가 “복원 가능하다”는 것만 보이며 모델이 그것을 계산에 쓴다는 인과 증거는 아니다(analyst-inference).
• 확인 상태: viewed

Figure 9 — 과제별 b의 영향 (T=10)

• 무엇을 보여주는가: (a) sparse linear, (b) 결정트리, (c) 2층 ReLU NN에서 T=10, 여러 b로 학습한 looped transformer의 시험 반복별 오차.
• 읽는 법: Figure 4와 같은 방식. (a)(c)의 b는 12–50, (b)의 b는 12, 50, 70, 100. 세로 점선은 각 b.
• 핵심 관찰: (a) b=12만 학습 반복 뒤 약 0.5까지 올라가고 나머지는 낮게 평평하다. (b) b=12는 반복 뒤 약 0.5로 오르고, b=50·70은 약 0.08–0.1, b=100은 약 0.12 근처에서 머문다. (c) 모든 b가 약 0.15–0.2 근처로 수렴하고 발산하지 않는다(그래프 판독, analyst-inference).
• 의미와 주의: 저자는 Figure 4와 비교해 T=10에서 선형회귀는 b=20에서 발산하지만 sparse linear는 수렴한다는 점, 2층 NN은 어려운 과제인데도 더 작은 b, T로 충분하다는 점을 들어 b, T가 transformer 관점의 과제 난이도를 반영한다고 본다(author-stated).
• 확인 상태: viewed

Figure 10 — 더 복잡한 함수 클래스에서의 비교

• 무엇을 보여주는가: (a) sparse linear, (b) 결정트리, (c) 2층 ReLU NN에서 in-context 예시 수에 따른 오차를 looped transformer, transformer, 전통 해법과 비교.
• 읽는 법: x축 in-context examples((a) 0–40, (b)(c) 0–100), y축 squared error. 점선 1.0. 범례: (a) Transformer, Least Squares, Averaging, Lasso, Looped Transformer. (b) Transformer, 3-Nearest Neighbors, Greedy Tree Learning(± sign preproc.), XGBoost(± sign preproc.), Looped Transformer. (c) Transformer, Least Squares, 3-NN, 2-layer NN GD, Looped Transformer.
• 핵심 관찰: (a) looped transformer가 예시 약 10개에서 0 근처로 가장 먼저 떨어지고, transformer는 약 20개에서 0 근처, Averaging은 높게 남는다. (b) looped와 transformer가 가장 낮은 무리에서 비슷하게 내려간다. (c) looped와 transformer가 약 0.2 수준으로 겹친다(그래프 판독). 본문 수치(예시 40개): Lasso 1.32e-04, 최소제곱 8.21e-14, looped 6.12e-04, transformer 0.0017.
• 의미와 주의: 저자는 looped가 모든 과제에서 transformer와 대등하고, sparse linear에서는 예시 수에 따른 추세 면에서 transformer와 Lasso를 넘는다고 말한다. 최종 정밀도 우위가 아님을 명시한다(author-stated). 성능 향상의 원인은 “더 희소한 해를 선호하는 귀납 편향”이라는 추측이다.
• 확인 상태: viewed

Figure 11 — sparsity 수준별 비교

• 무엇을 보여주는가: d=20 sparse linear 함수에서 0이 아닌 성분 수 s=1, 3, 6, 9, 12별로 looped transformer(실선)와 transformer(점선)의 오차.
• 읽는 법: x축 in-context examples(0–40), y축 squared error. 색은 s(진한 파랑 s=1부터 빨강 s=12까지).
• 핵심 관찰: 같은 s에서 실선이 점선보다 먼저 떨어진다. s가 클수록 두 곡선 모두 0에 도달하는 예시 수가 늘어 약 20개 근처에서 0이 된다(그래프 판독).
• 의미와 주의: 저자 해석은 모든 sparsity 수준에서 looped가 같거나 더 낫고, 특히 예시가 적을 때 더 정확하다는 것(author-stated). b=30, T=15로 학습했다.
• 확인 상태: viewed

Wiki 연결

• Looped Transformers as Programmable Computers — 이 논문이 looped transformer 구조의 출발점으로 인용하는 Giannou et al.(2023). 표현력(구성적 증명) 쪽이 그 논문, 실제 학습법과 경험적 비교 쪽이 이 논문이다.
• Deep Equilibrium Models — 원문이 implicit model의 극단으로 인용하는 Bai et al.(2019). input injection과 “weight-tied·input-injected 네트워크의 보편성”(Theorem 3)이 이 논문의 설계 근거다.
• End-to-end Algorithm Synthesis with Recurrent Networks - Logical Extrapolation Without Overthinking — 원문이 input injection이 없으면 반복이 무작위해진다는 근거로 인용하는 Bansal et al.(2022). recall = input injection이라는 같은 처방을 합성곱 반복 모델에서 먼저 제시했다.
• Reasoning with Latent Thoughts - On the Power of Looped Transformers — looped transformer의 추론 능력과 귀납 편향을 이론·실험으로 잇는 후속 계보.
• Fixed-Point Reasoners - Stable and Adaptive Deep Looped Transformers — looped transformer가 학습 반복 밖에서도 고정점에 머물게 하는 안정화 문제를 직접 다루는 후속 연구.

BibTeX

@inproceedings{yang2024looped,
  title     = {Looped Transformers are Better at Learning Learning Algorithms},
  author    = {Yang, Liu and Lee, Kangwook and Nowak, Robert D. and Papailiopoulos, Dimitris},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2024},
  eprint    = {2311.12424},
  archivePrefix = {arXiv}
}