PiD : 픽셀 확산으로 더 빠르게 만드는 고해상도 이미지 디코딩 기술

 

이미지 생성 AI는 겉으로 보면 문장을 넣고 이미지를 받는 단순한 도구처럼 보입니다. 하지만 내부에서는 먼저 작고 압축된 공간에서 이미지를 설계한 뒤, 마지막에 다시 사람이 볼 수 있는 픽셀 이미지로 바꾸는 과정을 거칩니다. 이 마지막 변환 과정이 생각보다 중요합니다. 아무리 좋은 구도를 만들었어도, 디코더가 세부 질감이나 작은 글자, 머리카락 같은 부분을 잘 복원하지 못하면 결과물은 흐릿해질 수 있습니다.

 

최근 생성 AI 산업은 더 큰 해상도와 더 짧은 생성 시간을 동시에 요구받고 있습니다. 광고 이미지, 제품 컷, 영화 콘셉트 아트, 게임 배경처럼 실제 작업에 쓰려면 2K나 4K 수준의 선명도가 필요하죠. 이번 논문은 기존 방식처럼 낮은 해상도로 먼저 복원하고, 다시 업스케일러로 키우는 여러 단계를 거치는 과정을 하나로 합쳐, 더 빠르고 선명하게 이미지를 만드는 방법을 제안합니다. 핵심은 '잠재공간을 바로 고해상도 픽셀로 디코딩'한다는 생각입니다.

---

 

 

1) 기존 방식의 문제점

 

기존 고해상도 이미지 생성 파이프라인은 보통 두 단계를 거칩니다. 먼저 VAE 디코더가 잠재표현을 낮은 해상도 이미지로 바꾸고, 그다음 초해상도 모델이 이미지를 크게 키웁니다. 그러다보니 다음과 같은 문제가 발생합니다. 

• VAE 디코더가 '새로운 디테일을 그리는 모델'이라기보다 '압축된 정보를 복원하는 모델'에 가깝기 때문에 작은 글자, 피부 질감, 나뭇잎, 머리카락 같은 고주파 디테일이 약해질 수 있습니다.
• 낮은 해상도 디코딩, 초해상도 확산 모델, 다시 고해상도 디코딩 같은 단계가 이어지면 시간이 늘고 메모리도 많이 쓰게 되어 작업 비용이 증가됩니다. 

논문에서는 이런 문제 때문에 다단계 방식 대신, 잠재표현에서 바로 고해상도 픽셀 이미지를 생성하는 PiD를 제안합니다. PiD는 512×512에 해당하는 잠재표현을 2048×2048 이미지로 1초 이내에 디코딩하고, GB200 GPU에서는 약 210ms까지 줄였다고 보고합니다.

 

2) 접근 방식

 

PiD는 'Pixel diffusion Decoder'의 약자입니다. 말 그대로 픽셀 공간에서 작동하는 확산 디코더입니다. 기존 디코더처럼 잠재표현을 단순히 복원하지 않고, 잠재표현을 조건으로 삼아 고해상도 이미지를 직접 생성합니다. 여기서 잠재표현은 전체 구도와 의미를 알려주는 설계도 역할을 하고, 픽셀 확산 모델은 표면 질감과 세부 묘사를 채우는 화가 역할을 합니다.

 

그림 : PiD의 전체 아이디어 (상) 및 성능/속도 비교 (하)

 

이 방식의 핵심은 디코딩과 업스케일링을 분리하지 않는 것입니다. PiD는 4배 또는 8배 확대 결과를 하나의 생성 과정에서 만듭니다. 또한 노이즈가 조금 남아 있는 잠재표현도 처리할 수 있게 훈련되어, 원래의 잠재 확산 모델을 끝까지 돌리지 않고 중간에 멈춘 뒤 PiD가 나머지를 고해상도 이미지로 완성할 수 있습니다.

 

3) 모델/기술 아키텍처

 

PiD의 큰 구조는 세 부분으로 볼 수 있습니다.

