10장 — 어텐션과 트랜스포머: 멀리 떨어진 칸끼리 읽기
크로스 어텐션: 그림 칸이 글 토큰을 읽는다
지금까지는 한 줄 안의 토큰들이 서로를 읽었다. 그런데 글로 조건을 거는 모델에서는 읽어야 할 것이 다른 줄에 있다. Stable Diffusion 1.x 의 글 인코더는 프롬프트를 토큰 77개, 토큰마다 수 768개의 표로 넘긴다. 정육면체가 그려질 칸은 「red」와 「cube」를, 공이 그려질 칸은 「blue」와 「sphere」를 많이 읽어야 한다. 두 줄은 길이도 다르고(4096칸 대 77토큰) 폭도 다르다(320 대 768). 셀프 어텐션의 계산을 어떻게 고쳐야 그림 칸이 글 토큰을 골라 읽을 수 있을까?
역사: 번역기의 디코더가 원문을 읽던 자리
사실 처음 나온 어텐션이 이 꼴이었다. 바다나우와 동료들의 번역기는 번역문을 쓰는 디코더가 원문을 읽은 인코더의 낱말 벡터들을 골라 읽었다. 바스와니와 동료들의 트랜스포머도 원문 쪽 블록 더미(인코더)와 번역문 쪽 블록 더미(디코더)로 나뉘어, 디코더 블록마다 「질의는 디코더의 앞 층에서, 열쇠와 값은 인코더의 출력에서」 오는 어텐션을 하나씩 두었다. 논문은 이것을 인코더–디코더 어텐션이라 불렀다.
2022년 롬바흐(Robin Rombach)와 동료들은 잠재 디퓨전 모델(Stable Diffusion 의 바탕)에서 이 자리를 U-Net 안에 만들었다. 글이든 분류 라벨이든 조건을 따로 둔 인코더로 토큰 줄로 바꾼 뒤, U-Net 중간 층의 특징에서 질의를, 조건의 토큰 줄에서 열쇠와 값을 만들어 읽게 한 것이다. 논문 초록은 크로스 어텐션 층을 넣어 디퓨전 모델을 「글이나 경계 상자 같은 일반적인 조건」을 받는 생성기로 바꿨다고 적었다.
두 칸이 네 낱말을 읽기
장난감으로 보자. 그림 칸이 둘(정육면체가 그려질 위 칸, 공이 그려질 아래 칸), 글 토큰이 넷(red, cube, blue, sphere)이다. 글 인코더를 지난 토큰은 앞뒤 문맥을 담고 있다고 하고, 열쇠의 세 수를 「위쪽 물건에 관한 말인가, 아래쪽 물건에 관한 말인가, 색 낱말인가」로 두자. red는 (1, 0, 1), cube는 (1, 0, 0), blue는 (0, 1, 1), sphere는 (0, 1, 0)이다. 위 칸의 질의는 (2, 0, 1), 아래 칸은 (0, 2, 1)이라 하자. 점수를 √3으로 나눠 소프트맥스를 씌우면 위 칸은 red, cube, blue, sphere에 0.487, 0.273, 0.153, 0.086을 주고, 아래 칸은 0.153, 0.086, 0.487, 0.273을 준다. 같은 글 토큰 줄을 읽지만, 칸마다 질의가 달라 서로 다른 낱말에 무게를 둔다.
달라진 것은 질의를 만드는 줄과 열쇠·값을 만드는 줄이 다르다는 것 하나다.
비중 표는 칸 × 토큰, 곧 4096 × 77이고 줄마다 소프트맥스를 한다. 글 토큰은 읽히기만 하고 이 계산으로 바뀌지 않는다. 이렇게 질의는 한 줄에서, 열쇠와 값은 다른 줄에서 만들어 한쪽이 다른 쪽을 골라 읽는 어텐션을 크로스 어텐션 (다른 줄을 골라 읽기 / cross-attention)이라 한다.
Stable Diffusion 1.x 의 어텐션 블록은 트랜스포머 블록(셀프 어텐션과 MLP 를 층 정규화·잔차로 감싼 한 벌)의 두 부품 사이에 이 부품을 하나 끼운 꼴이다.
flowchart LR X["U-Net 특징 지도<br/>64 × 64 × 320"] --> G["그룹 정규화<br/>1 × 1 합성곱"] G --> S["층 정규화 → 셀프 어텐션<br/>칸끼리 읽기 → 더하기"] S --> C["층 정규화 → 크로스 어텐션<br/>글 토큰 줄 c를 읽기 → 더하기"] T["글 인코더<br/>77 × 768"] --> C C --> F["층 정규화 → MLP<br/>칸마다 곱씹기 → 더하기"] F --> O["1 × 1 합성곱<br/>+ 블록 입력"]
ML에서: 글 쪽 폭이 들어가는 자리는 두 행렬뿐
ComfyUI 의 CrossAttention 은 질의 행렬을 그림 쪽 폭에서, 열쇠·값 행렬을 context_dim(Stable Diffusion 1.x 는 768)에서 만든다. 글 인코더를 바꿔 글 쪽 폭이 달라지면 U-Net 에서 크기가 바뀌는 행렬은 크로스 어텐션의 열쇠·값 행렬뿐이다. 그래서 같은 U-Net 뼈대가 SD 2 에서는 1024, SDXL 에서는 2048 폭의 글 토큰을 받는다. 어텐션 블록마다 자기 열쇠·값 행렬을 따로 가지므로, 같은 글 토큰 줄이라도 블록마다 다른 열쇠와 값으로 읽힌다. 64 × 64 단계의 블록은 세밀한 결에 필요한 낱말을, 16 × 16 단계의 블록은 큰 배치에 필요한 낱말을 읽도록 따로 배울 수 있다.
문제 12. 글 토큰이 하나뿐이면
크로스 어텐션에 들어오는 글 토큰 줄이 토큰 하나 c1뿐이라고 하자. (가) 그림 칸마다의 비중은 얼마인가? (나) 칸마다 읽어 가는 o는 어떻게 되는가? 질의 행렬 WQ는 무슨 일을 하는가? (다) 토큰이 77개라도 모든 토큰이 같은 벡터라면? (라) 이것은 시간 벡터를 채널마다 더하는 것과 어떻게 닮았는가?
토큰이 하나면 점수가 하나고, 소프트맥스를 하면 무조건 1이에요. (가)는 1.
(나)는 칸마다 비중 1로 같은 값 c1WV를 읽으니까, 모든 칸이 똑같은 벡터를 받아. 질의 행렬은… 아무 일도 안 하네. 점수가 얼마든 비중은 1이니까.
(다)는요? 77개인데도 같아요?
77개면 비중이 나뉘니까 다르지 않을까요?
열쇠가 다 같으면 점수도 다 같아서 비중이 77분의 1씩이야. 그런데 값도 다 같으니까 섞어 봤자 같은 벡터 하나지. 결국 (나)와 똑같아. 토큰이 몇 개든 서로 다른 것이 없으면 「어느 칸이 무엇을 읽나」가 사라져.
그럼 (라)는요?
모든 칸에 같은 벡터를 더하는 거니까, 시간 벡터를 채널 수로 바꿔 모든 칸에 똑같이 더한 것과 같은 꼴이에요. 크로스 어텐션이 칸마다 다르게 읽을 수 있는 건 글 토큰이 여럿이고 서로 다를 때뿐이에요. 문장을 벡터 하나로 줄여서 넘기면 크로스 어텐션도 결국 「그림 전체를 한꺼번에 미는」 길이 되네요.
그래요. 글 인코더가 토큰마다의 벡터 줄을 넘기는 까닭이 여기 있어요.
조별 과제 피드백을 「다들 수고했어요」 한 줄로만 받으면, 누가 무엇을 고쳐야 하는지 알 수 없는 거랑 같네요.
문제 13. 어텐션 블록의 매개변수는 어디에 쓰이나
ComfyUI 로 Stable Diffusion 1.x 의 어텐션 블록을 만들어 세면, 채널 320인 블록 하나는 2,546,240개다. 그 가운데 셀프 어텐션 409,920, 크로스 어텐션 696,640(그중 열쇠·값 행렬 491,520), MLP 1,231,680, 층 정규화 셋 1,920, 들고 나는 층(그룹 정규화와 1 × 1 합성곱 둘) 206,080이다. 어텐션 블록 16개를 모두 합하면 267,239,360개이고, 셀프 어텐션 49,574,080, 크로스 어텐션 43,962,560, MLP 148,797,120이다. (가) 채널 320 블록에서 크로스 어텐션이 셀프 어텐션보다 큰 까닭은? (나) 16개를 합친 몫(%)은? 「어텐션 블록」의 매개변수는 주로 어텐션에 쓰이는가? (다) 글 인코더를 폭 1024짜리로 바꾸면 열쇠·값 행렬은 모두 몇 개가 되는가(합친 열쇠·값 행렬은 지금 19,169,280개)?
(가)는 크로스 어텐션의 열쇠·값이 768에서 오니까요. 셀프 어텐션은 네 행렬이 다 320 × 320이라 409,600에 치우침 320이고, 크로스는 열쇠·값이 768 × 320씩 491,520, 거기에 질의와 출력 행렬 205,120을 더하면 696,640이에요.
(나)는 MLP 가 55.7%, 셀프 어텐션 18.6%, 크로스 어텐션 16.5%야. 이름은 어텐션 블록인데 절반 넘게 MLP 네.
그런데 1280 채널 블록에서는 순서가 바뀌어요. 셀프 어텐션 6,554,880, 크로스 어텐션 5,244,160이에요. 왜일까요?
셀프 어텐션은 네 행렬이 다 1280 × 1280이라 폭의 제곱으로 크는데, 크로스의 열쇠·값은 768 × 1280이라 폭에 한 번만 비례해요. 폭이 768보다 커지면 셀프 쪽이 더 커지죠.
(다)는 블록마다 2 × 폭 × 1024를 다시 세면 돼요. 320이 다섯, 640이 다섯, 1280이 여섯이니까 25,559,040개예요. 640만 개 남짓 늘어요.
글 인코더를 바꾸면 이 행렬들은 크기부터 달라지니 그대로 쓸 수 없어요. 그래서 글 인코더를 바꾼 모델은 적어도 이 행렬들을, 대개는 U-Net 전체를 새로 학습해요. 이 부분이 2,500만 개쯤으로 U-Net 전체 8억 6,000만 개의 3%가 안 되는데도요.
콘센트 모양이 다른 나라에서 가전제품을 쓰려면 플러그만 바꾸면 되는 줄 알았는데, 전압이 달라서 속까지 손봐야 하는 거네요.