다차원 트랜스포머를 위한 축별(Axial) 어텐션
Digest: 이미지·비디오와 같은 다차원 텐서(tensor, 여러 축을 가진 배열)에 대한 autoregressive 모델링은 full self-attention의 O(N²) 복잡도 때문에 현실적으로 적용이 어려웠다. 기존 해결책인 Sparse Transformer는 희소 패턴을 설계해 계산량을 줄였지만 custom sparse kernel이 필요하고 receptive field가 제한적이었다. 본 논문은 “다차원 텐서의 축(axis)이 본래 자연스러운 구조적 분해를 제공한다”는 핵심 통찰을 바탕으로, N×N 이미지에서는 한 번은 같은 row 위의 픽셀끼리, 다음엔 같은 column 위의 픽셀끼리만 attention을 수행하는 Axial Attention을 제안한다. d차원 텐서에서 복잡도는 O(N·N^{1/d}) = O(N^{(d+1)/d})로, 이미지의 경우 O(N√N)이 되며, 여러 축별 attention을 번갈아 쌓으면 full receptive field를 유지한다. 추가로 semi-parallel decoding 스킴(upper context는 병렬로 계산, row 단위 샘플링)으로 추론 속도를 약 √N배 가속한다. 실험적으로 ImageNet-32 3.76 bpd, ImageNet-64 3.44 bpd, BAIR Robot Pushing 1.29 bits/dim으로 Sparse Transformer·Image Transformer 수준 또는 그 이상에 도달한다. 한계로는 (i) autoregressive 전용 구조라 discriminative task로의 직접 적용은 별도 설계가 필요하고, (ii) d가 커질수록 layer 수가 선형 증가한다는 점이 있다. 이 논문은 이후 MaxViT, HaloNet, CCNet 등 “attention을 축/국소 축별로 분해한다”는 계열 연구의 원형이 되었으며, “attention을 축 방향으로 separable하게 분리할 수 있는가”라는 열린 질문을 제기한다.
섹션별 요약
Introduction
- 이미지/비디오 같은 다차원 데이터에서 likelihood 기반 autoregressive 모델링은 품질이 높지만, 전역 self-attention은 길이 N에 대해 O(N²) 메모리·연산을 요구하여 고해상도 이미지에 적용하기 어렵다.
- 기존 대안(Sparse Transformer, Image Transformer의 local attention, PixelCNN의 conv)은 (a) custom kernel이 필요하거나 (b) receptive field가 제한되거나 (c) 구현이 복잡했다.
- 본 논문은 **“텐서의 각 축별로 독립적인 attention을 순차 적용”**하는 단순한 아이디어만으로 full receptive field, sub-quadratic 복잡도, 표준 프레임워크 호환성을 동시에 달성함을 보인다.
Methods
- Axial Attention: d차원 텐서 x ∈ R^{N_1×…×N_d×D}에 대해 축 k를 따라 길이 N_k 벡터들만 모아 독립적으로 self-attention 수행. 나머지 축은 batch로 취급. 축 k에 대한 복잡도는 O((∏_{j≠k} N_j) · N_k²) = O(N · N_k).
- N_1 = N_2 = … = N_d = N^{1/d}일 때 한 축 attention의 복잡도는 O(N^{(d+1)/d}). 이미지(d=2)는 O(N√N), 비디오(d=3)는 O(N·N^{1/3}).
- Axial Transformer 구조 (이미지용):
- Channel Encoder: 이전 channel 정보를 unmasked row + column axial attention으로 집약.
- Outer Decoder: masked column attention + unmasked row attention으로 “현재 row의 윗쪽 context”를 병렬 계산.
- Inner Decoder: 4블록의 masked row attention으로 현재 row를 좌→우로 autoregressive 샘플링.
- Semi-Parallel Decoding: upper context는 row를 생성하기 전에 한 번만 병렬 계산해두고, row 내부만 autoregressive로 풀기 때문에 순수 pixel-by-pixel 대비 약 √N배 빠르다.
Results
| Model | Dataset | Metric | Score | vs. Baseline |
|---|---|---|---|---|
| Axial Transformer | ImageNet-32 | bits/dim | 3.76 | Sparse Transformer 3.77 수준 |
| Axial Transformer | ImageNet-64 | bits/dim | 3.44 | Sparse Transformer 3.44 동률, PixelCNN 3.57 대비 개선 |
| Axial Transformer | BAIR Robot Pushing (video) | bits/dim | 1.29 | 이전 SOTA(Video Transformer 등) 대비 개선 |
- 세 벤치마크 모두에서 SOTA 수준 likelihood 달성, 특히 custom sparse kernel 없이 표준 dense attention만으로 달성했다는 점이 핵심.
Discussion
- 한계: (i) autoregressive density 추정 전용으로 설계되어 분류/표현학습 적용엔 추가 설계 필요. (ii) 축 수 d가 증가하면 축별 block이 선형으로 늘어나 layer budget 관리가 필요. (iii) 축별 분해는 “대각선 방향 의존성”이 단일 axial layer에서는 직접 포착되지 않고 최소 2개의 axial layer를 거쳐야 연결됨.
- 향후 방향: axial attention을 discriminative backbone으로 확장, 더 고차원(d≥3) 텐서로의 일반화, non-uniform 축 길이 처리, axial + local 혼합 패턴 탐색.
Insights
- 주목할 점: sparse pattern을 “손으로 설계”하는 대신, 텐서의 자연스러운 축 구조가 그 자체로 sparsification을 제공한다.
- 연결 고리: Separable convolution(depthwise-separable)이 공간 축을 채널과 분리했듯, axial attention은 공간 축을 서로 분리한다. 개념적 쌍대.
- 시사점: custom CUDA kernel 없이 O(N^{(d+1)/d}) 스케일을 확보 → 표준 프레임워크로도 고차원 autoregressive 가능.
- 비판적 코멘트: “축별 분해”는 축이 의미 있는 구조(이미지/비디오)에서는 강력하지만, unordered set이나 그래프 같은 데이터에는 직접 적용 불가.
Discussion Points
- 논쟁점: axial decomposition이 receptive field는 커버하지만, 한 layer 내에서 대각선 의존성을 즉시 포착할 수 없는 것이 성능 ceiling으로 작용하는가?
- 검증 필요 가정: d가 더 커질 때(예: 4D 의료영상) 축별 분해가 여전히 유효한가?
- 후속 연구: Axial-DeepLab, MaxViT, HaloNet, CSWin, CCNet 등.
메타데이터
| 항목 | 내용 |
|---|---|
| 제목 | Axial Attention in Multidimensional Transformers |
| 저자 | Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, Tim Salimans |
| 소속 | Google Brain, UC Berkeley |
| 연도 | 2019 |
| 발표 | arXiv:1912.12180 |
| 링크 | arXiv, Code |
| 키워드 | axial attention, autoregressive model, multidimensional transformer, image/video density estimation |
왜 이 연구를 하는가?
핵심 질문
다차원 텐서(이미지·비디오)에 대해 full receptive field, sub-quadratic 복잡도, 표준 프레임워크 호환성을 동시에 만족하는 self-attention을 만들 수 있는가?
기존 접근법의 한계
| 한계 | 설명 |
|---|---|
| Full self-attention O(N²) | N = H·W·C·T가 커지면 메모리/연산 폭발 |
| Sparse Transformer | 손으로 설계한 sparse pattern + custom sparse GPU kernel 필요 |
| Image Transformer (local attention) | receptive field가 local block에 제한 |
| PixelCNN | conv 기반이라 long-range 의존성 약함 |
핵심 통찰
- 다차원 텐서의 각 축 자체가 자연스러운 factorization을 제공한다: “모든 쌍”이 아니라 “같은 row 내/같은 column 내”로 attention을 제한하고, 축을 번갈아 쌓으면 global receptive field가 복원된다.
- 축별 attention은 dense attention의 표준 구현을 그대로 사용할 수 있다 — batch 축을 재배치하기만 하면 된다.
방법 (Method)
프레임워크 개요
graph TB A["입력 이미지 x ∈ R^{H×W×D}"] --> B["Channel Encoder<br/>(unmasked row + column axial attention)"] B --> C["Outer Decoder<br/>masked column attention<br/>+ unmasked row attention"] C --> D["Upper Context U<br/>(row를 생성하기 전 한 번에 병렬 계산)"] D --> E["Inner Decoder<br/>4× masked row attention"] E --> F["현재 row를 좌→우로<br/>autoregressive 샘플링"] F --> G["다음 row로 이동"] G -.-> D subgraph AxialAttn["Axial Attention (핵심 연산)"] H["축 k를 따라 길이 N_k 벡터만 선택"] --> I["나머지 축은 batch dim으로 treat"] I --> J["표준 dense self-attention 수행<br/>복잡도 O(N · N_k)"] end
핵심 흐름: row axis attention → column axis attention 을 번갈아 쌓으며, masking으로 autoregressive 순서를 보장한다. 2D 이미지에서 한 row-attention + 한 column-attention이면 모든 픽셀 쌍이 2-hop으로 연결되어 full receptive field가 확보된다.
핵심 구성요소
- Axial Attention along axis k: 텐서를 축 k에 대해 “1D 시퀀스들의 묶음”으로 보고 표준 multi-head self-attention 수행. 복잡도 O(N · N_k).
- Masking: row/column 순서에 따라 causal mask 적용. Outer decoder의 column attention은 “현재 row 위쪽”만 보도록 mask, Inner decoder의 row attention은 “현재 위치 왼쪽”만 보도록 mask.
- Semi-Parallel Decoding: upper-context(현재 row보다 위의 모든 정보)를 row 단위로 한 번만 병렬 계산 → row 내부만 픽셀 단위 autoregressive. 샘플링 비용이 naive 방식 대비 약 √N배 감소.
발견 (Findings)
주요 결과
| 모델 | ImageNet-32 (bpd) | ImageNet-64 (bpd) | BAIR Robot (bits/dim) |
|---|---|---|---|
| PixelCNN++ / Gated PixelCNN | 3.83 | 3.57 | — |
| Image Transformer | 3.77 | — | — |
| Sparse Transformer | 3.77 | 3.44 | — |
| Video Transformer | — | — | 1.35 |
| Axial Transformer | 3.76 | 3.44 | 1.29 |
핵심 발견
- 축별 분해만으로 Sparse Transformer와 동등 또는 더 나은 likelihood를 custom kernel 없이 달성.
- 비디오(d=3)로의 확장이 자연스러우며, 시간 축을 하나의 축으로 추가하는 것만으로 SOTA 갱신.
- Semi-parallel decoding으로 샘플링 속도가 √N배 개선되어 고해상도 샘플링이 현실화.
이론적 의의
축별 Separability는 다차원 attention의 기본 정리
Depthwise-separable convolution이 “공간 축 × 채널”의 분리를 통해 계산을 줄였듯, axial attention은 “공간 축 × 공간 축”의 분리를 통해 self-attention 비용을 O(N^{(d+1)/d})로 낮춘다. 축 수 d가 클수록 상대적 이득이 커져 고차원 텐서 모델링의 기본 도구가 될 수 있다.
Full receptive field의 구성적 보장
단일 axial layer는 한 축만 커버하지만, d개의 axial layer를 번갈아 쌓으면 d-hop 경로로 모든 위치가 연결된다. 즉 깊이 = 축 수 × 상수만으로 global receptive field를 수학적으로 보장한다. 이는 이후 MaxViT(block + grid), CSWin(cross-shaped) 등 “축/줄 기반 attention” 설계의 이론적 출발점이다.
Autoregressive 병렬화의 새 축
AR 모델의 샘플링은 본질적으로 직렬이지만, “조건부 독립 가정 없이” 일부 context를 공유하면 row 단위 병렬화가 가능하다는 것을 보여줌 → 이후 diffusion/parallel decoding 아이디어와 연결.
재현성 및 신뢰도 평가
| 항목 | 등급 | 비고 |
|---|---|---|
| 코드 공개 | ✅ | google-research 저장소에 공개 |
| 데이터 공개 | ✅ | ImageNet-32/64, BAIR 공개 벤치마크 |
| 하이퍼파라미터 | ⚠️ | 본문/부록에 주요 값 보고, 일부 세부 스케줄은 코드 참조 |
| 실험 환경 | ⚠️ | TPU 사용, 구체적 자원 규모 일부만 명시 |
| 통계적 신뢰도 | ⚠️ | 단일 run 위주, 표준편차 미보고 |
| 종합 등급 | B | 공개성은 좋으나 statistical rigor는 제한 |
주장별 신뢰도
| # | 주장 | 근거 | 신뢰도 |
|---|---|---|---|
| 1 | Axial attention의 복잡도는 O(N^{(d+1)/d}) | 단순 counting argument, 공식적 유도 | 🟢 |
| 2 | ImageNet-32에서 3.76 bpd, ImageNet-64에서 3.44 bpd | Table 보고, 비교군 일관 | 🟢 |
| 3 | Semi-parallel decoding이 √N배 빠름 | 복잡도 분석 기반, 실측 벽시계도 일부 제시 | 🟡 |
| 4 | Custom kernel 없이 Sparse Transformer를 능가 | ImageNet-32에서 3.76 < 3.77로 근소 우위 | 🟡 |
읽기 난이도: ⭐⭐
Self-attention 기본과 autoregressive image modeling(PixelCNN, Image Transformer) 배경이 필요. 아이디어 자체는 매우 단순하여 배경만 갖추면 빠르게 읽힘.
관련 연구 비교 매트릭스
| 축 | Axial Transformer | Sparse Transformer (Child+ 2019) | Image Transformer (Parmar+ 2018) | PixelCNN++ | MaxViT (2022, 후속) |
|---|---|---|---|---|---|
| 핵심 접근 | 축별 독립 attention | 수작업 sparse pattern (strided/fixed) | local block attention | masked causal conv | block + grid axial attention |
| 문제 정의 | 다차원 AR density | 1D 긴 시퀀스 AR | 2D 이미지 AR | 2D 이미지 AR | 2D 이미지 discriminative |
| 복잡도 | O(N^{(d+1)/d}) | O(N·√N) | O(N·w²) (local) | O(N·k²) | O(N) |
| Receptive field | Full (multi-hop) | Full (approx) | Local block | Local kernel stack | Full |
| Custom kernel | ❌ 불필요 | ✅ 필요 | ❌ | ❌ | ❌ |
| 핵심 메트릭 | 3.76/3.44 bpd, 1.29 video | 3.77/3.44 bpd | 3.77 bpd | 3.83/3.57 bpd | 분류용 (bpd 아님) |
| 한계 | AR 전용, 대각선 2-hop | 수작업 pattern, kernel | local | long-range 약함 | AR 아님 |
| 코드 공개 | ✅ | ✅ | ✅ | ✅ | ✅ |
관련 연구
- Sparse Transformer — 수작업 sparse pattern으로 O(N√N) 달성한 직접 비교 대상. Axial은 같은 복잡도를 자연스러운 축 구조로 달성.
- Image Transformer — local block attention. Axial은 full RF로 확장한 버전.
- Axial-DeepLab — Axial attention을 discriminative segmentation backbone에 적용한 대표 후속.
- MaxViT / HaloNet / CSWin Transformer — “축/줄 기반” attention의 후속 계보.
- CCNet — criss-cross attention. Axial과 실질적 쌍대.
- Attention Methods — 어텐션 방법론 계열 노트 허브.
원자적 인사이트 (Zettelkasten)
💡 축별 Separability는 attention 비용 절감의 1차 근본 원리
출처: Axial Attention in Multidimensional Transformers (Ho et al., 2019)
유형: 이론적
다차원 텐서에서 self-attention의 O(N²) 비용은 “모든 쌍을 본다”는 가정에서 온다. 그러나 텐서의 d개 축은 이미 자연스러운 factorization을 제공하므로, 축별로 attention을 분리하고 번갈아 쌓으면 O(N^{(d+1)/d})로 낮추면서 full RF를 유지할 수 있다. 이는 depthwise-separable convolution과 쌍대적이며, 이후 MaxViT, CSWin 등 “축/줄 기반” attention 설계의 공통 출발점이 된다.
핵심 조건/맥락: 데이터가 의미 있는 축 구조(이미지·비디오·의료영상)를 가질 때 성립. unordered set, 일반 graph엔 직접 적용 불가.
연결: Sparse Transformer, Depthwise Separable Convolution, MaxViT
활용 가능성: 고차원 텐서(4D MRI, 시공간 센서)에서 O(N^{(d+1)/d})만으로 full-RF attention 설계 가능.
💡 Full receptive field는 “축 수 × 상수” 깊이로 구성적으로 보장된다
출처: Axial Attention in Multidimensional Transformers (Ho et al., 2019)
유형: 이론적
단일 axial layer는 한 축만 연결하지만, d개의 축별 layer를 차례로 쌓으면 임의의 두 위치가 d-hop 안에 연결된다. 따라서 “깊이 = d × 상수”라는 아주 얕은 스택만으로도 수학적으로 global RF가 보장된다. 이 논리는 receptive field를 layer 수로 명확히 환원시키는 드문 예시다.
핵심 조건/맥락: 각 축의 attention이 unrestricted (full-axis) 해야 함. local-axis로 제한하면 추가 조건 필요(HaloNet).
연결: HaloNet, CCNet, Receptive Field 분석
활용 가능성: 새로운 attention 구조 설계 시 “몇 layer면 RF 도달?”을 깊이 상한으로 보장.
💡 Semi-Parallel Decoding — AR 샘플링에서 context 공유 병렬화
출처: Axial Attention in Multidimensional Transformers (Ho et al., 2019)
유형: 방법론적
Autoregressive 샘플링은 본질적으로 직렬이지만, “현재 row보다 위의 context”처럼 한 row 전체에 공통인 정보는 한 번만 병렬 계산해 재사용할 수 있다. 이렇게 조건부 독립 가정 없이 context 계산만 공유해도 순수 pixel-AR 대비 약 √N배 빨라진다. 이는 이후 parallel decoding / block-wise autoregression 연구와 연결된다.
핵심 조건/맥락: 모델이 “global context vs. local AR”을 분리할 수 있어야 하며, 분리 지점이 병렬 단위가 된다.
연결: Parallel Decoding, Speculative Decoding, Blockwise Parallel Decoding
활용 가능성: LLM 추론 최적화에서 “공유 가능한 prefix context”를 식별하여 재계산 비용 절감.
💡 Custom kernel 없는 sub-quadratic attention — 표준 dense 연산 재구성만으로 충분
출처: Axial Attention in Multidimensional Transformers (Ho et al., 2019)
유형: 실험적
Sparse Transformer는 O(N√N)를 위해 전용 CUDA sparse kernel을 작성했다. Axial은 동일 복잡도를 batch 축 재배치 + 표준 dense self-attention만으로 얻는다. 즉, 복잡도 이득의 원천은 “kernel 최적화”가 아니라 “연산 그래프의 reshaping”이며, 이는 “시스템 최적화 전에 수학적 factorization을 찾아라”라는 엔지니어링 교훈을 제공한다.
핵심 조건/맥락: 해당 factorization이 데이터 구조와 일치할 때만 성립.
연결: FlashAttention, Sparse Transformer
활용 가능성: 새로운 효율적 attention 제안 시 “factorization vs kernel” 두 축 중 factorization을 먼저 탐색.
💡 대각선 의존성은 axial attention의 근본적 2-hop 비용
출처: Axial Attention in Multidimensional Transformers (Ho et al., 2019)
유형: 실패/한계
Axial attention에서 (i,j)와 (i’,j’) (i≠i’, j≠j’)는 단일 axial layer 안에서 직접 연결되지 않고 반드시 중간 위치를 경유해야 한다. 즉 대각선 정보 전달은 항상 2-hop 비용이 든다. 이 제약은 이후 MaxViT의 grid attention, CSWin의 cross-shaped window처럼 “대각선을 한 layer 안에 넣는” 변형들의 설계 동기가 된다.
핵심 조건/맥락: 순수 axial 구조에 국한. dilated/grid 변형에서는 완화됨.
연결: CSWin Transformer, MaxViT
활용 가능성: 새 효율 attention 구조의 receptive-field 분석 시 “대각선 hop 수”를 점검 기준으로 사용.
핵심 용어 정리
| 용어 | 정의 |
|---|---|
| Axial Attention | 다차원 텐서의 한 축만 따라 self-attention을 수행하고, 나머지 축은 batch로 취급하는 연산 |
| Receptive Field (RF) | 한 출력 위치가 의존하는 입력 위치의 집합 |
| Autoregressive (AR) | p(x) = ∏ p(x_i |
| bits/dim (bpd) | 픽셀 당 negative log-likelihood를 bit 단위로 환산한 값 (낮을수록 좋음) |
| Masked Attention | 미래 정보를 못 보도록 softmax 전 -∞ mask를 씌운 attention |
| Sparse Transformer | Child et al. 2019, 수작업 sparse pattern으로 O(N√N) 달성 |
| Semi-Parallel Decoding | upper context는 병렬 계산, 좁은 단위(row) 내부만 AR로 순차 샘플링 |
| Separable Operation | 고차원 연산을 축별 저차원 연산의 합성으로 분해하는 기법 (예: depthwise-separable conv) |
| Axial Transformer | row/column(+time) axial attention을 쌓은 AR density 모델 |
| Channel Encoder / Outer / Inner Decoder | 이전 채널, row간, row내부 의존성을 각각 담당하는 Axial Transformer의 세 서브모듈 |
태그
paper #2019 attention axial multidimensional image-transformer autoregressive sparse-attention