12장 — 디퓨전 트랜스포머: U-Net 자리에 트랜스포머를

자주 하는 실수와 요약

자주 하는 실수

실수 나온 문제 바로잡는 법
멀리 닿는 말을 여러 번 하면 정밀도도 올라간다고 봄 1 10cm 단위의 말은 몇 번을 해도 10cm 단위다. 싣는 굵기와 마지막 손질과의 거리가 함께 정한다
거친 칸에 실린 정보는 그 칸 크기 단위로만 자리를 안다고 봄 2 채널에 칸 안의 자리를 수로 적을 수 있다. 문제는 그 수가 올라오는 길에서 이웃만 보는 합성곱을 거쳐 끝까지 실려 가느냐다
매개변수의 87%가 선형 층이면 선형 층 가속이 계산의 87%에 먹힌다고 봄 3 합성곱은 칸마다 같은 커널을 되풀이해 매개변수 몫(13.0%)보다 계산 몫(27.2%)이 크다. 빨라지지 않는 몫이 남아 SDXL 은 1.57배에서 멈춘다
패치 한 변을 반으로 줄이면 토큰이 두 배라고 봄 4, 6 가로·세로가 모두 두 배라 토큰은 네 배, 어텐션 쌍은 16배
패치를 작게 자르면 모델이 가벼워진다고 봄 5 바뀌는 것은 패치를 옮기는 층뿐(전체의 0.04%). 매개변수는 거의 그대로이고 연산량이 네 배씩 는다
그림을 두 배로 키우면 어텐션 쌍도 두 배라고 봄 6 한 변이 두 배면 토큰이 네 배, 쌍은 16배
문을 0으로 둔 블록은 빼 둔 블록과 같다고 봄 7 출력은 같아도 블록은 꽂혀 있어, 문이 열릴 만큼의 학습 신호가 그 블록에서 나온다
0에서 시작하면 기울기가 모두 0이라 배우지 못한다고 봄 8 문 자신의 기울기는 0이 아니다(f × c). 문이 곱하는 f 까지 0으로 묶으면 그때는 멈춘다
변조 층이 매개변수의 27%면 계산도 27%라고 봄 9 변조 층은 조건 벡터 하나에 그림마다 한 번 걸린다. 토큰마다의 계산에는 들지 않는다
두 줄 블록은 가중치가 두 벌이라 어텐션 쌍도 두 배라고 봄 10 질의·열쇠·값을 이어 어텐션은 한 번. 토큰 하나는 제 줄의 가중치만 지난다
크로스 어텐션 모델에서도 글 토큰이 몸통을 지나며 바뀐다고 봄 11 크로스 어텐션에서 글은 읽히기만 한다. 글 쪽 일은 몸통 앞의 어댑터가 맡는다
토큰이 4배면 연산량도 정확히 4배라고 봄 12 행렬 곱은 토큰 수에 비례하지만 어텐션 쌍은 제곱으로 는다. DiT-XL/2 를 512 × 512 로 쓰면 4.4배
두 줄 블록의 매개변수 36d² 를 토큰마다의 계산에 그대로 씀 13 그림 토큰은 그림 쪽 한 벌, 그중에서도 변조 층을 뺀 12d² 만 지난다
같은 그림 크기면 어느 모델이나 몸통을 지나는 토큰 수가 같다고 봄 14 합동 어텐션 모델은 글 토큰도 몸통을 지난다. 긴 프롬프트일수록 그 몫이 커진다
공개된 몸통을 빌려도 블록 구조를 마음대로 바꿀 수 있다고 봄 15 블록 수·폭·어텐션 구조·위치 축·잠재 채널은 가중치에 묶여 있다. 바꿀 수 있는 것은 글 인코더와 그 사이의 어댑터다
받는 쪽이 이미 가진 것을 잊고 보낼 것의 길이를 셈 16 젖은 사진에서 원래 송장을 빼면 얼룩이 나온다. 좁은 곳에서 움직이는 쪽만 보내면 된다
출력이 D개의 수면 D개 방향을 모두 맞힐 수 있다고 봄 17 마지막 층 앞의 폭이 d 면 출력은 d차원의 평평한 곳에 갇힌다. 모든 방향에 퍼진 잡음·속도는 손실을 D − d 아래로 줄이지 못한다
원래 그림도 D개의 수라 잡음처럼 막힌다고 봄 17 원래 그림은 낮은 차원의 면 위에서만 움직인다. 잡음 몫은 식 속의 xₜ 가 채운다
원래 그림을 맞히면 짙은 잡음에서 흐릿한 평균이 나와 나쁘다고 봄 18 끝까지 배운 답은 출력 꼴과 상관없이 같은 조건부 평균이다. 넉넉한 신경망에서 다른 것은 손실의 가중치뿐이다
패치를 좁은 층으로 줄이면 정보를 버려 해롭다고 봄 19 원래 그림을 맞힐 때는 그림이 사는 좁은 면만 지나가면 된다. 리·허는 768개의 수를 16개로 줄여도 무너지지 않았다

요약

