[CVPR 2024] OMG-Seg: Is One Model Good Enough For All Segmentation?

안녕하세요. Vision 분야에서 Recognition에 대해 Detection에 대해 연구를 하고, 최근에는 Segmentation 분야에 대해 관심을 가지고 있습니다. 제 생각에선 Detection이든, Segmentation이든, 그 어떤 Task이든 연구와 실용성의 관점에서 차이점은 미미하되 Segmentation 코드도 다뤄보고 싶어 본 논문을 찾게 되었습니다. 정확히는 연구를 진행하며 Open-World, Open-Vocabulary 태스크에 주 관심이 있지만, 결국 생각해보면 해당 태스크들에서 추구하는 바는 “실 생활에서 활용 가능한 Vision 기술”에 해당합니다. 그러한 관점에서 예전부터 “Supervised에서 시작된 Vision 연구에서 연구자들이 새로운 문제를 찾고, 그 문제를 해결하다 보니 또 새로운 문제를 찾고, 그렇게 발전되어 가다 보면 결국 모든 연구들이 이상하는 바는 궁극적인 하나의 모델을 만들어 내는 것”이라고 생각을 하였는데, 본 논문이 하나의 모델로 모든 Segmentation 태스크를 수행하고자 하는 시작점으로 보입니다.

Introduction

본 논문의 Architecture는 위 Mask2Former를 기반으로 구성되어 있습니다. 본 논문 이전에도 Mask2Former라는 논문에서 Segmentation에서 Semantic, Instance, Panoptic을 한 번에 진행하려는 시도가 존재하였습니다. Mask2Former 논문에 대해 짧게 이야기하자면 해당 논문에선 SOTA specialized architecture, 즉 특정 태스크를 목적으로 하는 모델들은 ㅈ해당 태스크에서 최상의 성능을 보이나, 하나의 태스크에서만 활용 가능하다는 한계가 존재하며 성능 향상을 위해 모델 구조가 고도화되다 보니 모델 학습을 위해 비용이 크게 증가되며 학습 시 모델의 수렴에도 긴 시간이 필요하다는 단점이 존재합니다. 특히, 앞서 말한 “하나의 태스크에만 국한됨”의 한계를 극복하고자 MaskFormer라는 구조도 소개되었으나, 해당 모델은 특정 태스크에서 (Instance의 경우) 아쉬운 성능을 보입니다. Mask2Former에서는 핵심적으로 아래 그림의 masked attention 구성으로, mask r-cnn에서 mask prediction 내의 영역에 대해서만 attention을 수행하는 방식을 통해 segmentation 분야에서 메모리 소비를 줄이며 학습 수렴을 빠르게 하여 좋은 성능을 보였습니다.

본 논문에서도 동일한 입장을 보입니다. 통합된 하나의 모델이 Segmentation에서 가지 뻗은 모든 태스크에 대응할 수 있다면, 심지어 Mask2Former에선 하지 않은 Video segmentation까지, 특정 태스크에 특화된 모델 구조를 가지지 않아도 되며 이는 실용성 측면에서 매우 합리적임을 언급합니다. 그렇게 하기 위한 모델 구성에 있어, DETR은 태스크 통합에 있어 굉장히 유망한 구조이며, 최근의 연구들은 DETR을 기반으로 하여 Closed-set과 Open-set, 즉 Open-vocabulary segmentation까지 포괄할 수 있는 모델 구성으로 좋은 성능을 보입니다. 하지만 이 또한, Mask2Former가 당대에는 각 태스크에 특화된 모델에 비해 좋은 성능을 보였을지언정, 이후 등장한 SOTA 성능은 훨씬 높은 성능을 보이곤 합니다. 결국 하나의 통합된 모델이 그간 활발히 연구되지 않은 결정적 이유는 “성능”이라는 측면에서 설득시키기 어려웠기 때문입니다. 본 논문에서 저자는 OMG-Seg, One Model Good Enough, 구조를 설계하였으며, 뒤 방법론에서 한번 더 언급하겠지만 모델 구성은 Mask2Former를 대부분 따르되 Video segmentation까지, 그리고 Image와 Video, 그리고 Interactive segmentation을 한 번에 학습한 측면에서 Contribution을 주장합니다. 해당 모델이 Segmentation에 있어 어느 수준까지 커버하는지, 아래의 표로 확인해보겠습니다. 쉽게 말하자면, Segmentation에 있어 저도 처음 들어본, 정말 모든 태스크를 커버하네요.

Methodology

