사진의 75%를 가렸더니 더 잘 배웠다: MAE의 계산과 학습

사진의 네모 조각을 대부분 가리고 남은 부분만 보여 준다. 사람은 의자의 다리나 강아지의 귀를 보고 가려진 대상을 대략 짐작한다. 신경망에도 같은 문제를 주면, 정답 라벨 없이 유용한 시각 표현을 배울 수 있을까. Masked Autoencoders, 줄여서 MAE는 이 질문을 픽셀 복원이라는 간단한 작업으로 시험한다.

그런데 핵심은 가리는 행위 하나에 있지 않다. MAE에서는 정보를 더 많이 숨기는 선택과 계산을 더 적게 하는 선택이 맞물린다. 어려운 문제를 만들어 학습을 유도하면서, 큰 모델이 실제로 처리할 토큰은 오히려 줄인다. 이 두 효과가 한 설계에서 나온다는 점이 재미있다.

빈칸을 거대한 인코더에 넣지 않는다

이미지를 일정한 크기의 패치로 나누고 일부를 무작위로 선택한다. 인코더는 보이는 패치만 처리한다. 가려진 자리를 나타내는 mask token을 큰 인코더 입력에 함께 넣지 않는다. 이 선택 덕분에 사전 학습 중 인코더의 시퀀스가 짧아진다.

그다음 가벼운 디코더에서 보이는 패치의 특징과 가려진 자리의 mask token을 합친다. 위치 정보를 붙여 어떤 빈칸을 예측해야 하는지 알려 주고, 각 패치의 픽셀을 복원한다. 디코더는 전체 위치를 보지만 인코더보다 작게 만들 수 있다. 원문의 비대칭 구조는 이 계산 분담을 명확히 보여 준다.

이 구조를 단순히 가벼운 오토인코더라고 부르면 목적을 놓치기 쉽다. 최종적으로 필요한 자산은 주로 인코더다. 디코더는 학습 문제를 수행하기 위한 임시 장치이며, 인식 작업에서는 버린다. 복원기를 얼마나 화려하게 만들지가 아니라, 복원 문제를 통해 인코더가 어떤 정보를 배우게 할지가 중심이다.

실제 구현도 드문드문한 토큰을 처리하는 특별한 연산을 요구하지 않는다. 토큰 순서를 섞고 일부를 선택해 인코더에 넣은 뒤, 디코더 앞에서 빈칸을 추가하고 원래 순서를 복원할 수 있다. 빈칸을 생략하는 개념과 sparse 전용 커널을 써야 한다는 주장은 별개다.

이미지에서는 많이 가려야 문제가 어려워진다

문장의 단어 하나를 가리면 의미가 크게 달라질 수 있다. 이미지에서는 인접한 픽셀과 패치가 상당한 정보를 공유한다. 조금만 가리면 주변의 색과 질감을 이어 붙이는 것으로 문제를 쉽게 풀 수 있다. 논문은 높은 마스킹 비율, 대표적으로 75%가 효과적이라는 결과를 보여 준다.

예를 들어 196개 패치 중 49개만 보여 준다고 생각해 보자. 이 숫자는 비율을 설명하기 위한 계산이다. 어느 부분이 보일지 매번 달라지므로 모델은 한 번은 바퀴를 보고, 다른 번에는 차창을 보고 숨겨진 부분을 예측해야 한다. 단순한 국소 보간에만 의존하기 어려워진다.

동시에 인코더가 처리할 패치 수는 줄어든다. 학습 과제가 더 어려워지는데 계산이 더 쉬워지는 특이한 조합이다. 물론 마스킹 비율을 계속 높이면 반드시 좋아지는 것은 아니다. 너무 적은 단서만 주면 예측할 정보 자체가 부족해진다. 논문은 마스킹 비율을 바꾼 실험으로 이 선택을 확인한다.

또한 무작위 패치 마스킹과 커다란 영역을 통째로 가리는 방식은 서로 다른 문제를 만든다. 분산된 단서는 물체 여러 부분을 조금씩 보여 주고, 큰 빈 영역은 주변 문맥에서 전체 부분을 추측하게 한다. 가리는 방식은 단순한 데이터 처리 옵션이 아니라 학습할 정보의 종류를 정하는 설계다.

손실은 보이지 않았던 곳에서 계산한다

디코더는 픽셀을 예측하고, 가린 패치에서 원본과의 평균제곱오차를 계산한다. 보이는 부분을 그대로 복사하는 데 점수를 주기보다, 숨겨진 정보를 추론하는 데 학습 신호를 쓴다. 논문은 패치별로 정규화한 픽셀을 목표로 사용하는 변형도 비교한다.