• 고해상도 픽셀 이미지를 직접 다루는 PixelDiT 기반 확산 모델입니다.
• 낮은 해상도의 잠재표현을 모델 내부 토큰 형태로 바꿔 넣어주는 경량 어댑터입니다.
• 잠재표현이 얼마나 깨끗한지에 따라 참고 강도를 조절하는 sigma-aware gate입니다.

 

 

그림 : 모델 아키텍처

 

위의 아키텍처에서 보면, PiD는 '흐릿한 고해상도 캔버스'와 '작은 설계도'(잠재표현정보)를 동시에 봅니다. 캔버스는 점점 노이즈가 제거되며 이미지가 되고, 설계도는 인물 위치, 사물 형태, 장면 의미를 알려줍니다. sigma-aware gate는 '이 설계도를 얼마나 믿을지'를 조절합니다. 그래서 완전히 깨끗한 잠재표현뿐 아니라, 아직 일부 노이즈가 남은 중간 잠재표현도 처리할 수 있습니다. 

 

 

 

4) 세부 적용 기술

 

1️⃣ 조건부 픽셀 확산 디코딩

PiD는 잠재표현을 바로 고해상도 픽셀 이미지로 바꿉니다. 기존 방식은 '잠재표현 → 저해상도 이미지 → 초해상도 이미지' 순서였지만, PiD는 '잠재표현 → 고해상도 이미지'로 바로 갑니다. 조건부라는 말은 모델이 텍스트와 잠재표현을 참고해 이미지를 만든다는 뜻입니다.

 

2️⃣ PixelDiT 기반 고해상도 픽셀 생성 prior

prior는 모델이 이미 학습해 둔 '그럴듯한 이미지에 대한 감각'을 의미합니다. PiD는 PixelDiT라는 픽셀 공간 확산 트랜스포머를 기반으로 사용합니다. 덕분에 단순 복원보다 더 자연스러운 질감, 선명한 경계, 작은 패턴을 만들 수 있습니다. 

 

3️⃣ Noisy latent conditioning

훈련 중 PiD는 깨끗한 잠재표현만 보지 않으며 일부러 노이즈가 섞인 잠재표현도 보여줍니다. 이렇게 하면 모델이 잠재표현을 너무 맹신하지 않고, 부족한 세부 정보를 스스로 생성하는 능력을 갖게 됩니다. 또 실제 추론 때 잠재 확산 모델을 중간에 멈춰도 PiD가 그 상태를 받아 처리할 수 있습니다.

 

4️⃣ Sigma-aware gate

sigma는 노이즈의 양을 뜻합니다. sigma-aware gate는 잠재표현에 노이즈가 많으면 적게 믿고, 깨끗하면 더 많이 믿도록 조절하는 장치입니다. 초보자에게 비유하면, 흐릿한 지도는 참고만 하고, 선명한 지도는 더 믿는 방식입니다. 

 

5️⃣ Latent projection and injection

낮은 해상도의 잠재표현은 그대로 PixelDiT에 넣을 수 없습니다. 그래서 PiD는 잠재표현을 크기에 맞게 조정하고, 합성곱 블록으로 특징을 뽑은 뒤, 트랜스포머가 읽을 수 있는 토큰으로 바꿉니다. 그런 다음 여러 DiT 블록 사이에 이 정보를 주입합니다. 쉽게 말해, 작은 설계도를 모델 내부의 여러 작업 단계에 나눠 전달하는 방식입니다.

 

그림 : 잠재표현 조건 삽입

 

6️⃣ Rectified Flow 학습

Rectified Flow는 노이즈 이미지를 깨끗한 이미지로 바꾸는 방향을 학습하는 방법입니다. 확산 모델이 여러 번 조금씩 노이즈를 줄이는 것과 비슷하지만, 여기서는 '어느 방향으로 이동하면 깨끗한 이미지에 가까워지는지'를 예측합니다. PiD는 이 방식으로 고해상도 픽셀 공간의 변화 방향을 배웁니다.

 

7️⃣ DMD2 기반 4-step distillation

