U-Net

 

 

segmentation과 diffusion은 서로 다른 과제지만, 네트워크에 요구하는 공간적 성질이 상당히 비슷해서 U-Net이 diffusion에 잘 맞습니다.

먼저 segmentation도 형식상으로는 per-pixel classification입니다. 하지만 각 픽셀을 독립적으로 분류하는 단순한 문제가 아닙니다.

예를 들어 어떤 픽셀 하나만 보고 그것이 고양이인지 배경인지 판단하기 어렵습니다. 주변의 윤곽, 질감뿐 아니라 “이 영역 전체가 고양이 몸체인가?”라는 넓은 문맥까지 필요합니다. 따라서 segmentation에는 동시에 다음이 필요합니다.

• 정확한 픽셀 위치와 경계
• 국소적인 질감과 윤곽
• 물체 전체에 대한 넓은 문맥
• 입력과 출력 사이의 공간적 대응

이 요구사항이 diffusion의 denoising과 매우 유사합니다.

Diffusion도 위치별 출력을 만들어야 한다

Diffusion U-Net은 노이즈가 섞인 입력 xtx_txt​ 또는 latent ztz_tzt​를 받아, 보통 같은 공간 크기의 노이즈를 예측합니다.

ϵθ(zt,t,c)\epsilon_\theta(z_t,t,c)ϵθ​(zt​,t,c)

예를 들어 입력이

64×64×464\times64\times464×64×4

크기의 noisy latent라면, 출력도 일반적으로

64×64×464\times64\times464×64×4

입니다.

즉, U-Net은 각 공간 위치에 대해 “이 위치에 포함된 노이즈를 어느 방향으로 얼마나 수정해야 하는가?”를 출력합니다. 분류가 아니라 위치별 회귀 예측에 해당합니다.

과제입력위치별 출력

Segmentation	이미지	각 픽셀의 클래스
Diffusion	noisy image/latent	각 위치의 noise, vvv, 또는 관련 예측값

출력의 의미는 다르지만, 둘 다 입력과 공간적으로 정렬된 dense output을 만든다는 공통점이 있습니다.

왜 단순한 CNN보다 U-Net이 적합한가

Diffusion에서 노이즈를 제거하려면 국소 정보와 전체 정보를 모두 봐야 합니다.

예를 들어 특정 위치가 단순한 회색 노이즈처럼 보이더라도, 주변과 전체 이미지를 보면 그것이 사람의 눈이 될 위치인지, 배경이 될 위치인지 판단할 수 있습니다.

U-Net의 좌측 downsampling 경로는 해상도를 줄이며 receptive field를 넓힙니다.

• 높은 해상도: 모서리, 질감, 미세 패턴
• 중간 해상도: 물체의 부분과 형태
• 낮은 해상도: 전체적인 배치와 의미적 관계

따라서 diffusion에서는 좌측 경로가 현재 noisy latent를 여러 공간 규모에서 해석하는 역할을 합니다.

그런데 downsampling하면 위치 정보가 손실된다

해상도를 계속 줄이면 전체 구조를 이해하기는 쉬워지지만, 정확한 위치와 미세한 패턴은 약해집니다.

예를 들어 다음처럼 압축한다면,

64×64→32×32→16×16→8×864\times64 \rightarrow32\times32 \rightarrow16\times16 \rightarrow8\times864×64→32×32→16×16→8×8

8×88\times88×8 특징에서는 이미지 전체를 넓게 볼 수 있지만, 원래 64×6464\times6464×64 공간의 세부 정보는 많이 압축됩니다.

이 상태에서 단순히 다시 확대하면 다음이 흐려지기 쉽습니다.

• 정확한 윤곽
• 가는 선과 작은 질감
• 서로 가까운 구조의 구별
• 위치별 노이즈 패턴

이 문제를 해결하는 것이 U-Net의 skip connection입니다.

Skip connection이 diffusion에서도 중요한 이유

U-Net은 좌측의 높은 해상도 특징을 동일한 해상도의 우측 블록으로 직접 전달합니다.

우측 특징=깊은 의미 정보+좌측의 공간적 세부 정보\text{우측 특징} = \text{깊은 의미 정보} + \text{좌측의 공간적 세부 정보}우측 특징=깊은 의미 정보+좌측의 공간적 세부 정보

개념적으로 다음 두 정보를 결합합니다.

• 깊은 경로: “전체적으로 무엇을 만들어야 하는가?”
• skip connection: “각 위치에 현재 어떤 세부 패턴이 있는가?”

Segmentation에서는 이것이 정확한 물체 경계를 복원하는 데 사용됩니다. Diffusion에서는 현재 noisy latent의 위치별 세부 신호를 유지하면서, 전체 문맥에 맞는 denoising 결과를 계산하는 데 사용됩니다.

특히 diffusion은 한 번에 완성된 이미지를 출력하는 것이 아니라, 여러 timestep에 걸쳐 조금씩 수정합니다. 매 단계에서 입력 ztz_tzt​의 공간 구조를 유지한 채 적절한 수정 방향을 예측해야 하므로 skip connection이 매우 유용합니다.

공통점과 차이점

두 과제의 연관성은 다음과 같이 정리할 수 있습니다.

특성Segmentation U-NetDiffusion U-Net

공간적으로 정렬된 출력	픽셀별 클래스	위치별 noise 또는 vvv
국소 정보	경계·질감 판단	미세 패턴·노이즈 판단
전역 정보	물체 종류·영역 판단	이미지 구성·의미 판단
Downsampling	넓은 문맥 추출	넓은 문맥과 이미지 구조 파악
Upsampling	클래스 맵 해상도 복원	위치별 denoising 출력 복원
Skip connection	경계와 위치 보존	입력의 세부 공간 신호 보존
출력 성격	주로 분류	연속값 회귀

따라서 diffusion이 U-Net을 채택한 핵심 이유는 “segmentation 기술을 그대로 사용하기 위해서”라기보다는,

U-Net이 입력과 같은 공간 크기의 정교한 출력을 만들면서, 전체 의미와 위치별 세부 정보를 동시에 처리하는 구조이기 때문입니다.

Stable Diffusion에서는 여기에 조건 정보가 추가된다

Stable Diffusion의 U-Net은 원본 segmentation U-Net을 그대로 복사한 것은 아닙니다. 기본적인 down–up 구조와 skip connection을 가져오고 다음을 추가하거나 변경했습니다.

• Max pooling 대신 stride convolution 등의 downsampling
• 단순 Conv 블록 대신 residual block
• timestep embedding으로 현재 노이즈 단계 ttt 전달
• self-attention으로 떨어진 위치 사이의 관계 처리
• cross-attention으로 텍스트 조건 전달
• 픽셀이 아닌 VAE latent 공간에서 작동

텍스트 조건은 주로 “무엇을 그릴 것인가”를 알려주고, U-Net의 다중 해상도 구조는 그것을 “어떤 위치와 규모에서 어떻게 구현할 것인가”를 계산합니다.

결론적으로 segmentation과 diffusion의 직접적인 목적은 다르지만, 두 문제 모두 다음 형태를 공유합니다.

공간 입력⟶전역 문맥과 국소 정보의 결합⟶입력과 정렬된 공간 출력\text{공간 입력} \longrightarrow \text{전역 문맥과 국소 정보의 결합} \longrightarrow \text{입력과 정렬된 공간 출력}공간 입력⟶전역 문맥과 국소 정보의 결합⟶입력과 정렬된 공간 출력

이 구조적 공통점 때문에 segmentation에서 성공했던 U-Net이 diffusion의 denoising backbone으로도 매우 적합했던 것입니다.