하나의 모델, 다양한 예산: 확산 트랜스포머를 위한 탄력적 잠재 인터페이스 (ELIT)

기존 Diffusion Transformer는 이미지 크기에 따라 연산량이 고정되어 효율성이 떨어졌으나, ELIT는 이미지의 중요 부위에만 연산을 집중하고 추론 시점에 연산량을 자유롭게 조절할 수 있게 한다. 이를 통해 동일한 품질을 유지하면서도 하드웨어 제약이나 지연 시간 요구사항에 맞춰 유연하게 대응할 수 있다.

왜 중요한가

핵심 기여

Elastic Latent Interface (ELIT) 제안

입력 이미지 크기와 연산량을 분리하여 추론 시점에 토큰 수를 조절함으로써 유연한 품질-비용 트레이드오프를 제공하는 메커니즘을 도입했다.

Read/Write Cross-Attention 구조 도입

공간적 토큰에서 중요한 정보를 추출하는 Read 레이어와 처리된 정보를 다시 공간 도메인으로 복원하는 Write 레이어를 통해 효율적인 정보 흐름을 구현했다.

Importance-Ordered Representation 학습

학습 시 무작위로 후미 토큰을 누락시키는 Tail Dropping 방식을 통해 앞쪽 토큰은 전역 구조를, 뒤쪽 토큰은 세부 사항을 담당하도록 계층적 표현력을 학습시켰다.

Cheap Classifier-Free Guidance (CCFG) 구현

연산량이 적은 하위 예산 모델을 가이드로 활용하여 품질 저하 없이 전체 추론 비용을 약 33% 절감하는 기법을 적용했다.

핵심 아이디어 이해하기

Transformer 기반의 DiT는 이미지를 패치 단위로 나누어 처리하는데, 모든 패치에 동일한 연산량을 할당한다. 이는 하늘이나 배경처럼 단순한 영역과 얼굴처럼 복잡한 영역을 구분하지 못해 자원을 낭비하는 결과를 초래한다. ELIT는 '잠재 인터페이스(Latent Interface)'라는 가변 길이의 토큰 세트를 도입하여 이 문제를 해결한다.

Read 레이어가 Attention 연산을 통해 이미지의 복잡하고 중요한 영역에서 더 많은 정보를 가져와 잠재 토큰에 담고, 핵심 연산은 이 압축된 잠재 토큰들 사이에서만 수행된다. 이 과정은 마치 고해상도 사진에서 중요한 피사체만 골라내어 정밀하게 보정하는 것과 유사한 원리로 작동한다.

특히 학습 과정에서 'Tail Dropping' 기법을 사용해 잠재 토큰의 순서에 따라 중요도를 부여한다. 결과적으로 추론 시 사용자가 토큰 개수를 줄이면 연산 속도가 빨라지고, 늘리면 세부 묘사가 정교해지는 유연한 구조를 갖게 된다. 이는 하나의 모델이 저사양 기기부터 고사양 서버까지 다양한 환경에 맞춰 스스로를 최적화할 수 있음을 의미한다.

방법론

전체 아키텍처는 Spatial Head, Latent Core, Spatial Tail의 세 부분으로 나뉜다. Spatial Head는 입력 패치를 정제된 공간 표현으로 변환하고, Latent Core에서 대부분의 Transformer 연산이 이루어지며, Spatial Tail은 최종 출력을 생성한다.

Read 레이어는 [잠재 토큰 l과 공간 토큰 s를 입력으로] → [Cross-Attention 연산을 통해 s의 정보를 l로 병합하여] → [업데이트된 잠재 토큰 l_O를 얻고] → [이 값이 공간상의 핵심 정보를 압축한 인터페이스 역할을 수행]하는 방식으로 작동한다. 이 과정에서 모델은 손실값이 높은 영역에 더 많은 주의를 기울이도록 학습된다.

Grouped Cross-Attention을 적용하여 연산 효율성을 극대화한다. [공간 토큰을 G개의 그룹으로 나누고 잠재 토큰도 이에 맞춰 분할하여] → [그룹 내에서만 Cross-Attention 연산을 수행해] → [연산 복잡도를 O(NK)에서 O(NK/G)로 감소시키고] → [해상도 증가에 따른 연산량 증가를 억제]한다.

주요 결과