PiD는 처음에는 여러 단계로 동작하는 teacher 모델을 학습한 뒤, DMD2로 4단계 student 모델로 줄입니다. distillation은 큰 모델이나 느린 과정을 작은 모델 또는 빠른 과정으로 압축하는 기술이라고 생각하시면 됩니다. 

 

그림 : DMD2 증류 과정 설명

 

8️⃣ LDM early termination

기존 잠재 확산 모델은 모든 denoising step을 끝까지 수행합니다. PiD는 중간 단계의 잠재표현도 받을 수 있기 때문에, 기본 LDM을 조금 일찍 멈출 수 있습니다. 이렇게 하면 전체 생성 시간이 줄어듭니다. 다만 너무 일찍 멈추면 장면 구조가 덜 잡혀 품질이 떨어질 수 있습니다.

 

9️⃣ VAE latent와 RAE semantic latent 모두 지원

PiD는 FLUX, SD3 같은 VAE 잠재표현뿐 아니라 DINOv2, SigLIP 기반 RAE 잠재표현에도 적용됩니다. RAE 계열은 의미 정보는 강하지만 세부 질감이 부족할 수 있는데, PiD는 이 부족한 부분을 픽셀 확산 모델의 생성 능력으로 보완합니다.

 

 

6) 제약사항

• 훈련 비용이 큽니다. 2K 모델 학습에는 대규모 H100 GPU가 사용되며, 4K 확장에는 GB200 GPU 여러 대가 필요합니다. 실제 서비스에 적용하려면 학습 비용과 운영 비용을 따져야 합니다.
• PiD는 생성 능력이 강한 만큼 원본 잠재표현에 없던 디테일을 상상할 수 있습니다. 이는 장점이지만, 문서 복원이나 의료 영상처럼 원본과 정확히 일치해야 하는 작업에서는 조심해야 합니다.
• early termination은 항상 좋은 것이 아닙니다. 너무 일찍 멈추면 잠재표현의 의미 구조가 덜 만들어져 이미지 품질이 떨어질 수 있습니다. 반대로 끝까지 돌리면 PiD가 세부 질감을 새로 만들 여지가 줄고, 전체 시간도 늘어납니다.
• 4K 디코딩은 가능하지만 메모리 부담이 사라지는 것은 아닙니다. 

---

 

PiD의 의의는 디코더를 단순한 복원 장치에서 생성 장치로 바꿨다는 점입니다. 기존 VAE 디코더는 압축된 정보를 되돌리는 데 초점이 있었지만, PiD는 잠재표현을 바탕으로 고해상도 픽셀 이미지를 직접 합성합니다. 그래서 디코딩과 업스케일링을 하나의 과정으로 묶고, 속도와 품질을 함께 개선하려고 합니다.

 

실생활 적용 측면에서보면, 고해상도 결과가 필요한 이미지 작업에서 후처리 시간을 줄일 수 있다는 장점이 있을 것입니다. 

---

 

참고자료

• 사이트) PiD 프로젝트 사이트 (링크) - Overview of PiD 영상을 참고하세요.
• 논문) PiD: Fast and High-Resolution Latent Decoding with Pixel Diffusion (링크)
• 코드) NVDIAdml PiD github 저장소 (링크)

---

 

Q&A

 

Q. PiD는 기존 VAE 디코더와 무엇이 다른가요?

PiD는 잠재표현을 단순히 복원하지 않고, 픽셀 확산 모델을 이용해 고해상도 이미지를 직접 생성합니다. 그래서 디코딩과 업스케일링을 한 번에 처리합니다.

 

Q. sigma-aware gate는 왜 필요한가요?

잠재표현이 깨끗하면 더 많이 참고하고, 노이즈가 많으면 덜 참고하기 위해 필요합니다. 흐릿한 설계도를 그대로 믿으면 결과가 흔들릴 수 있기 때문입니다.

 

Q. PiD가 항상 원본에 더 충실한가요?

항상 그렇지는 않습니다. PiD는 더 선명하고 자연스러운 디테일을 만들 수 있지만, 때로는 원래 잠재표현에 없는 디테일을 추가할 수 있습니다. 그래서 정확한 복원이 중요한 작업에서는 검증이 필요합니다.