저번에는 LLM / VLM 모델의 count에 관한 논문을 읽었고,
이번에 읽는 이 논문은 diffusion 모델의 count 관련한 논문을 읽어보려고 한다.
CVPR 2025 논문이다.
0. Abstract
Text-to-image 모델이 전례없는 성공을 거두었음에도 불구하고 텍스트를 이용해서 이미지 속 객체 수를 control하는 것은 놀라울 정도로 어렵다.
이게 아동도서, 요리 레시피 삽화 등의 다양한 분야에서 중요하다.
정확한 개수의 객체를 세는 것은 근본적으로 어려운 문제이다.
생성모델이 계속(여러 객체가 비슷하게 생기거나 겹쳐 있더라도) 객체의 separate한 identity를 유지한 다음, 생성 과정에서 암묵적으로 전체 개수를 계산해야 하기 때문이다.
즉, 목표가 6개이고 현재 생성된 객체가 4개일 때, 모델은 아직 2개가 부족하다는 것을 내부적으로 계산해야 한다.
이러한 표현이 존재하는지는 아직 알려져 있지 않다.
정확한 객체 생성을 위해서 먼저 diffusion 모델 안에서 object identity 정보를 가지고 있는 feature를 찾아냈다.
그리고 이 feature를 사용해서 denoising 과정에서 객체를 분리하고 개수를 세서 over생성과 under생성을 탐지한다.
Under 생성은 기존에 생성된 layout을 기반으로 빠진 객체의 모양과 위치를 예측하는 모델을 학습시켜서 해결한다.
이렇게 정확한 개수를 갖도록 denoising 과정을 유도할 수 있음을 보였다.
저자들이 고안한 CountGen 방식은 object layout를 결정하기 위해서 외부 소스를 사용하지 않는다.
Diffusion 모델 자체의 prior를 사용한다! (우와)
평가 결과 CountGen이 기존 방법들보다 높은 counting 정확도를 보였다.
3. Our Approach: CountGen
CountGen의 목표는 text를 조건으로 하는 이미지 생성기가 복잡한 프롬프트에도 제대로 의도한 개수의 object를 만들도록 하는 것이다.
방법은 2-step으로 구성된다.
1. 먼저 text 조건으로부터 object가 어디에 어떻게 나와야 자연스러운지를 나타내는 자연스러운 layout을 생성한다.
2. 그 다음 이 layout를 청사진으로 하여 최종 이미지를 생성한다.
3.1 Discover Layout During Early Generation
생성 과정에서 object수를 세기 위해서는 먼저 서로 다른 object들의 개별적인 정체성을 표현하는 내부의 representation을 찾아야 한다.
근데 이러한 내부 representation이 SDXL과 같은 diffusion 모델에 존재하는지조차 밝혀지지 않았다.
그래서 이러한 내부 representation을 설명하고, 생성 초기에 어떻게 object들의 layout를 탐지할지를 보인다.
An emerging instance-identity representation in SDXL.
SDXL 안에서의 'objectness' 개념이 무엇인가?
기존 연구들은 cross-attention을 이용해서 생성된 이미지 안에서 특정 클랙스의 객체 위치를 찾았다.
근데 이 객체 하나하나의 위치를 구별하는지에 대한 연구는 거의 없다.
저자들은 서로 다른 timestep, layer에서 feature 분석을 해서 이 문제를 해결한다.
아래는 여러 layer의 self-attention feature들을 PCA로 시각화 한 결과이다. 이 중 l_52&{up} layer가 가장 outperform 한다.
가장 outperform 한다는 게 무언인가? 서로 다른 instance들이 다른 색을 띄고 있다. 즉 서로 다른 feature를 생성한다는 것이다.
살짝 충격적인 것은, 같은 instance면 같은 feature로 여길 거라고 생각했는데 아니었다. 다른 애로 인식을 한다!
Identifying object instances.
Cross-attention으로 각 객체애 대한 foreground mask를 생성한다.
이렇게 생성된 foreground mask하고 self-attention feature를 결합해서 객체 픽셀이랑 배경 픽셀을 분리한다.
그리고 self-attention map에서 나온 객체 픽셀들을 클러스터링해서 각 객체별 mask를 생성한다.
이 접근은 객체 표현을 정교화하고 생성된 layout의 정확도를 높여준다.
수식으로 각각 $A_{l,t}^{self}$, $A_{l,t}^{cross}$라고 표현한다.
입력된 프롬프트에 있는 객체에 해당하는 토큰들의 cross-attention map을 합친 다음 *dynamic thresholding을 이용해서 foreground mask M = **Otsu(A_{l,t}^{cross})을 추출한다.
* dynamic thresholding: Attention 스케일이 timestep, layer마다 계속 바뀌기 때문에 각각에 따라서 threshold를 설정함
** Otsu: 유명한 자동 threshold 방법. 그룹 안의 분산은 최소화하고 그룹 간 분산은 최대화 하는 threshold!
Foreground mask M에 의해서 추출된 self-attention feature들을 p_k라고 하고, 이를 클러스터링 시킨다.
Clusters = *DBSCAN(p_k, \epsilon)
* DBSCAN: 밀도 기반 클러스터링. \epsilon은 탐색 반경 거리
이렇게 해서 나온 클러스터들을 grouping해서 최종적으로 initial layout $L = \bigcup_{C \in \mathrm{Clusters}} C$이 생긴다.
3.2 ReLayout: Correcting the Number of Objects
전체적인 장면 구성은 유지하면서 object의 개수를 수정하는 ReLayout를 소개한다.
인풋은 3.1에서 생성된 object-layout이다. 이 object-layout에 있는 object의 수를 센다.
Over-Generation의 경우에는 그냥 제일 작은 인스턴스를 제거한다.
Under-Generation의 경우에는 아래처럼 ReLayout이 기존의 장면구조를 해치지 않는 범위에서 새로운 인스턴스를 추가한다.
딱 맞게 object들이 생성된 경우에는 ReLayout가 원래의 layout을 그대로 유지한다.
3.2.1에서 Under-Generation을 다루는 방법을 더 자세하게 소개할 것이다.
3.2.1 Handling Under-Generation
Under-Generation 문제를 다루기 위해서 기존 layout으로부터 새로운 layout을 예측하도록 하는 U-Net을 학습시킨다.
U-Net 안에서 forward pass가 1번 진행될 때마다 object 1개가 추가된 새로운 mask가 생성된다. 개수가 맞을 때까지 반복된다.
이제 아래에서 U-Net의 구조랑 학습방법을 자세하게 설명한다.
Q: 왜 하필이면 U-Net을 사용하는가?
Under-Generation 문제를 해결할 때 왜 다른 모델도 아니고 하필이면 U-Net 을 사용하는지에 대한 의문이 들었다.
U-Net이 위치 정보를 잃지 않으면서 전체 구조까지 보는 모델이기 때문이다!
CNN은 풀링되면서 해상도가 줄어들면서 무슨 이미지인지는 잘 파악하는데, 물체가 구체적으로 어디에 있는지는 잘 파악하지 못한다. U-Net은 애초에 이 문제를 해결하기 위해 제안된 구조이다.
U-Net의 encoder에서는 쭉쭉 풀링하면서 큰 구조를 이해하고, decoder에서는 다시 쭉쭉 원래 크기로 키우면서 픽셀 단위로 다시 그린다. 그리고 skip connection이 정확한 경계선과 위치를 보존한다.
이러한 원리로 ReLayout에서 사용되는 U-Net은 encoder에서 '고양이들이 삼각형으로 배치되어 있구나'를, decoder에서 '새 고양이는 여기에 넣자'를, skip connection은 기존 고양이들의 위치와 경계를 유지하도록 돕는다.
Creating a training dataset.
ReLayout의 U-Net을 학습하기 위해서는 비슷한 layout을 유지하면서 object수만 k개랑 k+1개로 차이나는 layout pair 데이터셋이 필요하다.
시드를 고정하고 프롬프트만 2 cats에서 3 cats로 바꾸면 비슷한 layout의 이미지가 생성된다는 것을 관찰했다. (우와 관찰 미쳤다)
이러한 관찰을 바탕으로 약 1만 개의 layout pair 데이터셋을 생성했다.
이를 검증할 때는 3.1에서 설명한 object mask를 추출하는 방법을 상요했다.
Matching objects.
U-Net을 학습하기 위해서는 Mk에서의 각 object와 Mk+1에서의 각 object 간의 대응관계를 알아야 한다.
예를 들어서 ㅁ ㅁ / ㅁ ㅁ ㅁ 일 때 ㅁ ㅁ / ㅁ ㅁ ㅁ 인지 알아야 한다는 것이다.
대응 관계를 찾기 위해서 *hungarian 알고리즘을 사용하여 object의 위치 변화가 최소가 되도록 하는 matching을 찾는다.
*hungarian 알고리즘: 전체 비용을 최소로 하도록 1:1 매칭하는 알고리즘
Training the U-Net module.
U-Net 구조를 수정해서 input 채널 9개, output 채널 10개까지 처리할 수 있도록 학습했다.
CountGen이 최대 10개 object까지 다루기 때문이다.
아래 2개의 loss function을 이용해서 U-Net을 학습했다.
(1) predicted mask랑 target mask 사이의 *dice loss
(2) mask끼리 겹치지 않도록 하는 overlap loss
*dice loss: predicted mask랑 target mask가 얼마나 겹치는가를 측정함
Inference.
추론 시 전처리 단계에서 먼저 input mask 주변에 패딩을 추가한다. 그릴 공간을 넉넉하게 주기 위해서이다.
다 그린 다음에는 끝에 조금 깎아서 붙어있는 object들을 조금 떼준다.
3.3 CountGen: Layout-Based Image Generation
3.2에서 얻은 object mask layout이 주어졌을 때, 이제 이미지 생성 과정이 이 layout을 따르도록 해야 한다.
이때 inference time optimization을 사용한다. 즉 모델을 다시 학습시키는 게 아니라 inference 과정에서 loss를 계산하고, 그 gradient를 이용해서 현재 latent를 직접 수정하는 것이다.
2가지 방법(object layout loss, self-attention masking)를 사용한다.
Object layout loss는 foreground 영역에서 object가 생성되도록 유도하고, self-attention masking은 background 영역에 object가 생성되지 않도록 유도한다.
Object layout loss.
가중치가 적용된 BCE Loss를 이용해서 layout 내의 object 위치를 최적화한다.
합쳐진 cross-attention score를 $c$라고 하고 object의 존재 여부를 나타내는 binary mask를 $m$이라고 할 때 수식은 아래와 같다.
픽셀 단위로 계산되고, 가중치 $w_i$는 object 픽셀에서 10 background에서 1이다.
원래 SDXL 생성 과정에서는 각 timestep마다 noisy latent $X_t$를 인풋으로 받는다.
근데 여기서는 생성 초기 25 step동안 object layout loss를 $X_t$로 전달해서 loss가 감소하도록 X_t를 업데이트 한다.
(생성 초반에만 하는 이유: 처음만 잘 잡아주면 그대로 생성하기 때문)
Self-attention masking.
Object layout loss만 적용하면 object가 mask 밖에서도 생성될 수 있다.
이를 해결하기 위해서 background 픽셀에서 foreground 픽셀로 이어지는 self-attention connection을 masking한다.
이렇게 해서 object에서 이미지의 나머지 부분으로 전달되는 정보 흐름을 막아서 object가 background에 생성되는 것을 막는다.
수식으로 표현하면 아래와 같다. i: attention을 계산하는 픽셀(Query), j: attention을 참조하는 픽셀(Key) 이다.
즉, query가 background이고 key가 foreground이면 attention을 0으로 만드는 것이다!
Q: 이렇게 하면 background에 object가 어케 안 생기게 되는 것인가?
Foreground에 있는 고양이 feature가 background로 퍼져나가면 background도 그 고양이 feature를 이용해서 고양이의 일부를 생성할 가능성이 생긴다!
이 통로를 끊어서 고양이 feature 정보가 foreground 안에만 유지되도록 하는 것이다!
(필요할 때 이 방법 적절히 응용하면 유용할 것 같다. 기억해놓자)
그래서 전체 모델 구조는 아래와 같다.
7. Limitations
3.3 때문에 object 1개만 있어야 하는 layout에 여러 개의 object가 생성되기도 한다.
즉, layout은 맞았는데 실제 이미지 생성 과정에서 같은 위치에 object가 중복으로 생성되는 경우가 있다는 것이다.
그리고 SDXL보다 CountGen이 더 단순한 배경을 생성하기도 한다. 제약을 걸기 때문이다.
한편으로는 실험 범위를 object 10개로 제안했기 때문에 조금 범위가 좁아보일 수 있다.
그럼에도 불구하고 section 5에서 보면 최신 모델들에게 이 task는 매우 어려운 문제임을 알 수 있다.
특히 생성해야 하는 object 수가 늘어날수록 성능이 크게 감소한다.
8. Conclusions
개수를 딱 맞게 생성하는 것은 매우 어려운 task이다.
모델이 objectness를 이해해야 할 뿐만 아니라 공간적 배치를 만족하면서 자연스러운 이미지를 생성해야하기 때문이다.
특히 3개 이상의 객체를 생성할 때 성능이 크게 떨어진다.
이를 해결하기 위해서 크게 3가지 단계를 사용한다.
1) diffusion 모델의 self-attention layer에서 objectness를 추출함
2) layout 구조를 유지하면서 객체 수를 생성하는 U-Net 학습시킴
3) 수정된 layout을 이용해서 이미지를 생성하기 위한 layout-guidance optimization
이렇게 해서 SDXL의 counting 정확도를 26%에서 52%로 향상시켰다.
향후에는 CountGen으로 multiple object를 다룰 수 있도록 확장하고, 다른 생성 모델 구조에도 CountGen을 적용시키려고 한다.
---
이미지 생성 모델의 counting을 다룬 몇 안 되는 논문이라서 자세하게 읽어봤다.
자세하게 읽으니 시간은 오래 걸리지만 그래도 실험에 쓸만한 아이디어를 얻게 되는 것 같다.
중요한 논문들은 이렇게 자세히 읽어보려고 한다.
재밌다 counting task!