통합된 모델 구성을 위한 핵심은 Task-specific한 요소는 최소화하면서도 태스크 간 공유되는 파라미터의 수는 늘리는 전략입니다. 저자의 목적은 최상의 성능을 추론함이 아닌, 모델의 범용성을 확대함에 있으며 위 예시에서 본 표와 같이, Video분야의 Segmentation까지 책임짐에 있습니다. 저자가 설계한 핵심 요소는 DETR 구조에서의 Object query를 활용하며, 이 Query가 태스크에 고유한 엔터티들을 표현하는 단일 객체로 활용될 수 있다는 점입니다. 하나의 예로, Video segmentation은 Image segmentation에 비해 추가적인 하나의 ID (프레임)만 필요로 합니다. 해당 사실은 하나의 Decoder가 각 Object query와 관련된 Feature에 대한 연산을 대부분 공유하여 처리할 수 있음을 의미하며, 다음 구조에서 Cross-attention layer (Mask2Former에서 활용된)에 입력되는 Feature만 차별점이 있습니다. 즉, 해당 논문에서 말이 어렵게 적혀져 있는데 한 마디로 “각 태스크들이 결국 큰 관점에서 보면 한, 두 가지의 세세한 차이점만 있으므로 Encoder-Decoder에서 대부분을 공유할 수 있음”을 의미합니다.

모델 구조를 보면, CLIP의 Visual Encoder를 Freeze한 상태로 Feature map을 추출합니다. 이후 Pixel Decoder로 1/8, 1/16, 1/32의 Multi-scale feature map을 생성하며, 이 때는 Image 또는 Video에 대해 해당됩니다. CLIP을 Visual Encoder로 택한 이유는, Open-vocabulary에 대응하기 위함입니다. DETR에서 차용한 개념의 Semantic query Q_{obj}^{s} 와 Interactive segmentation (e.g., SAM)을 위해 설계된 Location query에 해당하는 Q_{obj}^{l} 를 가진 채, Mask Decoder의 입력으로 들어가게 됩니다. 위 (a)와 (b)의 구성은 모두 Mask2Former에서 차용하였습니다. 다만 CLIP의 Visual backbone을 활용하며, Interactive에 대한 Prompts Encoder를 활용하여 Location query를 생성한다는 점을 제외하면 나머지는 모두 동일합니다.

Mask Decoder로 넘어와서, 결국 핵심은 맨 처음의 Masks (Binary Classification으로 구할 각 클래스에 대한 Binary Mask: e.g. Mask RCNN)에 대한 Feature가 Key,Value로 들어가게 되며 (Location+Semantic)Object query와의 Cross-attention 연산을 진행합니다. mask-attention의 결과로 정제된 query로 (c)의 연산이 시작됩니다. (c)에서는 Object query들이 예측하는 값이 결국 특정 태스크에 따라 다양해지며, Image의 경우 Semantic, Instance, Panoptic에 따라 일부 달라지며 Video의 경우에는 Image에 추가적인 프레임 ID와 예측 시 Time-series를 활용한다는 점 정도만 차이가 있겠습니다. 특히, Interactive segmentation의 경우 사용자가 클릭한 지점 또는 박스에 대해 해당 지역을 입력으로 하므로, 그 지점에 대해 연관된 영역에 대한 예측만을 진행한다는 차이가 있습니다.

Architecture는 Image의 경우 Mask2Former를 그대로 활용한 모습이며, Video의 경우에도 예측하는 점만 다르게 활용할 뿐 (물론 그걸 위해 앞단에서 Visual Backbone을 변경하거나, 뒤 편에서 각 태스크에 대한 Loss를 Classification, mask Cross Entropy, Dice Loss로 통합하여 활용하는 등의 차이는 있습니다) 큰 특이점을 보진 못했는데, 학습 시가 특이합니다. 저는 하나의 태스크에 대해 학습하고, 다른 태스크에 대해 학습하고, 그들을 어떻게 합칠까에 주목해서 보았는데 (Pretraining 하는 방식으로), Co-training, 즉 모든 태스크를 합쳐 한 번에 학습시킨다는 점입니다. 이는 Entity(Query)가 크게 보았을 때 세 개의 다른 태스크 (Image, Video, Interactive)에 대해 하나로 치환될 수 있기 때문입니다.Entity는 Thing, Stuff, Clas-agnostic mask, 어떤 것이든 될 수 있으며 학습 시에는 헝가리안 매칭으로 학습되기에 태스크에 구애받지 않아도 괜찮다고 본 저자의 시선이 조금 새로웠습니다. 다 설명하고 나니, 결국 모델 자체가 Mask2Former다 보니 모델을 보려 했는데 그냥 Transformer 기반의~ 요즘 모델들이 다 이런 Architecture만 취하지 않냐하는 생각도 듭니다.

Experiments

