Mask-aware inference with State-Space Models

이 논문은 State-Space 모델인 Mamba 아키텍처가 무작위 형태의 결손 데이터를 처리할 수 있도록 '부분 합성 (Partial Convolutions)'의 원리를 도입한 새로운 구성 요소인 '부분 비전 마바 (Partial Vision Mamba, PVM)'를 제안하고, 이를 통해 깊이 완성, 이미지 인페인팅, 결손 데이터 분류 등 다양한 작업에서 효과성을 입증합니다.

핵심

이 논문은 **"눈이 가려진 상태에서도 세상을 똑똑하게 이해하는 새로운 AI"**에 대한 이야기입니다.

기존의 AI(컴퓨터 비전) 는 완벽한 사진만 보고 학습합니다. 하지만 현실 세계는 다릅니다. 안개가 끼거나, 센서가 고장 나거나, 누군가 얼굴을 가렸을 때 데이터가 일부 누락되거나 '잘못된 값'으로 채워져 있습니다. 이런 불완전한 데이터를 기존 AI 에게 주면, AI 는 그 잘못된 값까지 진짜 정보인 줄 알고 혼란에 빠집니다.

이 논문은 **Mamba(마마)**라는 최신 AI 기술에 **'가려진 부분을 무시하고 진짜 정보만 골라보는 능력'**을 심어주었습니다. 이를 **PVM(부분적 비전 마마)**이라고 부릅니다.

이 내용을 일상적인 비유로 설명해 드릴게요.

---

1. 문제 상황: "눈이 가려진 그림자 놀이"

상상해 보세요. 여러분이 어두운 방에서 그림자 놀이를 하고 있습니다.

• 기존 AI (CNN 이나 일반 Mamba): 그림자가 비친 벽을 보는데, 누군가 손으로 가린 부분 (검은색) 을 그냥 '검은색 벽'이라고 믿고 전체 그림을 그립니다. 그래서 가려진 부분까지 검은색으로 칠해버려서 얼굴이 뭉개지거나 엉뚱한 모양이 됩니다.
• 현실의 문제: LiDAR(거리 측정 센서) 나 안개 낀 사진처럼 데이터가 일부 뚫려 있거나 (구멍이 있거나), 0 으로 채워져 있는 경우가 많습니다. 기존 AI 는 이 '구멍'을 진짜 정보로 착각합니다.

2. 해결책: "가려진 부분을 무시하는 스마트한 눈 (PVM)"

이 논문이 제안한 PVM은 마치 스마트한 탐정과 같습니다.

• 마스크 (Mask) 라는 안경: 탐정에게 '어떤 부분이 가려져 있는지' 알려주는 안경을 씌워줍니다.
• 진짜 정보만 골라먹기: 탐정은 가려진 부분 (검은색) 이 있는 것을 알기 때문에, 그 부분을 무시하고 가려지지 않은 진짜 정보 (얼굴, 나무, 건물 등) 만 집중해서 분석합니다.
• 결과: 가려진 부분이 있어도, 나머지 정보만으로 전체 그림을 완벽하게 복원하거나 분류할 수 있습니다.

3. 어떻게 작동할까? (세 가지 핵심 기술)

이 탐정 (PVM) 이 어떻게 작동하는지 세 가지 비유로 설명합니다.

① '구멍 난 퍼즐'을 맞추는 법 (패치 임베딩)

기존 AI 는 퍼즐 조각을 다 합쳐서 보는데, 조각이 하나 빠진 상태면 그 빈 공간도 퍼즐의 일부라고 착각합니다.

• PVM 의 방법: "여기 조각이 없구나!"라고 먼저 확인합니다. 그리고 빈 공간에는 '아무것도 없는 상태'를 의미하는 특별한 마커를 붙입니다. AI 는 이 마커를 보고 "아, 여기는 무시해야겠다"라고 학습합니다.

② '소문'이 퍼지는 방식 (상태 공간 모델)

Mamba 는 정보를 한 줄로 이어가며 처리합니다. (소문이 한 사람에서 다음 사람으로 퍼지는 것처럼요.)

• 기존 방식: 소문을 전할 때, 가짜 소문 (잘못된 데이터) 을 전하는 사람이 있으면 그 소문이 전체에 퍼져버립니다.
• PVM 의 방법: "이 사람은 가짜 소문만 전하네?"라고 판단하면, 그 사람의 소문을 차단하고 진짜 소문만 다음 사람에게 전달합니다. 그래서 전체 소문의 흐름이 깨끗하게 유지됩니다.

③ '가려진 부분 채우기' (Depth Completion & Inpainting)

• 깊이 측정 (Depth Completion): 안개 낀 도로에서 차의 거리를 재는데, 센서 데이터가 일부 뚫려 있습니다. PVM 은 뚫린 부분을 무시하고, 남은 데이터로 도로의 모양을 완벽하게 재구성합니다. (기존 AI 는 뚫린 부분을 0 으로 채워 도로가 갑자기 끊긴 것처럼 오해했습니다.)
• 그림 채우기 (Inpainting): 사진에서 얼굴 일부를 가려도, PVM 은 가려진 부분을 '무시'하고 나머지 얼굴 특징을 보고 눈, 코, 입을 자연스럽게 다시 그려냅니다.

