HiRAE가 깊이별 잔차 예산으로 계층 융합 토크나이저를 학습한다

무엇인가

사전학습된 비전 인코더의 표현은 이미지 생성의 기반이 되지만, 마지막 층의 특징만 쓰면 충실한 재구성에 필요한 미세 디테일이 빠진다. 중간 층에는 보완적인 시각 정보가 남아 있지만, 이를 재구성용으로 융합하도록 학습하면 생성 모델이 모델링하기 어려운 잠재 분포가 나올 수 있다. 기존 방식은 층 선택을 경험적으로 튜닝하거나(RAEv2의 고정 집계), 융합과 디코딩을 단계적으로 최적화해야 했다(DRoRAE). 이 논문은 층 부분집합을 수동으로 고르지 않고 인코더 전체 계층을 한 번에 학습하는 HiRAE를 제안한다.

어떻게 동작하나

HiRAE-24는 동결된 DINOv3-L의 24개 층을 모두 사용한다. 각 층의 특징 H_ℓ(256개 토큰, 1024 채널)은 토큰별 MLP 전문가 E_ℓ를 거쳐 U_ℓ로 변환되고, 가장 깊은 층 H_23의 선형 투영으로 만든 라우팅 점수를 L2 정규화(부호 유지)해 공간 위치마다 층별 가중치를 얻는다. 어느 층의 정보를 어느 위치에서 얼마나 쓸지를 데이터가 학습하는 구조다.

무엇과 다른가

핵심은 라우팅 가중치가 아니라 그 결과 생기는 보정량을 깊이별로 묶어 제한하는 데 있다. 24개 층을 얕은 그룹(0~7), 중간(8~15), 깊은(16~23) 세 그룹으로 나누고, 각 그룹의 가중 합 R_g를 잔차 제어 모듈 C_g에 통과시켜 보정 Δ_g를 만든 뒤 가장 깊은 특징 H_23에 더하고 층 정규화를 적용해 융합 잠재 Z를 얻는다. 학습 중 잔차 드롭아웃 확률은 얕은 그룹 0.50, 중간 0.25, 깊은 0.10이고, 프로베니우스 노름 상한은 각각 H_23 노름의 0.025배, 0.075배, 0.150배다. 삼각부등식에 따라 세 그룹 보정의 총합은 H_23 노름의 0.250배를 넘지 않는다. 얕은 층일수록 예산이 빡빡해, 깊은 의미 표현을 앵커로 두고 세부 정보를 제한적으로만 주입한다. 토크나이저 단계에서 융합 모듈과 디코더를 함께 학습하고(픽셀·지각·적대 손실, 디코더 입력 노이즈), 이후 토크나이저를 동결한 채 DiT 생성기를 학습한다. DRoRAE가 두던 융합 전용 적응 단계를 없앤 것이 차이다. 융합 잠재는 토큰 수와 채널 차원(16×16×1024)을 그대로 유지한다.

어떻게 쓰나

ImageNet-256에서 재구성 FID는 RAEv2의 0.299에서 0.209로 약 30% 줄었고, 5,000장 매칭 부분집합에서 PSNR은 22.667에서 26.377dB로, LPIPS는 0.074에서 0.043으로 개선됐다. 80에폭 생성기 학습 후 가이드 생성 gFID는 1.060에서 1.038로, IS는 255.300에서 257.823으로 좋아졌다. 다만 가이드 없이 쓰면 gFID 2.129로 RAEv2 K=23의 3.010보다는 낫지만 RAEv2의 1.650보다는 나쁘다.

전제와 한계

텍스트-이미지에서는 JourneyDB와 BLIP3o 캡션 부분집합으로 256×256 사전학습을 한 뒤 BLIP3o-60k로 지도 미세조정(SFT)한다. 사전학습 후 HiRAE-24는 RAEv2 대비 GenEval +4.52점, DPG-Bench +1.51점을 얻었고, SFT 후에는 GenEval +2.84점, DPG-Bench +1.45점, GenAI-Bench +0.97점이며 SFT 후 GenEval은 84.86에서 87.70으로 올랐다. 절제 실험에서 24개 층을 그대로 쓰는 DRoRAE식 융합을 RAEv2에 맞춰 적용하면 rFID는 0.065로 훨씬 낮아지지만 가이드 gFID는 1.551로 HiRAE-24의 1.038보다 나쁘고 FD_r^6도 3.375 대 1.856으로 뒤진다. 재구성 수치만 보고 융합 방식을 고르면 생성 품질을 놓친다는 뜻이다. 정규화를 빼면 rFID는 0.023까지 내려가지만 20에폭 가이드 gFID가 7.905로 폭발한다. 깊이 그룹은 3개일 때 가이드 gFID가 가장 낮았고, 어느 그룹의 잔차를 제거해도 LPIPS가 나빠져 세 그룹 모두 재구성에 기여한다.

잠재 공간 분석에서 5,000장 기준 10-최근접 이웃의 동일 클래스 비율은 RAEv2 74.454%에서 HiRAE-24 79.536%로 오히려 높아졌고, HiRAE-24 내부에서도 깊은 앵커 79.634%에서 융합 후 79.536%로 거의 유지된다. 공간 유효 랭크는 130.360에서 154.893으로 100장 모두에서 증가했고 평균 공간 CKA는 0.985로 패치 관계 구조는 비슷하게 남는다. 표준화 잠재 노름의 10% 크기 섭동에서 HiRAE-24의 디코딩 LPIPS 변화는 RAEv2의 21~22% 수준이고, 생성기도 7개 log-SNR 구간 중 내부 5개 구간에서 RAEv2보다 깨끗한 잠재를 더 정확히 예측한다.

실무에서는 고정된 비전 인코더 위에 토크나이저를 얹어 생성 모델을 학습할 때 참고할 만하다. 층 선택을 손으로 고르는 대신 전체 계층을 넣고 깊이별 노름 상한만 정하면 되므로 설정 부담이 줄고, 잠재 차원이 커지지 않아 기존 생성기 파이프라인을 그대로 쓸 수 있다. 다만 원문에는 별도의 한계 절이 없고 결론에서도 명시적 한계를 밝히지 않는다. 확인해야 할 전제는 실험이 256×256 해상도와 DINOv3-L 단일 인코더, ImageNet-1K 및 특정 T2I 데이터 조합에 한정된다는 점, 그룹 수와 예산·드롭아웃 확률 같은 하이퍼파라미터가 실험적으로 정해졌다는 점(2·3·4그룹 비교에서 3개 선택), 그리고 가이드 없는 생성에서는 RAEv2보다 나쁘고 재구성 FID만 보면 DRoRAE식 융합이 더 낮다는 상충이 남는다는 점이다.