U-Net 의 멀리 읽는 길은 해상도를 낮춘 몇몇 단계에만 있고, 칸 하나 단위의 자리는 마지막에 이웃만 보는 합성곱이 정한다. 디퓨전 트랜스포머의 정보 통로는 하나다. 모든 블록에서 모든 토큰이 서로 읽는다. 들어오는 꼴도 좌표가 붙은 토큰의 줄 하나라서 가로세로 비, 글, 참조 그림, 영상 프레임을 같은 줄에 이어 붙이고, 언어모델을 위해 다듬은 어텐션·저정밀 행렬 곱 같은 연장을 그대로 빌려 쓴다(FLUX.1 은 층 계산이 모두 선형 층, SDXL 은 27.2%가 합성곱). 작게 그린 손가락과 글자가 무너지는 것은 몸통이 아니라 압축기의 되살리기 천장 몫이다. 트랜스포머는 토큰 줄을 받으므로, 잠재 그림을 한 변 a 칸짜리 패치로 잘라 줄로 펴고, 조각마다 좌표를 붙여 위치를 알린다. 패치 한 변은 매개변수가 아니라 토큰 수, 곧 연산량을 정하며, 한 변을 반으로 줄이면 토큰은 네 배, 어텐션 쌍은 16배가 된다. 오늘의 모델들은 1024 × 1024 그림을 모두 토큰 4,096개로 만든다. 시간은 블록마다 정규화 뒤의 배율·치우침과 출력의 문으로 넣고, 그 층을 0에서 시작해 블록이 항등 함수로 출발하게 한다(adaLN-Zero). DiT 의 비교에서 이 블록은 토큰을 붙이는 블록의 거의 절반 FID 를 냈다. 다만 변조 층이 매개변수의 3분의 1을 차지해서, 뒤의 모델들은 한 벌을 나눠 쓰는 쪽으로 덜어 냈다. 글은 글 토큰 줄과 그림 토큰 줄을 이어 셀프 어텐션 하나로 함께 처리하는 합동 어텐션으로 넣는다. 두 줄 블록은 줄마다 가중치를 따로, 한 줄 블록은 한 벌로 같은 일을 하며, 어텐션 쌍과 토큰마다의 계산은 같고 매개변수만 다르다. 몸통을 키울 때 품질을 따라가는 것은 매개변수 수보다 연산량이고, 블록 몸통은 대략 18d² × 깊이의 매개변수를 갖고, 토큰 하나는 그 가운데 변조 층을 뺀 12d² × 깊이 짝의 곱셈·덧셈을 지난다. FLUX.1·FLUX.2·Qwen-Image·Z-Image·Krea 2·Anima 는 이 네 부품을 저마다 다르게 조립했고, 흐름은 두 줄에서 한 줄로, 변조 층은 가볍게, 글 인코더는 크게 가고 있다. 이 모델들의 토큰 하나는 폭의 수십 분의 1인 수만 담아, 잡음이나 속도를 맞혀도 문제가 드러나지 않았다. 압축기 없이 픽셀을 16 × 16칸 같은 큰 패치로 자르면 토큰 하나의 수가 폭에 닿는다. 마지막 층은 폭만큼의 방향밖에 내지 못하므로 모든 방향에 퍼진 잡음과 속도는 맞힐 수 없지만, 낮은 차원의 면 위에 있는 깨끗한 그림은 같은 신경망으로도 맞힌다. 신경망이 깨끗한 그림을 내놓고 속도는 식으로 바꿔 얻는 원래 그림 예측(리·허의 JiT)은 그래서 압축기 없이도 버틴다. 끝까지 배운 답은 세 출력 꼴이 같고, 폭이 넉넉하면 남는 차이는 손실의 가중치다.

flowchart LR
  U["U-Net<br/>멀리 읽는 길은 거친 단계에만"] -->|"통로 하나, 토큰 줄 하나"| D
  A["잠재 그림<br/>128 × 128 × 16"] -->|"패치 2 × 2<br/>좌표 붙이기"| B["그림 토큰 4,096개"]
  T["시간 t"] --> M["조건 벡터 하나"]
  M -->|"배율·치우침·문<br/>0에서 시작"| D
  G["글 인코더의 토큰 줄"] -->|"이어 붙이기"| D["트랜스포머 블록 × L<br/>두 줄 → 한 줄"]
  B --> D
  D -->|"패치를 되돌리기"| O["속도·잡음·원래 그림 예측"]
  X["픽셀 큰 패치<br/>토큰 하나의 수 ≥ 폭"] -.->|"원래 그림을 맞혀야 버틴다"| O
  S["키우기: 폭 d · 깊이 L · 토큰 N"] -.->|"연산량 ≈ 12d²L × N + 2N²dL"| D

막힌 곳

이제 글 토큰 줄과 시간을 받아 잡음 섞인 잠재 그림의 속도를 내놓는 트랜스포머를 지을 수 있다. 글은 합동 어텐션으로 층마다 그림과 섞이고, 몸통이 크면 클수록 더 정교한 그림을 만든다. 잠재 대신 픽셀을 큰 조각으로 잘라 넣어도, 신경망이 깨끗한 그림을 내놓게 하면 같은 몸통이 버틴다.

그런데 이렇게 학습한 모델에 「파란 공 위의 빨간 정육면체」을 주고 그림을 뽑아 보면, 글이 말한 것을 반쯤만 따르는 그림이 자주 나온다. 학습 손실은 글이 주어졌을 때 평균적으로 맞는 속도를 배우게 할 뿐, 그 글을 「꼭」 따르라고 하지는 않는다. FLUX.1 [dev] 의 코드를 열어 보면 시간 말고도 숫자 하나를 더 받아 시간 벡터에 더하는 입력이 있다(guidance_in). 학습이 끝난 모델이 조건을 더 세게 따르게 하려면 예측의 어디를 어떻게 바꿔야 할까? 그리고 그 숫자는 무엇을 조절하려고 들어간 것일까?