4. 실험 결과: "왜 PVM 이 더 잘할까?"

논문은 세 가지 시험을 치렀습니다.

1. 깊이 측정 (Depth Completion): 안개 낀 도로에서 차의 거리를 재는 실험.
   • 결과: 기존 AI 보다 23% 더 정확해졌습니다. 마치 안개 속에서도 선명한 도로를 보는 것과 같습니다.
2. 그림 채우기 (Inpainting): 얼굴 사진의 일부를 지우고 다시 그리는 실험.
   • 결과: 가려진 부분 (머리, 턱선 등) 을 훨씬 자연스럽고 선명하게 복원했습니다. 기존 방식은 흐릿하거나 이상한 선이 생겼지만, PVM 은 실제 사진처럼 완벽하게 만들었습니다.
3. 사진 분류 (Classification): 얼굴 일부가 가려진 사진을 보고 "이게 누구야?"라고 맞추는 실험.
   • 결과: 정답률이 36%나 급상승했습니다. 눈만 보여도 "아, 이건 고양이구나!"라고 맞출 수 있게 된 것입니다.

5. 결론: "불완전한 세상에서도 완벽한 판단"

이 논문의 핵심 메시지는 다음과 같습니다.

"세상은 완벽하지 않습니다. 데이터는 항상 일부가 빠지거나 망가져 있습니다. 하지만 PVM이라는 새로운 기술을 쓰면, AI 는 그 '빠진 부분'에 일희일비하지 않고, 남은 진짜 정보만 믿고 훨씬 똑똑하고 정확한 판단을 내릴 수 있습니다."

이 기술은 자율주행차 (안개 낀 도로), 의료 영상 (부정확한 스캔), 보안 카메라 (가려진 얼굴) 등 현실 세계의 불완전한 데이터를 다루는 모든 분야에서 혁신을 가져올 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

• 배경: 실제 컴퓨터 비전 작업 (LiDAR 데이터, 센서 오류, 익명화 처리 등) 에서는 입력 데이터에 임의의 모양으로 결손되거나 유효하지 않은 영역 (invalid data) 이 존재하는 경우가 많습니다.
• 기존 한계:
  • CNN: 부분 합성곱 (Partial Convolutions, PConvs) 을 통해 유효한 픽셀만 고려하여 재정규화하는 방식으로 이 문제를 해결했습니다.
  • State Space Models (SSM, 예: Mamba): 최근 Transformer 를 대체할 수 있는 선형 복잡도의 고성능 아키텍처로 부상했으나, 추론 (inference) 시 임의의 모양을 가진 유효하지 않은 데이터를 처리할 수 있는 내재적 메커니즘이 부재합니다.
  • 핵심 문제: 기존 Mamba 기반 모델은 결손된 데이터를 '0'과 같은 플레이스홀더로 처리하여 모든 토큰을 유효한 것으로 간주합니다. 이로 인해 특징 추출이 왜곡되고, 단일 결손 토큰이 전체 시퀀스의 상태를 오염시켜 모델 성능이 급격히 저하됩니다.

2. 방법론 (Methodology)

2.1. 마스킹 인식 프레임워크 (Mask-aware Framework)

저자들은 유효하지 않은 데이터를 처리하기 위해 입력을 데이터 텐서 $(x)$ 와 이진 유효성 마스크 $(m)$ 의 튜플로 정의하고, 네트워크 전체에서 다음 원칙을 따르는 마스킹 인식 (Mask-aware) 처리를 제안합니다.

• 동적 마스크 업데이트: 데이터가 네트워크를 통과함에 따라 마스크도 동적으로 업데이트됩니다.
• 연산 규칙:
  • 요소별 연산 (Element-wise): 두 텐서의 연산 시 결과의 유효성은 입력 마스크의 논리곱 (AND) 으로 결정됩니다.
  • 특징 연결 (Concatenation): 채널 축 연결 시에도 입력 마스크의 논리곱을 따릅니다.
  • 수용 영역 연산 (Receptive Field):
    • 기존: 수용 영역 내 하나라도 유효하지 않은 데이터가 있으면 결과도 유효하지 않음.
    • 제안 (Partial): 수용 영역 내 최소 하나의 유효한 데이터가 있으면 결과를 유효한 것으로 간주 (부분 합성곱, 부분 풀링 등 적용).
  • 시퀀스 모델링 (SSM): 시퀀스 역사 (history) 가 수용 영역 역할을 하므로, 하나라도 유효한 토큰이 있어야 유효한 출력을 생성해야 합니다.

2.2. Partial Vision Mamba (PVM)

기존 Vision Mamba (VM) 아키텍처를 마스킹 인식 방식으로 변형한 새로운 구성 요소입니다.

