대규모 모델 학습에서 점점 더 큰 병목이 되는 것은 연산량만이 아닙니다. 모델이 깊어질수록 역전파에 필요한 중간 활성값을 저장해야 하고, 이 과정에서 메모리 요구량이 빠르게 커집니다.
Sakana AI가 소개한 DiffusionBlocks는 이 문제를 “전체 네트워크를 한꺼번에 학습해야 한다”는 전제에서 다시 바라봅니다. 논문은 arXiv:2506.14202로 공개됐고, 코드는 GitHub의 SakanaAI/DiffusionBlocks 저장소에 올라와 있습니다.12
DiffusionBlocks는 무엇을 바꾸나
DiffusionBlocks의 기본 아이디어는 신경망을 여러 블록으로 나누고, 한 번에 한 블록씩 학습하는 것입니다. 표면적으로는 단순한 블록 단위 학습처럼 보이지만, 핵심은 블록을 나누는 이유와 학습 목표를 확산 모델의 관점에서 설명한다는 데 있습니다.
논문은 Transformer 계열 네트워크의 잔차 연결을 확산 모델의 잡음 제거 과정으로 해석합니다. 각 블록은 전체 문제를 한꺼번에 해결하는 대신, 특정 잡음 구간을 맡아 학습합니다.1
이렇게 하면 전체 네트워크에 대한 그래디언트와 활성값을 동시에 저장해야 하는 부담을 줄일 수 있습니다. 논문은 블록 수에 비례해 메모리 요구량을 낮출 수 있다고 설명합니다.1
왜 지금 이 접근이 중요한가
블록 단위 학습 자체가 완전히 새로운 발상은 아닙니다. 문제는 기존 접근들이 블록별 목표를 어떻게 정당화할지, 그리고 특정 작업을 넘어 얼마나 일반화될 수 있을지에서 한계를 보였다는 점입니다.
DiffusionBlocks가 흥미로운 이유는 여기 있습니다. 단순히 “메모리를 아끼기 위해 네트워크를 잘라서 학습하자”가 아니라, 블록별 학습이 왜 성립할 수 있는지를 확산 모델의 score matching 관점으로 연결합니다. 즉, 메모리 효율화 기법이면서 동시에 학습 방식을 다시 설계하는 시도입니다.
이 흐름은 최근 AI 인프라 연구의 방향과도 맞닿아 있습니다. Hugging Face의 Infrastructure 컬렉션에는 메모리 효율화와 학습·추론 인프라 관련 논문들이 포함돼 있습니다.3 예를 들어 2024년 5월 21일 공개된 Reducing Transformer Key-Value Cache Size with Cross-Layer Attention는 Transformer의 KV 캐시 크기 축소를 다루고, 2024년 5월 28일 공개된 VeLoRA: Memory Efficient Training using Rank-1 Sub-Token Projections는 메모리 효율적 학습을 다룹니다.3
DiffusionBlocks도 같은 질문에 답하려는 연구로 볼 수 있습니다. 더 큰 모델을 만들기 위해 단순히 더 많은 GPU를 쓰는 대신, 학습 과정 자체를 더 메모리 효율적으로 바꿀 수 있는지를 묻는 것입니다.
적용 가능성을 볼 때 확인할 점
DiffusionBlocks를 평가할 때 가장 먼저 봐야 할 지점은 범용성입니다. Sakana AI는 DiffusionBlocks를 ViT, DiT, masked diffusion, autoregressive transformers, recurrent-depth transformers에서 검증했다고 밝혔습니다.1
다만 이를 곧바로 모든 모델에 적용 가능한 만능 최적화로 보기는 어렵습니다. 실제 적용에서는 몇 가지 판단이 필요합니다.
네트워크를 어떤 기준으로 블록화할 것인가
각 블록에 잡음 구간을 어떻게 배정할 것인가
기존 학습 파이프라인을 얼마나 수정해야 하는가
메모리 절감 효과가 성능, 학습 안정성, 구현 복잡도와 어떤 균형을 이루는가
따라서 DiffusionBlocks의 가치는 “바로 가져다 쓰는 트릭”이라기보다, 대규모 모델 학습을 설계하는 새로운 관점에 가깝습니다.
배울 점
DiffusionBlocks가 던지는 메시지는 분명합니다. 앞으로의 모델 확장은 단순히 더 깊고 큰 네트워크를 만드는 문제만이 아닙니다. 그 네트워크를 어떻게 나누고, 어떤 단위로 학습하며, 메모리 부담을 어디서 줄일 것인지가 점점 더 중요해집니다.
DiffusionBlocks는 “전체를 한 번에 학습하지 않는 방법”을 확산 모델의 언어로 정리한 사례입니다. 대규모 학습 비용을 줄이려는 연구자나 인프라 설계자라면, 이 접근을 메모리 최적화 기법이자 학습 목표 재구성의 사례로 함께 살펴볼 만합니다.
1DiffusionBlocks: Block-wise Neural Network Training via Diffusion Interpretation
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.