DiffusionBlocks: Sakana AI의 블록 단위 학습, ICLR 2026

DiffusionBlocks는 스텝마다 잔차 블록 하나씩 학습해 활성화 메모리를 B× 절감하면서 동등하거나 더 높은 정확도를 달성한다.

DiffusionBlocks: Sakana AI의 블록 단위 학습, ICLR 2026

R 2026 채택으로 DiffusionBlocks는 훈련 효율성 문헌에서 동료 검토 지위를 확보하게 되었습니다. 연구소와 프로덕션 AI 팀의 인프라 결정은 보수적인 경향이 있으며, 5가지 아키텍처-태스크 쌍에 걸쳐 검증된 동료 검토 결과는 프리프린트보다 더 큰 신뢰를 얻습니다. MarkTechPost를 포함한 매체의 보도는 초기 실무자들의 관심을 반영합니다. 후속 연구에서는 스킵 연결 구조가 없거나 더 복잡한 비잔차 아키텍처 및 희소 혼합 전문가(MoE) 모델로 ODE 프레임워크가 확장될 수 있는지 검토할 것으로 예상됩니다.

메모리 최적화 대상 메모리 범주 DiffusionBlocks와 호환? 비고
DiffusionBlocks 활성화 메모리 (B× 감소) AdaLN 컨디셔닝 및 잔차 아키텍처 필요
그래디언트 체크포인팅 활성화 메모리 (재계산 방식) 직교 동일 메모리 범주; 절감 효과 중복, 충돌 없음
ZeRO 1~3단계 (DeepSpeed) 옵티마이저 상태, 그래디언트, 파라미터 예 — 가산적 다른 메모리 범주; 이론적으로 조합 가능
텐서 병렬처리 장치 간 파라미터 + 활성화 샤딩 다른 세분화 수준에서 작동; 알려진 충돌 없음
파이프라인 병렬처리 장치 간 레이어 분산 주의 필요 블록 경계는 파이프라인 단계 경계와 일치해야 함
CPU 활성화 오프로딩 GPU 활성화 메모리 직교 둘 다 최대 활성화 메모리 감소; 중첩 시 수익 감소

DiffusionBlocks는 ZeRO 샤딩 및 텐서 병렬처리와 직교 관계에 있는데, 각각과 서로 다른 메모리 범주를 대상으로 하기 때문입니다. ZeRO 1~3단계는 데이터 병렬 랭크 전반에 걸쳐 옵티마이저 상태, 그래디언트, 파라미터를 줄입니다. 텐서 병렬처리는 장치 간에 가중치 행렬을 샤딩합니다. DiffusionBlocks는 반복당 하나의 블록만 실행하여 스텝당 활성화 메모리를 줄입니다. 세 가지를 모두 결합하면 이론적으로 가산적입니다 — 각각이 전체 메모리 예산의 서로 다른 구성 요소를 다룹니다. 논문에서는 이를 직교적으로 다루지만 조합된 구성을 직접 테스트하지는 않습니다.

"로컬 스코어 매칭 목적함수가 분류, 생성, 언어 모델링 아키텍처 전반에 걸쳐 전역 역전파에 필적할 수 있다면, 메모리 효율적인 훈련의 설계 공간은 이 분야가 가정해 온 것보다 훨씬 넓습니다." — Makoto Shing, Masanori Koyama, Takuya Akiba (source: Sakana AI, DiffusionBlocks)

가장 즉각적인 적용 대상은 규모 확장 시 활성화 메모리가 제약 요인인, 깊이가 깊은 잔차 트랜스포머 또는 이미지 디퓨전 모델을 운영하는 팀입니다. 활성화가 제약 요인이 아닌 경우 — 예를 들어 대규모 모델의 옵티마이저 상태가 장치 메모리 대부분을 소비하는 경우 — DiffusionBlocks만으로는 해당 구성을 해결할 수 없습니다. 어떤 구성 요소가 병목인지 파악하기 위한 메모리 프로파일링이 도입 결정에 앞서 이루어져야 합니다.