픽셀 MSE는 이해하기 쉽지만 의미 자체를 직접 측정하지는 않는다. 강아지 털 색을 맞히는 것과 강아지를 알아보는 것은 다른 목표다. 그래서 복원 이미지가 보기 좋아졌다는 사실만으로 인코더가 좋은 인식 표현을 배웠다고 결론 내릴 수 없다. 반대로 복원 결과가 흐릿해도 downstream 인식에는 유용한 정보가 있을 수 있다.

사고실험으로 가려진 의자 좌석을 생각해 보자. 보이는 다리와 등받이만으로 좌석의 존재는 짐작할 수 있어도, 원래 천의 무늬는 알기 어렵다. 제곱오차를 줄이는 예측이 여러 가능성의 평균에 가까워지면 무늬가 흐릿해질 수 있다. 이 설명은 손실의 일반적 성질에 대한 직관이며, MAE가 내부에서 확률분포를 명시적으로 출력한다는 뜻은 아니다.

따라서 이 논문은 복원 사진의 감상보다 학습한 표현을 어떻게 평가했는지 따라 읽는 편이 좋다. 어떤 학습 문제가 어떤 능력으로 이어졌는지가 핵심 질문이다.

잘 배웠는지는 인코더를 꺼내 본 뒤 알 수 있다

저자들은 사전 학습 후 인코더를 사용해 ImageNet 분류와 여러 전이 작업을 평가한다. 전체 모델을 미세 조정하는 방식과, 인코더를 고정하고 선형 분류기만 학습하는 방식은 서로 다른 성질을 본다. 전자는 새 작업에 맞게 표현을 바꿀 수 있는 출발점의 품질이고, 후자는 현재 특징에서 정답이 얼마나 쉽게 분리되는지에 더 가깝다.

그래서 두 평가를 하나의 정확도로 합치면 안 된다. MAE가 좋은 사전 학습이라고 할 때, 어떤 평가에서 얼마나 좋아졌는지 확인해야 한다. 특히 큰 ViT를 처음부터 지도학습하는 설정의 최적화와 정규화도 비교 결과에 영향을 준다. 모델 이름만 같다고 모든 학습 조건이 같은 것은 아니다.

라벨을 사용하지 않는 사전 학습과 라벨 없는 최종 시스템 역시 다르다. 이후 분류나 검출 작업에서는 정답 데이터가 쓰일 수 있다. MAE의 기여는 그 데이터가 없던 단계에서 이미지를 활용해 좋은 초기 표현을 준비하는 데 있다.

내 데이터에 적용할 때 생기는 선택

촬영된 상품이나 제조 영상이 많이 쌓여 있다면, MAE는 라벨을 붙이기 전에도 그 데이터를 학습에 활용할 아이디어를 준다. 다만 원래 작업이 아주 작은 흠집을 찾는 것이라면, 큰 패치나 입력 축소로 그 정보가 사라지지 않는지 먼저 확인해야 한다. 전체 물체를 알아보는 과제와 미세 결함을 찾는 과제는 필요한 해상도가 다르다.

내가 권하고 싶은 비교는 작은 지도학습 기준 모델과 MAE 초기화 모델을 같은 라벨 예산으로 평가하는 것이다. 복원 손실 감소만 보고 도입을 결정하지 않는다. 촬영일이나 제품 종류를 분리한 평가에서 실제 전이 이득이 남는지 확인하면 더 설득력 있다. 이것은 논문의 직접 결론이 아니라 적용을 위한 실험 제안이다.

학습 비용을 계산할 때도 주의할 점이 있다. 사전 학습에서는 일부 토큰만 인코더에 넣지만, 최종 인식에서는 보통 전체 이미지를 사용한다. 사전 학습의 속도 개선을 최종 추론 속도 개선과 같은 비율로 기대해서는 안 된다.

MAE는 빠진 정보를 채우는 문제가 결국 남아 있는 정보를 어떻게 읽을지 가르칠 수 있다는 사실을 보여 준다. 더 많은 라벨을 만드는 대신, 이미 있는 이미지에서 정답을 만들고 계산을 필요한 곳에 배분한다. 빈칸을 채우는 디코더보다, 빈칸이 있을 때도 장면을 읽는 인코더가 남는다는 관점으로 읽으면 설계의 의도가 훨씬 선명해진다.

논문 정보: Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, Ross Girshick, Masked Autoencoders Are Scalable Vision Learners, 2021 공개, CVPR 2022. 논문 원문 · 공식 구현