오늘은 In Pursuit of Pixel Supervision for Visual Pre-training이라는 논문에 대해서 리뷰해보려고 한다. 아이디어부터 실험설계까지 전부 정말 좋은 논문이어서 강추한다. 특히 Introduction과 Related Work부분에서 Self-supervised learning의 역사를 compact하게 정말 잘 정리해놓았다고 생각한다. 이 내용을 읽는 것만으로도 큰 가치가 있다고 생각한다.
https://arxiv.org/abs/2512.15715
In Pursuit of Pixel Supervision for Visual Pre-training
At the most basic level, pixels are the source of the visual information through which we perceive the world. Pixels contain information at all levels, ranging from low-level attributes to high-level concepts. Autoencoders represent a classical and long-st
arxiv.org
Background
최근의 self-supervised learning 방법들은 objective(loss)가 어느 공간에서 정의되느냐에 따라 크게 두 가지 범주로 나뉘어 왔다. 첫 번째는 MAE와 같이 입력 이미지의 일부 정보를 제거한 뒤 원본 이미지를 복원하도록 학습하는 방식으로, loss를 pixel level에서 정의하는 방법이다. 두 번째는 CLIP이나 DINO처럼 서로 다른 view에서 얻은 동일 객체의 representation이 일치하도록 학습하는 방식으로, loss를 latent space에서 정의한다.
오늘날 pretraining 분야에서 사실상 가장 강력한 방법으로 여겨지는 것은 단연 DINO이며, 이로 인해 최근까지는 latent space에서 loss를 거는 방식이 pixel level objective보다 본질적으로 우수하다는 인식이 강하게 자리 잡아왔다.
그러나 이 논문이 주장하는 핵심은, pixel level에서 loss를 정의하는 접근 자체가 열등한 것이 아니라, 기존 MAE가 여러 측면에서 sub-optimal한 설계를 가지고 있었을 뿐이라는 점이다. 이에 저자들은 MAE의 성능을 제한하던 주요 요인들을 하나씩 개선함으로써, DINO와 견줄 수 있는 수준의 SOTA 모델을 제시한다.
이를 통해 이 논문은 pixel level loss 역시 여전히 매우 강력하고 유효한 self-supervised learning 전략임을 설득력 있게 입증한다.
Method
MAE에 대해 항상 반복적으로 제기되어 온 비판이 있다.
“MAE에서 추출한 encoder feature는 semantic한 정보를 충분히 담고 있지 못하고, low-level detail에 지나치게 치중한 것 아니냐”는 것이다. 그렇다면 그 원인은 크게 두 가지 가능성으로 나눠볼 수 있다. 첫째, pixel-level reconstruction loss 자체가 모델로 하여금 low-level detail 위주의 표현을 학습하도록 강제했기 때문일 수 있고, 둘째, high-level semantic 정보를 학습할 수 있음에도 불구하고, 모델 구조나 학습 설계가 이를 충분히 끌어내지 못했기 때문일 수도 있다.
저자들의 입장은 후자에 가깝다. 물론 pixel-level reconstruction loss를 사용하는 이상, MAE가 다른 latent-level self-supervised 방법들에 비해 low-level detail에 더 많은 비중을 둘 수밖에 없다는 점은 인정한다. 그러나 그럼에도 불구하고, MAE가 high-level semantic 정보를 충분히 담고 있지 못하다고 느껴지는 것은 loss의 한계라기보다 설계상의 문제일 가능성이 크다고 본다. 그래서 이 논문에서는 MAE의 설계상의 문제를 해결하기 위해, 원래의 MAE architecture에서 4가지 부분을 수정한다. 하나하나 같이 봐보도록 하자.
Deeper decoder
일반적으로 우리는 인코더 layer의 깊이에 따라 표현이 점진적으로 추상화되기를 기대한다. 즉, 얕은 layer에서는 low-level 정보가, 깊은 layer로 갈수록 high-level semantic representation이 형성되기를 기대한다. 그러나 저자들은 MAE의 경우, decoder가 지나치게 얕게 설계되어 있어 encoder의 후반부가 semantic abstraction에 도달하기 전에 reconstruction을 잘 수행하기 위한 low-level 표현을 학습하고 있는 것은 아닐까라는 의심을 제기한다. 다시 말해, 1부터 N까지의 layer를 통해 점진적으로 high-level 정보를 학습하기를 기대했지만, 실제로는 encoder의 앞부분(예: layer 1~ layer 0.75N)에서만 의미 있는 표현 학습이 이루어지고, 이후의 layer들은 reconstruction을 보조하는 역할로 수렴했을 가능성이 있다는 것이다.
이에 대한 하나의 근거로 저자들은 위의 figure를 제시한다. 이 figure는 MAE encoder의 서로 다른 layer에서 추출한 feature를 각각 downstream task에 사용했을 때의 성능을 비교한 결과를 보여준다. 여러 downstream task에서 일관되게 encoder의 마지막 layer보다 중간 layer의 feature가 더 높은 성능을 보인다는 점을 확인할 수 있다. 이는 encoder의 후반부 layer들이 low-level과 high-level 정보를 균형 있게 담은 표현을 학습하기보다는, reconstruction을 잘 수행하기 위한 low-level detail에 치중하고 있을 가능성을 시사한다.
이러한 분석을 바탕으로 저자들은 reconstruction과 관련된 역할을 decoder 쪽으로 보다 명확히 이전하고, encoder의 후반부가 이미지의 semantic feature를 추출하는 데 집중하도록 유도하자는 접근을 제안한다. 이를 위해 decoder의 깊이를 늘려 reconstruction capacity를 강화하고, 그 결과 encoder가 semantic representation을 학습할 수 있는 여지를 확보하려는 것이다.
물론 이러한 설계는 매우 신중하게 조절되어야 한다. 원래 MAE에서 decoder를 작게 설계한 이유는, 최종적으로 downstream task에서 활용되는 것은 encoder이기 때문에 encoder가 최대한 많은 정보를 학습하도록 하기 위함이었기 때문이다. 그럼에도 불구하고, 저자들이 지적하듯, decoder를 지나치게 축소할 경우 오히려 encoder의 후반부가 reconstruction을 직접 떠안게 되어, 결과적으로 semantic representation 학습에 불리하게 작용할 수 있다는 점은 기존 MAE 설계에서 간과된 부분이었다.
Decoder Depth를 늘려서 MAE를 학습하고, downstream task에서 MAE encoder feature의 성능을 평가한 지표는 다음과 같다. 상당한 성능 향상이 있음을 알 수 있다.
Larger Mask Block
원본 MAE에서는 모든 patch를 독립적으로 random masking한다. 그러나 이러한 방식에는 분명한 한계가 있다. patch 단위가 지나치게 작을 경우, 각각의 visible patch가 담고 있는 정보가 매우 국소적이어서 local semantic을 파악하기 어렵다는 점이다.
이러한 문제의식에서 출발하여, 저자들은 masking 단위를 단일 patch가 아니라 4×4 정도의 patch block 단위로 확장하면 모델이 보다 의미 있는 지역적 문맥(local semantic)을 관찰할 수 있다고 볼 수 있다고 주장한다. 즉, encoder가 입력으로 관측하는 정보가 단순한 저수준 시각 신호를 넘어, 일정 수준의 구조와 의미를 포함하게 되면서 semantic representation을 학습하기에 보다 유리한 조건이 형성된다는 것이다.
다만 이 접근은 신중하게 다루어져야 한다. MAE의 핵심은 여전히 충분히 어려운 reconstruction task를 통해 의미 있는 표현을 학습하는 것이기 때문에, 문제 자체가 지나치게 쉬워져서는 안 된다. 따라서 일정 수준 이상의 masking ratio를 유지하는 것이 필수적인데, 이때 masking block의 크기를 과도하게 키우면, 모델이 이미지의 극히 제한된 영역만 관찰한 채 나머지 대부분을 예측해야 하는 상황이 발생할 수 있다. 이 경우 학습은 실제 관찰에 기반하기보다는 거의 ‘상상’에 가까운 추론으로 흐를 위험이 있다.
따라서 masking 단위를 확장하는 것은 local semantic을 강화하는 데에는 효과적일 수 있지만, 동시에 관측 정보와 예측 난이도 사이의 균형을 세심하게 조절해야 하는 요소임을 알 수 있다.
실험 결과는 다음과 같다.
More [CLS] Tokens
CLS token은 visual patch들 사이에서 전역적인 정보를 주고받는 일종의 hub 역할을 수행하는 것으로 잘 알려져 있다. 이러한 특성 덕분에, CLS token에 별도의 supervision이나 loss를 직접 걸지 않더라도, image classification과 같은 downstream task에서 널리 활용되어 왔다.
그러나 최근 연구들에서는 하나의 CLS token만으로는 충분하지 않다는 지적이 반복적으로 제기되고 있다. 여러 visual patch들 간의 정보 교환을 집약하고, 이미지 전체의 전역적인 semantic을 하나의 token에 모두 담기에는 표현 용량 측면에서 한계가 있다는 것이다.
이에 따라 일부 연구들은 CLS token의 개수를 늘리는 방식을 제안해왔다. 여러 개의 CLS token을 두어 전역 정보를 분산적으로 수집하게 하고, downstream task에서는 이들 CLS token을 average하거나 aggregation하여 최종 representation으로 사용했다. 이러한 접근은 단일 CLS token에 과도하게 정보가 집중되는 문제를 완화하고, 보다 풍부한 전역 semantic representation을 형성하는 데 도움이 될 수 있다.
따라서 trend에 맞게 이 논문도 CLS Token의 개수를 늘리는 방법을 사용했다.
cls token 개수에 따른 실험 결과는 다음과 같다.
Web-Scale Data and Curation
기존 MAE에서 사용되던 데이터셋은 semantic representation을 학습하기에는 규모와 다양성 측면에서 한계가 있다고 판단되어, 저자들은 보다 큰 규모의 데이터셋을 새롭게 구축했다. 먼저 웹에서 대용량 데이터를 크롤링하고, 이후 이 데이터셋을 curate한다. curation 과정에 대해 설명하면, raw data를 이용해 개선된 MAE 모델을 먼저 학습한 뒤, 학습 과정에서 loss가 지나치게 작은, 즉 모델이 너무 쉽게 복원할 수 있는 샘플들을 제거하였다. 또한 흰 배경에 검은 글씨가 있는 이미지처럼, 텍스트 위주의 단순한 구조를 가진 이미지들이 semantic 학습에 크게 기여하지 않는다고 보고, 이미지 내 색상 분포의 다양성을 기준으로 필터링을 수행했다. 구체적으로는 color entropy가 낮은 이미지들을 제거함으로써, 보다 풍부한 시각적 변화를 포함한 데이터만을 남겼다. 이와 같이 정제된 데이터셋을 최종적으로 사용하여, 효과적으로 MAE 모델을 학습했다고 한다.
Experiments
여러 실험들을 통해서, 개선된 MAE(Pixio)가 downstream task에서 dino에 버금가거나 더 좋은 성능을 내고 있음을 알 수 있다.
Comments
pixel level loss가 아직 죽지 않았다는 것을 알려주는 좋은 논문이라고 생각한다.