더 넓은 시사점은 아키텍처에 있습니다. 잔차-디퓨전 대응은 정밀한 수학적 결과입니다: 이것이 다른 아키텍처 계열로 확장될 수 있다면 — ODE 관점은 충분히 일반적이어서 확장이 타당합니다 — 로컬 학습 가능한 심층 신경망의 설계 공간은 이 분야가 가정해 온 것보다 넓습니다. Forward-Forward 알고리즘은 이론적 근거 없이 이 방향을 제시했으며, DiffusionBlocks는 이 접근 방식을 원칙적으로 만드는 유도를 제공합니다. 그 유도가 MoE 레이어, 상태 공간 모델, 또는 아직 설계되지 않은 아키텍처로 일반화될 수 있는지 여부가 후속 연구의 핵심 미해결 질문입니다.

자주 묻는 질문

DiffusionBlocks: Sakana AI's Block-Wise Training for ICLR 2026

DiffusionBlocks는 실제로 GPU 메모리를 얼마나 절약하나요?

DiffusionBlocks는 활성화 메모리를 전체 네트워크 깊이 L에 비례하는 O(L)에서 O(L/B)로 줄입니다. 여기서 B는 학습 블록의 수입니다. B=4일 때, 학습 스텝당 활성화 메모리가 4배 감소합니다. 핵심 제약은 옵티마이저 상태(Adam의 1차 및 2차 모멘트 추정치)와 파라미터 메모리는 영향을 받지 않는다는 점입니다. 총 피크 메모리 절감량은 특정 아키텍처와 배치 크기에서 피크 메모리 중 활성화가 차지하는 비율에 따라 달라집니다. 활성화가 병목인 구성 — 대형 배치 크기에서 폭이 넓고 깊은 트랜스포머 — 은 이론적 최대치인 B배에 가까워집니다. 파라미터가 병목인 구성은 총 절감량이 상대적으로 작습니다.

블록 하나씩 학습하면 항상 종단간 학습과 동일한 품질로 수렴하나요?

대부분의 보고된 벤치마크에서는 그렇습니다 — CIFAR-10의 DiT 이미지 생성(FID 30.59 대 32.84 E2E), ImageNet 256(FID 9.00 대 9.01), OpenWebText의 자기회귀 언어 모델링(MAUVE 0.71 대 0.50)에서 DiffusionBlocks는 종단간 역전파를 능가합니다. 예외는 CIFAR-100의 ViT 분류로, DiffusionBlocks는 59.30%를 달성한 반면 E2E는 60.25%를 기록해 0.95 퍼센트포인트 차이가 납니다. 품질 민감도는 B 값이 클수록 증가하며, 최적의 B는 작업 및 아키텍처에 따라 다르므로 절제 실험이 필요합니다. 매우 높은 B 값에서 E2E와의 동등성이 유지된다는 보장은 없으며, 논문은 최적 B 선택을 경험적 문제로 다룹니다.

현재 DiffusionBlocks를 사용할 수 있는 아키텍처는 무엇인가요?

잔차 연결(residual connection) 기반의 모든 아키텍처에 적용 가능합니다: ResNet, Vision Transformer(ViT), Diffusion Transformer(DiT), 마스크 확산 언어 모델, 표준 잔차 연결을 갖춘 자기회귀 트랜스포머. 이론적 유도는 스킵 연결이 연속 시간 ODE의 오일러 스텝을 구현한다는 것에 의존하며 — 이 구조적 특성이 확산 유사성과 블록별 스코어 매칭 목적함수를 가능하게 합니다. 비잔차 아키텍처 — 스킵 연결 없는 순수 어텐션, Mamba 스타일 상태 공간 모델, 기타 재귀 구조 — 는 현재 이론적 범위 밖에 있으며, DiffusionBlocks를 적용하려면 프레임워크를 확장하거나 수정해야 합니다.

등확률 분할이란 무엇이며 왜 중요한가요?