위 문단에서 언급한 바와 같이, 저자는 Co-training을 위해 COCO Panoptic, COCO SAM, VIP-Seg, Youtube-VIS-2019의 Image, Interactive, Video 데이터 셋을 함께 활용합니다. Open-set에 대한 성능 측정을 위해 ADE-20k, DAVIS-2017, Youtube-VIS-2021 데이터 셋을 평가 시 활용하며 궁금할 수 있는 COCO SAM이란 groud-truth bbox와 mask에 대한 center points가 visual prompts로 활용되는 데이터 셋입니다. 위 데이터 셋들은 후에 연구에 필요 시 한번 쯤 자세히 확인해봐야겠네요.

첫 번째는 System-level comparison입니다. 아래 표의 Mask2Former를 Architecture로 활용하고 있지만, 해당 모델은 Video와 Interactive 영역에서 활용하고 있지 못하기 때문에, 본 논문에서의 성능이 물론 Mask2Former에 비해 떨어지지만, 저자는 이 성능에 대해선 첫 문단에 언급한 바와 같이 “최적의 성능이 아닌 모든 태스크를 할 수 있음”을 강조합니다. 또한, 타 방법론들의 최적의 Epoch는 나와있지는 않지만, 저자는 학습 시 24Epoch만을 학습하였다고 설명합니다. (24Epoch이라 해서 짧은 시간은 아닐껍니다. GPU 수와 Batch 수, 그리고 데이터의 수를 생각하면 Iteration으로는 몇 만이 될테니까요). 저자가 강조하는 다음 문장을 보면, 본인 모델의 성능이 떨어짐에도 어떻게 문제 제기를 잘 하여 (특정 모델은 특정 태스크에만 국한되어 있음) 이를 설득시키고자 하였는지 알 수 있습니다.

This ability to effectively handle all forms of segmentation, despite the small perfor- mance gaps noted, reinforces our assertion that OMG-Seg is a robust and versatile model suitable for diverse segmen- tation scenarios.

흥미로운 Ablation 두 가지 정도를 들고 왔습니다. 저자는 Co-training 방식의, 모든 데이터 셋을 동일하게 학습하는 방식을 가져왔는데, 이 결과를 통해 예를 들어 (a), Panoptic segmentation 데이터 셋으로 학습한 모델은 COCO-PS, Panoptic segmentation에서는 co-training에 비해 좋은 성능을 보이지만 다른 태스크에서의 성능은 떨어지므로, 저자는 a + b + c, 즉 데이터도 많고 여러 태스크에 일반적으로 좋은 성능을 보일 수 있는 방안을 모색하였다고 합니다.

또한, Co-training 외에도 Decoder를 공유한다는 점이 저자의 핵심 Contribution인데, 이 때의 경우도 Image와 Video에 대해 Decoder가 따로 존재한다면 COCO-PS의 경우에는 좋은 성능을 보이지만, Video까지 비교하면 전체적인 평균 성능은 Shared decoder가 더 좋은 성능을 보일 뿐만 아니라, 파라미터 수에서 Shared의 특성 때문에 더 적은 파라미터 수를 가지며, 이는 또 다른 말로 태스크의 입력이 늘어남에 따라 파라미터 수는 계속 증가해야만 하는 상황을 방지할 수 있다고 합니다. 해당 실험 외에는 정성적 결과나, 백본 변경에 따른 성능 차이만을 리포팅하므로 정성적 결과를 마지막으로 살펴보고 마무리하도록 하겠습니다. 아직은 Segmentation 태스크에 대한 논문을 읽은 적이 처음이여서 아무래도 설명의 깊이가 다소 낮은 듯 합니다. 해당 지점 보완하여 더 좋은 리뷰를 작성할 수 있도록 해야겠습니다.

Author: 이 상인

