APEX가 채널 계층과 비동기 디노이징으로 극한 압축 확산을 실용화한다
APEX: Asynchronous Prefix Denoising for Extreme Compression
무엇인가
잠재 확산 모델은 이미지를 더 작은 공간 격자로 압축해 어텐션 비용을 줄인다. 그런데 공간 압축을 극단까지 밀어붙이면 재구성-생성 딜레마가 나타난다. 격자가 작아질수록 재구성 품질을 지키려면 채널을 늘려야 하는데, 채널이 늘면 잠재가 이미지 고유의 텍스처와 국소 디테일을 더 많이 담게 되어 확산 모델이 생성해야 할 정보량이 커진다. 결과적으로 채널을 줄여 재구성을 포기하거나, 디테일을 살리고 생성을 포기하는 양자택일이 된다. 128배 압축에서는 512² 이미지가 4×4 격자, 즉 16토큰으로 표현되므로 이 문제가 특히 심하다.
어떻게 동작하나
APEX의 첫 번째 구성 요소는 해상도 인식 오토인코더다. 인코더가 만든 잠재 z를 채널 블록 b0부터 bL-1까지 나누고, 앞에서부터 누적한 z(ℓ)를 하나의 레벨로 정의한다. 각 레벨은 점점 높은 해상도의 타깃으로 감독된다. 학습 시 레벨을 균등하게 샘플링하고 그 뒤 블록은 0으로 채운 뒤, 디코더 출력과 원본을 각각 안티앨리어싱 리사이즈로 해당 레벨 해상도에 맞춘 다음 재구성·지각·적대적 손실을 계산한다. 이렇게 하면 좁은 초기 레벨이 고립된 고주파 디테일에 용량을 낭비하지 않고 일관된 레이아웃과 형태를 잡고, 뒤 레벨이 빠진 디테일을 복원한다.
무엇과 다른가
두 번째는 채널 그룹에 대한 비동기 디노이징이다. 전역 시계 τ 위에서 각 레벨은 로컬 시간 tℓ(τ)=clip(1+ℓ·g−τ, 0, 1)을 갖고, 레벨 간격 g만큼 늦게 시작해 항상 앞 레벨보다 덜 노이즈한 상태로 정제를 조건화한다. 각 레벨은 미시작(입력에서 0), 디노이징(자기 로컬 시간으로 노이즈), 완료(깨끗한 컨텍스트) 중 하나의 상태에 있다. 하나의 공유 Transformer가 모든 블록의 velocity를 예측하지만 디노이징 중인 레벨만 감독·갱신되고, 미시작·완료 블록의 예측은 버려진다. 등간격 스케줄에서는 스텝마다 디노이징 블록에 대한 단일 오일러 업데이트로 축약되어, 스텝당 비용이 플랫 모델과 동일하다. 손실 가중치는 디노이징 상태 밖에서 0이고 디노이징 레벨 사이에서는 상수(합이 1이 되게 정규화)라, 넓은 정제 레벨이 채널 수가 많다는 이유만으로 학습을 지배하지 않는다. 학습 시계 샘플링은 L0가 디노이징하는 구간에 확률 p_L0를 몰아주고 나머지를 균등 분배하며, DiT 입력 투영에서 b0 행은 표준 SiT 초기화, 나머지 행은 0으로 시작해 L0 전용 생성기에서 점차 정제를 배우게 한다.
어떻게 쓰나
실험은 ImageNet-512에서 64배·128배 공간 압축으로 진행됐다. f64에서 APEX의 gFID는 c128 2.85, c256 3.15, c512 4.9로 완만하게 나빠지는 반면, DC-AE는 3.22에서 9.19, 26.1로 급격히 무너진다. f128에서는 APEX가 c512 8.98, c1024 18.77을 기록했고, DC-AE는 14.44와 89.47, DC-AE 1.5는 11.86과 51.47에 그쳤다. f128c1024에서 Spectrum Matching은 73.71, FreqWarm은 79.51이었다. 16토큰 구성에서 SiT-XL로 8.98 guided gFID를 달성한 반면 비교 방법들은 모두 두 자릿수에 머문다.
전제와 한계
컴퓨트를 맞춘 비교도 있다. APEX f64c8-32-128을 DC-AE f64c128, f32c32와 같은 디노이저 아키텍처, 동일 누적 학습 FLOPs로 비교했을 때, 1배 추론 예산에서 APEX는 250 NFE로 gFID 2.74를, f32c32는 64 NFE로 2.95를 냈다. 1000 NFE(4배)에서는 APEX가 2.67까지 내려간다. 압축만으로는 품질-컴퓨트 트레이드오프가 보장되지 않지만, APEX는 모든 구간에서 f32c32를 앞서며 교차점을 없앤다.
절제 실험은 계층과 비동기 디노이징이 함께 설계돼야 한다는 점을 보여준다. 계층만 일반 확산에 얹으면 gFID가 오히려 나빠지고, 누적 프리픽스 학습으로 일부 회복되지만, 매칭되는 확산 과정을 도입할 때 17.83에서 8.12로 크게 개선된다. 해상도 매칭 타깃을 제거하면 8.77로 나빠지고, DC-AE 1.5식으로 작은 균일 폭 레벨을 다수 두고 클록 샘플러를 바꾸면 52.89까지 악화된다. 스태거 갭은 g=0.25 근처가 최적이고, 주 모델은 p_L0=0.9를 쓴다. 레벨별 기여를 보면 1·2·3번째 레벨에서 멈출 때 gFID가 각각 19.83, 9.04, 8.12로, 뒤 레벨이 직접 감독을 거의 받지 않아도 최종 품질에 필수적이다.
실무 관점에서 이 논문은 토큰 수를 줄여 추론 비용을 낮추려는 파이프라인에 직접 relevant하다. 채널을 늘려 재구성을 지키는 순간 생성이 무너지는 문제를, 토큰을 추가하지 않고 채널 축에서 해결한다는 점이 핵심이다. 스텝당 비용이 플랫 모델과 같으므로 기존 DiT 스택에 통합하기 쉽다. 다만 채널 경계, 레벨별 타깃 해상도, 스태거 갭, 전역 클록 분포를 직접 정해야 하므로, 자기 데이터와 압축률에서 이 값들을 스윕해 봐야 한다.
저자들이 밝힌 한계도 이 하이퍼파라미터에 집중된다. APEX는 채널 경계, 레벨별 타깃 해상도, 스태거 갭, 전역 클록 분포를 여전히 골라야 하며, 최적 선택은 설정마다 달라질 수 있다고 인정한다. 논문은 실용적 지침을 제시하지만, 이 선택들을 오토인코더·디노이저와 함께 공동 학습하는 것이 유망한 다음 방향이라고 남겨두었다.