등확률 분할은 각 블록에 노이즈 레벨 범위를 할당할 때, 원시 노이즈 구간의 1/B가 아니라 로그 정규 노이즈 분포 아래 전체 확률 질량의 1/B를 각 블록이 처리하도록 합니다. 로그 정규 분포에서 확률 질량은 스코어 추정이 가장 어려운 중간 노이즈 레벨에 집중됩니다. 균일한 구간 분할은 이러한 어려운 레벨에 블록 용량을 충분히 할당하지 못합니다. 등확률 분할은 각 블록에 확률 가중 학습 신호를 균등하게 분배하여, 가장 중요한 중간 범위에 더 넓은 절대 노이즈 범위를 집중시킵니다. 실질적인 효과는 모든 B개 블록에 걸쳐 학습 난이도가 더 균일하게 분포되어 전체 학습 신호 품질이 향상된다는 것입니다.

DiffusionBlocks를 그래디언트 체크포인팅이나 ZeRO와 함께 사용할 수 있나요?

이론적으로는 가능합니다 — 각각 서로 다른 메모리 범주를 대상으로 하며 직교적으로 설계되어 있습니다. DiffusionBlocks는 스텝당 하나의 블록만 실행하여 활성화 메모리를 줄입니다. ZeRO(1~3단계)는 데이터 병렬 랭크 전반에 걸쳐 옵티마이저 상태, 그래디언트, 파라미터를 줄입니다. 그래디언트 체크포인팅은 활성화를 저장하는 대신 역방향 패스에서 재계산함으로써 연산과 활성화 메모리를 교환합니다. 논문은 이들이 직교적이라고 설명하지만 결합 구성을 직접 테스트하지는 않습니다. 기법을 조합할 계획인 팀은 가산적 절감을 가정하기 전에, 특정 아키텍처와 배치 구성에서 상호작용이 올바르게 동작하는지 경험적으로 검증해야 합니다.

DiffusionBlocks가 학습 인프라에 가져오는 변화

DiffusionBlocks는 블록별 학습 연구에서 오랫동안 존재해온 이론적 공백을 메웁니다. 잔차-확산 대응 관계는 휴리스틱이 아닌 정밀한 수학적 결과이며, 다섯 가지 아키텍처-작업 쌍에 걸친 경험적 범위는 여러 도메인에 걸쳐 신뢰할 수 있을 만큼 폭넓습니다. 활성화 메모리가 깊은 잔차 모델의 핵심 제약인 팀에게 DiffusionBlocks는 그래디언트 체크포인팅의 구체적인 대안으로, 연산 오버헤드가 적습니다 — 대신 AdaLN 컨디셔닝 수정과 적절한 B를 찾기 위한 절제 실험이 필요합니다.

제약도 분명하며 도입 결정 시 고려해야 합니다. 옵티마이저 상태 메모리는 변하지 않으므로, 활성화가 실제 병목일 때만 완전한 절감 효과가 나타납니다. 체크포인트 호환성을 위해서는 학습 시작 전 아키텍처 수정이 필요하며, 기존 체크포인트에서 파인튜닝하는 경로는 설명되어 있지 않습니다. ViT 분류 성능 차이는 작지만 프로덕션 규모에서는 문제가 될 수 있습니다. 비잔차 아키텍처는 적용 범위 밖입니다. 현재 참조 구현은 주로 CIFAR-100 ViT 예제를 다루며, 논문의 DiT 또는 언어 모델링 결과를 재현해야 하는 팀은 일정을 계획하기 전에 구현 가용성을 확인해야 합니다.

더 큰 의미는 아키텍처적 차원에 있습니다. 분류, 생성, 언어 모델링 전반에서 로컬 스코어 매칭 목적함수가 글로벌 역전파와 동등하다는 결과는 학습 효율성 분야가 발전시켜 나갈 토대가 됩니다. ODE 프레임워크가 잔차 아키텍처를 넘어 — MoE 레이어, 상태 공간 모델, 아직 설계되지 않은 아키텍처로 — 일반화될 수 있는지가 앞으로 주목해야 할 질문입니다. DiffusionBlocks는 그 탐구를 위한 원칙적인 출발점을 제공합니다.

최종 업데이트: 2026-05-30. arXiv:2506.14202(ICLR 2026 채택)의 DiffusionBlocks 논문과 Sakana AI 공식 저장소를 기반으로, 2026년 5월 말 기준으로 검토되었습니다.

최신 소식 받기

AI 도구, 에이전트, 그리고 이들을 잇는 프로토콜에 대한 현장 기록.

Creeta 둘러보기