CVPR 2024
https://arxiv.org/abs/2403.10071
Baoquan Zhang, Huaibin Wang, Luo Chuyao, Xutao Li, Liang Guotao, Yunming Ye, Xiaochen Qi, Yao He
Introduction
멀티모달 학습의 발전과 함께, 최근에는 모든 modality를 transformer로 통합하려는 시도가 점점 더 많은 관심을 받고 있다. transformer는 본래 discrete token seqence를 모델링하기 위해 제안된 구조로, 이미지와 같은 continuous에 적용하는 것은 매우 어렵다. 이러한 문제를 해결하기 위해 Vertor-Quantized Image Modeling(VQIM)이 제안되었으며, 이는 이미지를 텍스트처럼 discrete한 seqence로 인코딩하는 것을 목표로 한다.
일반적으로 이전 연구들은 VQIM을 수행하기 위해 처음부터 discrete codebook을 학습하고, continuous feature representation을 discrete token으로 양자화하는 방식을 채택해왔다. 예를 들어, VQ-VAE의 경우 encoder-decoder 구조를 제안하여, latent feature space를 학습하고, 코드북 내에서 가장 가까운 벡터를 선택해 토큰으로 양자화한다. 이때 모델은 단순한 reconstruction loss를 통해 모델을 학습되며, 이후 다른 연구에서는 adversarial loss를 추가로 도입하여 성능을 향상시키기도 하였다. 이러한 방법들은 좋은 성능을 보였지만 codebook collapse라는 문제가 존재한다. 이는 훈련 과정에서 소수의 code vector만이 사용되고, 대부분의 벡터는 전혀 업데이트되지 않아 전체 성능을 저하시키는 것을 의미한다. 최근 이 문제를 해결하기 위해 codebook update, quantization, regularization 등 다양한 관점에서 보다 robust하게 학습시키기 위해 연구가 진행되고 있다.
본 논문의 저자들 또한 코드북 collapse 현상을 완화하기 위해 codebook transfer이라는 새로운 관점을 제시한다. 이 접근은 코드벡터 간의 관계나 코드북 사전(prior)의 구조를 고려하지 않고 처음부터 코드북을 학습하는 것이 매우 어렵다는 통찰에서 출발한다. 이러한 관계가 무시되면 코드북 붕괴 현상이 발생하고, 이는 결국 일부 코드벡터가 업데이트되지 않는 die off현상으로 이어진다고 저자들은 주장한다. 저자들은 최근 pre-trained language model(PLM)에서 영감을 받아 일부 모델들이 이미 잘 학습된 코드북을 내포하고 있으며, 이 코드북 내의 벡터들은 완전히 독립적인 것이 아닌 vision과 language 간 전이 가능한 relationship을 포함하고 있다는 점에 주목한다.
예를 들어, 오른쪽 그림에서는 시각적으로 유사한 “orange”와 “yellow”가 “blue”보다 코드북 공간에서 더 가까운 거리에 있고, 왼쪽 그림에서는 “breast”와 “belly”가 유사한 시각적 의미를 갖는 반면, “belly”와 “wheel”은 더 멀리 떨어져 있는 것을 확인할 수 있다. 이러한 관계들을 효과적으로 활용하면, 특정 코드가 직접 선택되지 않더라도 그와 연관된 다른 코드들과의 관계를 통해 cooperative optimization될 수 있고, 결과적으로 코드북 붕괴 현상을 완화할 수 있다고 저자들은 설명한다.
이러한 아이디어를 기반으로 저자들은 품사 정보를 활용하여 PLM으로 부터 codebook의 풍부한 semantic knowledge를 transfer하는 VQCT를 제안한다. 구체적으로 PLM과 품사 정보를 코드북 prior으로 도입하고, 품사에 따라 시각과 관련없는 token들을 필터링하여 형용사와 명사 코드북만을 남긴 비전 관련 코드북을 구축한다. 그리고 언어에서 비전으로 transfer mapping 방식으로 VQIM을 모델링하기 위해, graph convoultion-based codebook transfer network를 저자들은 새롭게 디자인하였다. 이러한 방식은 사전 학습된 코드북 지식을 충분히 활용하여 우수한 prior를 제공하고, 코드 간의 semantic relationships을 온전히 활용함으로싸 cooperative optimization을 달성할 수 있는 강점을 가진다.
Methodology
앞서 언급했듯이, 코드북 collapse 문제를 해결하려는 여러 시도가 있었지만, 기존 연구들은 코드북을 scratch 방식으로 학습하고, 코드 간 독립적인 방식을 접근하는데 초점을 맞춰왔다. 이에 대해 저자들은 코드 벡터와 코드북 prior 간의 관계를 고려하지 않은 채 처음부터 학습하는 방식은 본질적으로 매우 어렵고, 오히려 collapse를 야기시킬 수 있다고 주장한다.
이에 기반하여, 저자들은 품사 정보를 활용한 graph convolution-based codebook transfer framework인 VQCT를 제안한다. 핵심 아이디어는 코드북을 처음부터 학습하는 것이 아닌 PLM로부터 잘 학습된 코드북과 품사 정보를 prior 지식으로 도입하고, 이 정보를 기반으로 한 semantic knowledge와 관계를 활용하여 VQIM codebook 학습을 향상시키는데 있다.
VQCT framework는 다음과 같이 총 3가지로 구성된다.
- Encoder $f_{\theta_e}(\cdot)$
- Codebook Transfer Module $f_{\theta_t}(\cdot)$
- Decoder $f_{\theta_d}(\cdot)$
Codebook Transfer Module
Step 1: Codebook Construction. 기존 PLM은 이미 codebook을 제공함에도 불구하고, 이러한 정보는 기존 VQIM 영역에서는 전혀 활용되지 않았다. 하지만 VQIM으로 모든 단어를 transfer 것은, 단어 수가 매우 방대하기에 현실적으로 불가능하다. 따라서 저자들은 형용사, 명사 외에 다른 품사들은 시각 정보와 무관하다는 직관에 기반하여, 품사 정보를 prior로 도입한다. 즉, 형용사와 명사만을 남겨 vision 관련 코드북을 구축한다.
- Adj. PLM Codebook $R^{adj} = \left\{ r^{adj}_i \right\}_{i=0}^{K_{adj} - 1}$
- Noun. PLM Codebook $R^{noun} = \left\{ r^{noun}_i \right\}_{i=0}^{K_{noun} - 1}$
여기서 $K_{adj}, \quad K_{noun}$는 각각 형용사 및 명사 코드 수를 의미한다.
Step 2: Codebook Transfer. 저자들은 형용사와 명사 코드북을 개별적으로 추출한 뒤 이를 단순히 concatenate하여 하나의 코드북으로 만들거나, MLP를 통해 VQIM 코드북을 생성하는 방식은 적절하지 않다고 주장한다. 이는 실제로 sharp beak처럼 형용사-명사 조합이 시각적 특징을 설명하는 데 필수적이며, 두 코드북이 상호 독립적이지 않고 공동으로 사용되기 때문이다. 이에 따라 저자들은 형용사와 명사 간의 modifying relationship을 모델링하기 위해 modifying graph $ \mathcal{G} = (R, A) $를 먼저 구축한다. (본 논문에서는 text corpus라고만 했지 어떤 것을 사용했는지 구체적으로 명시되어있지는 않음)
- $ R = \{ R^{adj}, R^{noun} \}$ : node representation for PLM
- $ A = \{ A_{ij} \}$ : modifying relationship matrix ($A_{ij}=1$, 형용사 $i$가 명사 $j$를 수식하는 경우, 그렇지 않은 경우는 0)
이후 modifying graph $ \mathcal{G}$를 입력으로 하여 Graph Convolution-based Codebook Transfer Network(GCCTN) $f_{\theta_t}( \cdot )$를 통해 형용사와 명사 코드북을 VQIM으로 transfer한다.
$$ (c_{adj}, c_{noun}) = f_{\theta_t}(\mathcal{G}), \quad \mathcal{G} = (R, A)$$
여기서 GCCTN은 총 3개의 graph convolution layer로 구성되며 각 레이어 뒤에는 ReLU activation fucn, 이 적용된다. 최족적으로 VQIM 코드북은 형용사 $C_{adj} = \left\{ (k, e_k \in \mathbb{R}^{n_c}) \right\}_{k=1}^{K_{adj}}$와 명사 $C_{noun} = \left\{ (k, e_k \in \mathbb{R}^{n_c}) \right\}_{k=1}^{K_{noun}}$ 두 개의 코드북이 생성된다.
step 3: Vector Quantization. 주어진 image $x$에 대해 encoder $f_{\theta_e}(\cdot)$를 사용하여 continuous vector를 추출한다. 이 vector는 두 부분으로 나눠어, 각각 adjective vector, noun vector로 사용된다. 이후 각 벡터에 대해, 형용사 코드북 $C_{adj}$과 명사 코드북 $C_{noun}$을 기반으로 quantization operation $q(\cdot)$을 적용하여 각각 discrete code sequence로 변환한다. 이때 $l_2 norm$을 통해 가장 가까운 neighbor를 선택한다. 양자화된 형용사/명사 벡터는 concatenate되어 최종 quantized vector를 구성하며, 이를 decoder $f_{\theta_d}(\cdot)$에 입력하여 이미지 $\hat{x}$를 복원한다. 저자들은 이 방식이 기존 코드북 학습 기법들과 쉽게 통합 가능하며, 실험에서는 GAN 구조를 결합하여 성능을 평가하였다.
Analysis of Codebook Optimization
저자들은 VQCT가 코드북 collapse 문제를 완화하는지를 분석하기 위해, 2차원 코드북 환경을 설정하여 그림과 같이 시각화하였다. 위에서 알 수 있듯이, 기존 VQ-VAE는 active code vector만 최적화되며, 선택되지 않은 vector들은 전혀 업데이트되지 않고 die off 되는 분제가 발생한다. 반면 VQCT의 경우 전반적으로 업데이트 전 VQ-VAE와 상황은 비슷하지만, 업데이트 후에는 코드 간 semantic relationship를 기반으로 cooperative optimization이 이뤄진다. 즉, 하나의 코드벡터만 선택되었더라도, 그와 의미적으로 연관된 다른 코드들 또한 함께 최적화된다. 저자들은 이러한 현상의 원인으로, VQ-VAE처럼 코드북 자체를 학습하는 방식이 아니라, GCCTN을 통해 코드북을 생성하는 구조를 사용했기 때문이라고 설명한다. 이로 인해 선택되지 않은 코드벡터도 semantic connection을 통해 간접적으로 업데이트가 가능해져, 결국 코드북 collapse 문제를 효과적으로 완화할 수 있다는 것이다.
Experiments
Quantitative Evaluation 결과 아래 표와 같이 VQCT가 VQIM에 효과적임을 시사한다.
Qualitative Evaluation 결과 전반적인 시각 품질이 가장 우수하여 이미지 복원 성능을 크게 향상시킨다는 것을 추가로 입증하였다.
Ablation Study
PLM codebook 도입에 대한 VQIM ablation study 결과 PLM 코드북 도입은 VQIM에 매우 효과적임을 입증한다.
코드북 transfer network효과에 대한 ablation study 결과 단순히 코드북을 가져다 쓰는 것보다, transfer network를 통해 학습 가능한 구조로 전환하는 것이 효과적이며, GCN 기반으로 구성하고 분리해 사용하는 것이 가장 효과적임.
VQCT method를 다른 VQIM에 대해 확장 실험 결과, VQ-VAE와 VQ-GAN 모두에서 성능이 크게 향상되었으며 특정 모델에만 국한되지 않고 다양한 VQIM 구조에 일반적으로 적용 가능함을 실험적으로 입증하였다.
PLM 별 codebook 실험 결과 CLIP 기반 코드북을 사용할 경우 가장 높은 점수를 보였다.
코드북 collapse 문제를 완화하였는지에 대한 시각화 결과 더 넓은 코드 분포를 사용함으로써 collapse 문제를 완화하는데 효과적임을 보여준다.