ImageNet-1K 512px 벤치마크에서 기존 DiT 대비 FID 점수를 35.3%, FDD 점수를 39.6% 개선했다. 모델 크기가 커질수록 성능 향상 폭이 더 크게 나타나 대규모 모델에 적합함이 확인됐다.

학습 수렴 속도 측면에서 256px 해상도에서는 3.3배, 512px에서는 4.0배의 가속화를 달성했다. 이는 ELIT가 정보 밀도가 높은 영역에 연산을 집중함으로써 더 효율적으로 학습하기 때문이다.

추론 시 잠재 토큰 수를 4096개에서 256개까지 조절하며 실험한 결과, 토큰 수를 줄여도 이미지의 전체적인 구조는 유지하면서 세부 표현만 단계적으로 변화하는 안정적인 품질-비용 트레이드오프를 보여주었다. 또한 CCFG 기법을 통해 품질 저하 없이 추론 비용을 33% 절감했다.

기술 상세

ELIT는 DiT, U-ViT, HDiT 등 다양한 확산 모델 백본에 즉시 적용 가능한 모듈형 설계를 채택했다. 기존 모델의 핵심 Transformer 블록 구조를 유지하면서 Read/Write 레이어만 추가하여 구현 복잡도를 최소화했다.

핵심 메커니즘인 'Importance-ordered Latent Interface'는 Nested Dropout 개념을 응용한 Tail Dropping 학습 전략을 통해 구현된다. 학습 시 매 스텝마다 무작위로 J개의 토큰만 남기고 나머지를 제거함으로써, 모델이 앞쪽 토큰에 가장 중요한 정보를 먼저 채우도록 강제한다.

Autoguidance를 활용한 CCFG(Cheap Classifier-Free Guidance)는 ELIT의 가변 예산 특성을 극대화한다. 무조건부 생성 시에는 아주 적은 수의 토큰(약 35% 비용)만 사용하는 하위 예산 모델을 가이드로 사용하여, 품질은 유지하면서 전체 FLOPs를 33% 절감한다.

구현 상세에서는 QK Normalization과 RoPE(Rotary Position Embedding)를 통합하여 대규모 학습의 안정성을 높였다. 또한 MLP 블록에서 차원 확장을 제거하여 추가적인 연산 오버헤드를 방지했다.

한계점

대규모 데이터셋을 활용한 처음부터의 학습(From-scratch training) 효과는 아직 검증되지 않았으며, 현재는 주로 기존 모델을 미세 조정하거나 증류하는 방식에 집중되어 있다.

실무 활용

모바일 기기부터 고성능 서버까지 다양한 하드웨어 환경에서 하나의 모델로 최적의 성능을 낼 수 있게 해준다. 실시간성이 중요한 서비스에서는 연산량을 줄이고, 고품질 결과가 필요한 경우에는 자원을 더 투입하는 식의 유연한 운영이 가능하다.

모바일 기기의 배터리 상태나 발열 정도에 따른 실시간 이미지 생성 품질 조절
웹 서비스에서 사용자 요금제 등급에 따른 차등화된 생성 속도 및 품질 제공
비디오 생성 시 프레임별 중요도에 따른 가변 연산 할당으로 전체 렌더링 시간 단축

코드 공개 여부: 공개

코드 저장소 보기

키워드

DiT(확산 트랜스포머)Elastic-Compute(탄력적 연산)Latent-Interface(잠재 인터페이스)Adaptive-Inference(적응형 추론)

용어 해설

확산 트랜스포머(DiT): — 이미지 생성 모델인 Diffusion Model의 구조를 기존의 U-Net 대신 Transformer 아키텍처로 설계한 모델이다. 패치 단위의 처리를 통해 확장성이 뛰어나지만, 이미지 해상도에 따라 연산량이 고정되는 한계가 있다.
부동 소수점 연산량(FLOPs): — 모델이 한 번의 추론이나 학습을 수행할 때 필요한 수학적 연산의 횟수를 의미한다. AI 모델의 효율성과 하드웨어 요구사항을 측정하는 핵심 지표로 사용된다.
분류기 없는 가이드(Classifier-Free Guidance): — 확산 모델에서 텍스트 조건부 생성과 무조건부 생성을 동시에 수행하여 결과물의 품질을 높이는 기법이다. 품질은 향상되지만 모델을 두 번 실행해야 하므로 연산 비용이 두 배로 증가하는 단점이 있다.

하나의 모델, 다양한 예산: 확산 트랜스포머를 위한 탄력적 잠재 인터페이스 (ELIT)

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드