• 부분 패치 임베딩 (Partial Patch Embedding):
  • 패치 (Patch) 내부에 유효하지 않은 픽셀이 섞여 있을 때 (Inner-patch invalidity), 이를 해결하기 위해 **부분 선형 레이어 (Partial Linear layer)**를 사용합니다. 이는 유효하지 않은 위치에 평균 패딩 (mean padding) 을 적용한 후 선형 변환을 수행하여, 결손된 패치도 유효한 토큰으로 변환합니다.
• 마스킹 토큰 (Masked Tokens):
  • 패치 간 유효하지 않은 데이터가 SSM 처리 과정에서 전파되는 것을 방지하기 위해 (Inter-patch invalidity), 유효하지 않은 토큰을 **학습 가능한 마스킹 토큰 (learned masked token, BERT 의 [MASK] 와 유사)**으로 대체합니다. 이를 통해 SSM 이 결손된 데이터를 식별하고 무시하도록 학습시킵니다.
• 구조적 보장: PVM 은 유효하지 않은 플레이스홀더 값에 구조적으로 무감각 (insensitive) 하도록 설계되어, 유효한 데이터 흐름에만 영향을 받습니다.

3. 주요 기여 (Key Contributions)

1. PVM 개발: Vision Mamba 아키텍처 내에서 유효하지 않은 데이터를 처리할 수 있는 최초의 마스킹 인식 구성 요소인 **Partial Vision Mamba (PVM)**를 설계했습니다.
2. 프레임워크 정립: PVM 을 사용하여 임의의 모양을 가진 결손 입력 데이터를 처리하는 아키텍처를 설계하기 위한 **마스킹 인식 프레임워크 (유효성 속성 및 규칙)**를 공식화했습니다.
3. 범용성 입증: 깊이 완성 (Depth Completion), 이미지 인페인팅 (Image Inpainting), 유효하지 않은 데이터가 포함된 이미지 분류 (Image Classification) 등 세 가지 서로 다른 작업에서 PVM 기반 모델이 기존 모델보다 우월함을 실험적으로 증명했습니다.

4. 실험 결과 (Results)

• 깊이 완성 (Depth Completion, KITTI-3D):
  • RGB 가이드가 없는 (unguided) 환경에서 Sparse LiDAR 데이터를 처리했습니다.
  • PVM-DC는 기존 VM 기반 모델 (VM-DC) 대비 RMSE 기준 23% 이상 개선 (1.80m $\to$ 1.38m) 을 보였습니다.
• 이미지 인페인팅 (Image Inpainting, FFHQ):
  • PConvs 기반 모델과 기존 VM-UNet 을 비교했습니다.
  • Mamba 기반 모델이 PConvs 보다 전역 문맥 모델링 능력으로 인해 더 좋은 성능을 보였으며, PVM-UNet-N(완전한 마스킹 인식 설계) 이 가장 우수한 FID 및 LPIPS 점수를 기록했습니다.
  • PConvs 는 아티팩트가 발생하고, VM-UNet 은 흐릿한 반면, PVM 은 더 사실적인 디테일을 복원했습니다.
• 이미지 분류 (Image Classification, ImageNet-1k):
  • 입력 이미지에 임의의 마스크를 적용한 상태에서 분류 정확도를 측정했습니다.
  • PVM-Cls는 기존 PlainMamba 대비 Top-5 정확도가 36% 상대적 향상 (25.60% $\to$ 34.93%) 을 보이며 결손 데이터에 대한 강건성을 입증했습니다.
• Ablation Study:
  • 결손 토큰을 채우는 전략 중 **학습 가능한 토큰 (Learned token)**이 0 패딩이나 평균 패딩보다 가장 우수한 성능을 보였습니다.
  • 다양한 마스크 정책 (Easy, Hard, Extreme) 에서도 PVM 이 일관되게 우월한 성능을 유지하여 일반화 능력을 입증했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 기술적 의의: 이 연구는 State Space Models (Mamba) 가 가진 효율성과 선형 복잡도의 이점을 유지하면서, 실제 세계의 불완전한 데이터 (결손, 가려짐 등) 를 처리할 수 있는 첫 번째 마스킹 인식 아키텍처를 제시했습니다.
• 실용적 가치: 기존 CNN 기반의 부분 합성곱 (PConvs) 원리를 최신 SSM 아키텍처로 확장하여, 고해상도 이미지 처리 및 효율성이 중요한 응용 분야에서 결손 데이터 처리의 한계를 극복했습니다.
• 미래 전망: PVM 의 설계 원칙은 깊이 향상 (Depth Enhancement) 등 다른 결손 데이터 관련 작업으로 확장 가능하며, 마스킹 인식 스캐닝 전략 및 위치 인코딩 개선 등을 통해 지속적으로 발전할 수 있는 기반을 마련했습니다.

요약하자면, 이 논문은 Mamba 기반 모델이 결손된 입력 데이터를 처리할 때 발생하는 성능 저하 문제를 해결하기 위해, PConvs 의 원리를 차용한 'Partial Vision Mamba (PVM)'를 제안하고, 이를 통해 다양한 비전 작업에서 획기적인 성능 향상을 이끌어냈음을 증명합니다.