5 thoughts on “[CVPR 2024] OMG-Seg: Is One Model Good Enough For All Segmentation?

  1. 안녕하세요 좋은 리뷰 감사합니다.

    결국 해당 방법론을 살펴보았을때 모든 데이터셋의 샘플을 가져와 학습 때 끊김 없이 쭉 모두 학습한 것이고 그걸 co-training이라고 칭한 것이 맞나요?
    그렇다면 사실 못해도 6개 이상 되는 태스크들의 기존 연구에서 고정적으로 사용하는 loss들이 각기 다를 것 같은데, 그런 점을 모두 고려하여 학습 중 적용하는 loss가 계속 변경되나요? 예를들어 coco panoptic 데이터셋으로 ps을 학습하다가 youtube vis 데이터셋으로 vis을 학습하는 그 전환점에서 적용하는 loss도 동시에 변경되는 것인지 궁금합니다. 물론 헝가리안 매칭이 기본으로 수행되겠지만 각 태스크의 기존 베이스라인에서부터 적용되던 loss가 각기 다를 것 같아 여쭤보게 되었습니다.

    1. 안녕하세요. 리뷰 읽어주셔서 감사합니다.
      Co-training이란 모든 태스크에 대한 학습을 한 번에 진행하는, 말씀해주신 바처럼 모든 데이터 셋의 샘플을 가져와 끊임 없이 쭉 학습하는 것도 동일한 표현입니다.
      Loss의 경우, 본문에서 말씀드린 바와 같이 CrossEntropy, Binary Cross Entropy (Mask), Dice Loss 세 개만을 활용합니다. 이는 이미지든, 비디오든 동일하게 Decoder에서 Object query가 하나의 개념으로 추후 활용되므로 Loss에 대해 각각 변형되지 않고 하나의 loss로 학습됩니다.

  2. 리뷰 잘 읽었습니다.

    method 초입부, 본 구조 설계의 핵심에 대해 “Task-specific한 요소는 최소화하면서도 태스크 간 공유되는 파라미터의 수는 늘리는 전략” 라고 표현해 주셨습니다. Task-specific 한 요소를 최소화하고 대부분의 parameter들을 agnostic하게 간주하는 것 같은데, 물론 이게 “범용성” 이라는 측면에서는 이론적으로 맞아보이긴 합니다만 각 task 별로 성능적으로 살펴봤을때는 그래도 specific 한 parameter 들을 개별적으로 어느정도 가져야 좋을 거 같긴 합니다. task간 공유되는 파라미터 수를 늘리는 전략에 대해서 뭔가 추가적인 언급이라던가 실험이 있을까요?

    감사합니다.

  3. 안녕하세요 상인님 리뷰 감사합니다.

    제목을 보고는 여러개의 모델이 필요해서 합쳐야 한다는 느낌인것 같았는데, 저 질문의 답이 Yes라고 제시하는 모델이었네요,, 하나의 모델로 거의 모든 segmentation task를 수행할 수 있고 그 비결이 Mask2Former구조와 Co-training이라고 이해했습니다!!

    헝가리안 매칭에 대해 좀 찾아봤지만, 헝가리안 매칭으로 학습되기에 태스크에 구애받지 않아도 괜찮다고 본 저자의 시선이 새로웠다고 하신 부분은 제가 속뜻을 이해 못 해서 혹시 어떤 이유에서 새로우셨는지 궁금합니다!!

    1. 안녕하세요.

      음, 본 질문에 대해서는 다소 제 사적인 견해가 포함될 수 있습니다.
      본 헝가리안 매칭에 대한 저자의 시선이 새롭다고 생각이 든 이유는, “저자가 그렇게 의도하였기” 때문입니다.
      어려운 말로 들릴 수 있는데, 다음의 실제 사례를 한 번 생각해보곘습니다.
      우리가 흔히 말하는 Language Embedding, 음, 또 다른 예로는 Prototype Vector Layer는 실제로 모델 구성이 복잡하고 어려울까요? 그렇지 않습니다. 해당 모듈, 네트워크의 구성은 코드 몇 줄의 Linear Layer이거나, 때로는 단순히 nn.Embedding()과 같은 변수일 때도 있습니다. 다만, “저자가 해당 코드에 어떤 의미를 부여하여, 그들이 어떻게 동작하길 원하며 그렇기에 그들을 위해 어떠한 손실함수를 활용/설계” 하였는지가 중요합니다.
      조금 더 쉬운 예시로, 우리가 차원을 1000->10으로 만드는 FC layer를 하나 만들고, 그들에게 “너는 분류를 위한 모듈이니, 10개에 대해 합이 1로 되게끔하여 GT와 비교하도록 해”라고 의도하여 손실함수를 만든다면, 이것은 우리가 매번 쓰는 Classifier와 CrossEntropy가 되겠죠.

      즉, 제 사견에는 문제 해결이나 새로운 접근을 위한 정의가 굉장히 중요하다 생각듭니다. 어떻게 의도하였는가가 핵심이어야 한다는 점이죠. 질문의 예시로 돌아와, 헝가리안 매칭은 제가 알기론 1990년대의 굉장히 예전의 매칭 알고리즘입니다. 헝가리안 매칭이 활용된 논문은 대표적으로 DETR이 있죠. 이 때, 헝가리안 매칭을 어떤 시선으로 보아 접근할 것인지에 대해 저자는 DETR 구성의 query들을 특정 태스크에 하나씩 설계한 점이 아닌, 그 태스크들이 어떠한 공통점이 있기에 그 공통점까지 (예를 들어 A->B->C, D->B->C 라면 A와 D는 C로 가기 위해선 B를 거쳐야 한다는 공통점) 동일하게 여길 수 있다는 접근입니다.
      아마 제 위의 문장에 대한 이해가 어려울 수 있습니다. 제 사견이기 때문이나, 때로 논문을 읽고 실로 대단해보이는 논문들의 코드를 보았을 때는 그 부분이 단 코드 몇 줄에 불과할 때가 있습니다.

답글 남기기

이메일 주소는 공개되지 않습니다. 필수 항목은 *(으)로 표시합니다