🧬 biology

KHiM-Mamba: Injecting Pathology Knowledge into Mamba via Hidden-State Modulation for Whole Slide Image Analysis

KhiM-Mamba는 병리학적 지식과 LLM 유래 시맨틱 기술(semantic descriptions)을 통합하여 Mamba의 은닉 상태(hidden states)를 조절함으로써, 무관한 정보의 축적을 방지하고 다양한 진단 작업에서 최첨단 성능을 달anim하는 전체 슬라이드 이미지(Whole Slide Image) 분석용 신규 다중 인스턴스 학습 아키텍처입니다.

원저자: Qixiang Zhang, Yi Li, Tianqi Xiang, Haonan Wang, Mengjiao Wei, Bo Xu, Xiaomeng Li

게시일 2026-08-18

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Qixiang Zhang, Yi Li, Tianqi Xiang, Haonan Wang, Mengjiao Wei, Bo Xu, Xiaomeng Li

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ⚕️ 이것은 동료 심사를 거치지 않은 프리프린트의 AI 생성 설명입니다. 의학적 조언이 아닙니다. 이 내용을 바탕으로 건강 관련 결정을 내리지 마세요. 전체 면책 조항 읽기

현대 의학의 세계에서 병리학자의 가장 중요한 도구는 인간 조직의 아주 작은 샘ակ이 담긴 유리 슬라이드이며, 이는 생명과 질병의 세포 구조를 드러낼 때까지 확대됩니다. 오늘날 이 슬라이드들은 수십억 개의 픽셀을 포함할 정도로 거대한 디지털 이미지로 스캔되는데, 그 규모가 너무 방대하여 인간의 눈으로는 한 번에 살펴볼 수 없습니다. 의사들이 이 디지털 풍경을 탐색하는 것을 돕기 위해, 과학자들은 지칠 줄 모르는 조수 역할을 하는 인공지능 시스템을 개발했습니다. 이 시스템은 거대한 이미지를 수천 개의 작고 관리 가능한 조각으로 나누고, 각 조각의 세포 질감과 형태를 분석한 다음, 이를 종합하여 전역적인 진단을 내리려고 시도합니다. 과제는 엄청난 규모와 노이즈에 있습니다. 슬라이드 하나에는 수백만 개의 세포가 포함될 수 있지만, 진단의 핵심을 쥐고 있는 것은 단 몇 개의 흩어진 클러스터뿐이며, 나머지는 건강한 조직이나 무관한 배경입니다. 만약 컴퓨터가 방대한 양의 불필요한 정보에 정신이 팔린다면, 구석에 숨겨진 결정적인 단서를 놓칠 수 있습니다.

연구자들은 최근 이러한 거대한 디지털 슬라이드를 처리하기 위해, 인간의 뇌가 시간에 따른 정보의 시퀀스를 처리하는 방식에서 영감을 얻은 새로운 유형의 인공지능 아키텍처로 눈을 돌렸습니다. 이 기술은 긴 시퀀스를 따라 이동하며 이전에 보았던 것을 기억하고, 앞으로 나아가면서 무엇을 메모리에 유지할지 결정하는 데 탁월합니다. 그러나 홍콩과 충칭의 과학자 팀은 이러한 시스템이 사용되는 방식에서 결함을 발견했습니다. 그들은 이 AI 모델들이 슬로를 스캔할 때 전적으로 픽셀이 어떻게 보이는지에만 의존한다는 것을 발견했습니다. 만약 슬라이드가 실험실마다 흔히 발생하는 변동인 약간 다른 색상의 염료로 염색되었다면, 모델은 혼란을 느껴 무해한 색상 변화를 중요한 의학적 단서로 취급할 수 있습니다. 이는 시스템이 불필요한 세부 사항을 메모리에 축적하게 만들어, 진단에 필요한 결정적인 증거를 희석시키는 결과를 초래했습니다. 이를 해결하기 위해 연구진은 실제 의학적 지식을 AI의 메모리 프로세스에 직접 주입하는 새로운 방법을 도입하여, 시스템이 노이즈를 무시하고 진정으로 중요한 것에 집중하도록 가르쳤습니다.

홍콩 과학기술대학교와 충칭대학교 연구진이 이끄는 팀은 KHiM-Mamba라고 부르는 시스템을 개발했습니다. AI가 슬라이드를 스캔할 때 스스로 무엇을 기억할지 결정하게 두는 대신, 그들에게 병리학의 언어로 쓰인 가이드북을 주었습니다. 이 가이드북에는 건강하거나 질병이 있는 상태에서 서로 다른 유형의 조직이 어떻게 보이는지에 대한 구체적인 설명이 담겨 있습니다. AI가 슬라이드의 한 작은 조각에서 다음 조각으로 이동할 때, 이 가이드북을 참조하여 어떤 정보를 메모리에 저장할 가치가 있고 어떤 것을 버려야 할지를 결정합니다. 마치 엄청난 양의 책 더미를 분류하는 사서에게 보관할 제목 목록과 무시할 제목 목록이 명시된 리스트가 전달된 것과 같습니다. 사서는 더 이상 어떤 책이 중요한지 추측할 필요가 없습니다. 이 경우, '사서'는 AI이고, '책'은 작은 이미지 패치이며, '리스트'는 대규모 언어 모델에 의해 제공된 의학적 지식입니다.

연구진은 이전에 언어를 사용하여 AI의 슬라이 진단을 도우려 했던 시도들이 종종 너무 간접적이었다는 점을 발견했습니다. 그러한 방법들은 분석이 시작되기 전에 어떤 슬라이드를 볼지 선택하거나, 분석이 끝난 후 최종 답안을 확인하는 데 언어를 사용했습니다. 새로운 접근 방식은 다릅니다. 왜냐하면 이 방식은 AI가 생각하는 핵심 메커니즘 자체를 바꾸기 때문입니다. 이는 메모리 시스템의 내부적인 '쓰기(write)' 및 '읽기(read)' 프로세스를 수정합니다. AI가 새로운 메모리를 쓸 때, 의학적 지식은 시스템이 단순히 오해를 불러일으킬 수 있는 색상이나 질감이 아니라, 질병과 관련된 특정 형태와 패턴에 집중하도록 지시합니다. AI가 결정을 내리기 위해 메모리에서 정보를 읽을 때, 지식은 지금까지 축적된 가장 관련성 높은 증거를 검색하도록 안내합니다. 이는 시스템이 단순히 시각적 변이에 휘둘리는 것이 아니라, 의학적으로 의미 있는 연결 관계를 바탕으로 슬라이드의 표현형을 구축하도록 보장합니다합니다.

이를 구현하기 위해 팀은 각 특정 부분에 적합한 의학적 설명을 생성하는 스마트한 시스템을 만들었습니다. 단일 슬라이드에는 많은 다양한 유형의 조직이 포함되어 있으므로, 일반적인 설명만으로는 충분하지 않습니다. 이 시스템은 강력한 언어 모델을 사용하여 작은 조직 패치를 살펴보고, 그것이 어떤 종류의 조직인지 식별한 다음, 해당 조직이 건강하거나 질병이 있을 때 어떻게 보여야 하는지에 대한 정밀하고 상세한 설명을 생성합니다. 만약 시스템이 조직 유형에 대해 확신이 없다면, 실수를 피하기 위해 더 넓고 안전한 설명을 사용합니다. 이 역동적인 과정은 AI가 해결하고 있는 퍼즐의 매 조각에 대해 가능한 한 가장 정확하고 구체적인 지침을 받을 수 있도록 보장합니다.

이 접근 방식의 결과는 다양한 암 유형과 환자 예후를 포함하는 11개의 공개 데이터셋을 통해 테스트되었습니다. 새로운 시스템은 기존의 최선책들을 일관되게 능가했습니다. AI가 유방암의 서로 다른 하위 유형을 구별해야 하는 작업에서, 이 시스템은 이전의 어떤 모델보다 높은 정확도를 달al성했습니다. 조직 샘플을 기반으로 환자의 생존 시간을 예측할 때, 이 시스템은 더 신뢰할 수 있는 위험 평가를 제공했으며, 환자들을 고위험군과 저위험군으로 더 잘 분류했습니다. 아마도 가장 중요한 점은, 이 시스템이 서로 다른 염색 기법을 사용하는 서로 다른 병원의 슬라이드로 테스트되었을 때도 견고함을 입증했다는 것입니다. 다른 모델들이 이러한 변동성에 어려움을 겪으며 색상이 약간만 변해도 자주 실패했던 반면, 새로운 시스템은 높은 성능을 유지했습니다. 이는 AI를 단순히 시각적 패턴이 아닌 의학적 지식에 근거하게 함으로써, 시스템이 슬라이드의 외관이 아닌 질병 자체를 보도록 학습했음을 시사합니다.

또한 연구는 레이블이 지정된 사례가 부족한 의료 분야의 흔한 문제인, 학습할 데이터가 매우 적은 상황에서 시스템이 얼마나 잘 작동하는지를 조사했습니다. AI가 질병을 배우기 위해 단 8개 또는 16개의 사례만을 보여준 테스트에서도, 이 새로운 방법은 경쟁 모델들을 크게 앞질렀습니다. 이는 시스템에 주입된 의학적 지식이 강력한 토대로 작동하여, 최소한의 훈련 데이터만으로도 효과적으로 학습할 수 있게 함을 나타냅니다. 연구진은 또한 이러한 지식 계층을 추가하는 것이 시스템의 속도를 늦추거나 표준 컴퓨터 하드웨어에서 실행하기에 너무 크게 만들지 않는다는 것을 확인했습니다. 시스템은 여전히 빠르고 효율적이었으며, 수천 개의 이미지 패치를 몇 초 만에 처리할 수 있었습니다.

인공지능이 의료 이미지를 처리하는 방식을 재고함으로써, 이 연구는 시각적 데이터와 명시적인 의학적 지식을 결в합하는 것이 더 신뢰할 수 있고 정확한 진단 도구를 만든다는 것을 보여줍니다. 연구진은 방대한 양의 노이즈 섞인 데이터를 해결하는 열쇠가 단순히 더 큰 컴퓨터를 만드는 것이 아니라, 컴퓨터에게 무엇을 찾아야 하는지 가르치는 것임을 입증했습니다. 이 접근 방식은 계산 병리학의 새로운 길을 제시하며, 여기서 목표는 불필요한 것을 걸러내고 결정적인 것을 강조함으로써 의사를 보조하여, 최종 진단이 가장 강력한 증거를 바탕으로 이루어지도록 하는 것입니다.

기술 요약: KHiM-Mamba

문제 정의

계산 병리학에서의 전체 슬라이드 이미지(WSI) 분석은 일반적으로 다중 인스턴스 학습(Multiple Instance Learning, MIL)으로 정식화됩니다. 여기서 슬라이드는 패치들로 분해되며, "슬라이드 인코더"는 패치 특징들을 슬라이드 수준의 표현으로 집합시킵니다. 이 과정, 즉 **슬라이드 인코딩 역학(slide encoding dynamics)**은 매우 중요한데, 진단적 근거가 멀리 떨어진 조직 영역들에 걸쳐 분포되어 있기 때문입니다.

최근에는 선택적 상태 공간 모델(Selective State Space Models, SSM), 특히 Mamba가 긴 시퀀스를 선형 복잡도로 모델링할 수 있는 능력 덕분에 이 작업에 효율적인 아키텍처로 부상했습니다. 그러나 저자들은 기존 Mamba 기반 MIL 방식의 근본적인 한계를 지적합니다. 즉, 이들은 선택적 상태 전이 과정에서 **오로지 시각적 특징(visual features)**에만 의존한다는 점입니다. 대규모 WSI에서는 진단적으로 결정적인 영역이 풍부한 무관한 정보(예: 염색 변이, 배경 조직)에 둘러싸여 있습니다. 순수하게 시각 중심적인 역학은 상태 업데이트를 잘못 할당할 수 있습니다. 이로 인해 진화하는 은닉 상태(hidden state)가 작업과 무관한 증거를 축적하게 되어, 긴 스캔 궤적 동안 중요한 진단적 단서가 희석되는 문제가 발생합니다.

시각-언어 모델(Vision-Language Model, VLM) 기반의 MIL 방식들이 존재하지만, 이들은 대개 언어를 핵심 인코딩 역학에 통합하기보다는 외부적인 부가 요소(예: 패치 선택 또는 보조 감독용)로 취급합니다. 결과적으로, 순수 시각적 상태 전이의 한계는 해결되지 않은 채 남아 있습니다.

방법론: KHiM-Mamba

저자들은 Mamba의 핵심적인 선택적 상태 공간 메커니즘에 병리학적 지식 사전(knowledge priors)을 직접 주입하는 프레임워크인 KHiM-Mamba(Knowledge-aware Hidden-state Modulation)를 제안합니다. 이 아키텍처는 두 가지 주요 구성 요소로 이루어져 있습니다.

1. 지식 인식 은닉 상태 변조 (KHiM)

KHiM-Mamba는 언어를 주변 신호로 사용하는 대신, 언어 유래 지식을 사용하여 은닉 상태의 진화를 변조하도록 SSM 레이어를 재설계합니다.

메커니즘: 표준 Mamba 재귀는 정보를 은닉 상태( $h_t$ )에 쓰고 읽는 방식을 제어하는 입력 투영( $B_t$ ) 및 출력 투영( $C_t$ ) 행렬을 포함합니다.
변조: 저자들은 저차원 변조(low-rank modulation) 메커니즘을 도입하여, 언어 임베딩( $e_t$ $e_{t}$ )이 MLP를 통해 계수 벡터( $\alpha_B, \alpha_C$ $α_{B}, α_{C}$ )를 생성하도록 합니다. 이 계수들은 원래의 투영 행렬에 더해질 저차원 증분( $\Delta B_t, \Delta C_t$ $Δ B_{t}, Δ C_{t}$ )을 생성하는 데 사용됩니다.
- 쓰기 변조 ( $\tilde{B}_t$ ): 어떤 시각적 증거가 진화하는 은닉 상태에 포함될지를 안내하며, 진단 관련 특징의 축적을 장려하고 무관한 특징을 억제합니다.
- 읽기 변조 ( $\tilde{C}_t$ ): 축적된 문맥적 증거 중 무엇을 추출하여 패치 표현을 풍부하게 할지를 안내하며, 출력이 진단적으로 의미 있는 패턴을 강조하도록 보장합니다.
결과: 이는 SSM의 시간적 감쇠(temporal decay) 동작을 변경하지 않으면서 "쓰기-읽기" 역학을 재구성하여, 선택적 메커니즘을 "진단 인식형(diagnosis-aware)"으로 만듭니다.

2. 국소 적응형 어휘 검색 모듈 (Local-Adaptive Vocabulary Retrieval Module)

변조를 위한 효과적인 언어 사전을 제공하기 위해, 저자들은 고정된 전역 프롬프트를 피하는 모듈을 설계했습니다.

과정:
1. 개념 생성: LLM에 작업 정의를 프롬프트로 제공하여, 작업 관련 조직 개념(예: 특정 종양 유형)과 그에 대한 세밀한 형태학적 설명을 생성합니다.
2. 신뢰도 게이트 기반 라우팅: 각 패치에 대해, 시각-언어 모델(CONCH)이 패치 이미지 임베딩과 개념 텍스트 임베딩 사이의 유사도를 계산합니다.
3. 적응형 할당: 특정 조직 개념에 대한 유사도가 임계값을 초과하고 다른 개념들과의 차이가 충분하다면, 해당 패스에 세밀한 조직 특화 언어를 할당합니다. 그렇지 않으면 노이즈 섞인 사전 정보를 도입하는 것을 방지하기 위해 거친 수준의 병리학적 설명을 사용합니다.
결과: 이를 통해 다양한 조직 구성에 걸쳐 의미론적으로 정밀하고 견고한 변조 신호를 얻을 수 있습니다.

주요 기여

한계 식별: 본 논문은 기존 Mamba 기반 MIL 방식이 순수 시각 중심의 상태 전이로 인해 긴 WSI 시퀀스에서 무관한 증거를 축적한다는 점과, 기존 VLM-MIL 방식이 핵심 인코딩 역학에 개입하지 않기 때문에 이 문제를 해결하지 못함을 강조합니다.
아키텍처 혁신: KHiM-Mamba는 Mamba를 단순히 정적인 백본으로 사용하는 것이 아니라, 지식 사전을 직접 선택적 상태 전이(쓰기 및 읽기 행렬)에 주입함으로써 WSI 분석을 위해 Mamba의 내부 모델 역학을 재구성한 첫 번째 연구입니다.
국소 적응형 어휘 검색: LLM을 통해 패치에 세밀한 조직 특화 언어를 동적으로 할당함으로써, 다양한 다운스트림 작업에 대해 정밀한 변조 신호를 제공하는 모듈을 제안했습니다.
실험적 성능: 광범-범주(cancer subtyping), 생존 예측, 도메인 적응, 퓨샷 학습(few-shot learning)의 4가지 작업에 걸친 11개의 공개 벤치마크에서 뛰어난 성능을 입증했습니다.

실험 결과

저자들은 4가지 작업(암 하위 유형 분류, 생존 예측, 도메인 적응, 퓨샷 학습)에 걸친 11개의 공개 벤치마크에서 KHiM-Mamba를 평가했습니다.

암 하위 유형 분류 (Cancer Subtyping):
- BRACS 데이터셋: 미세 분류 7-클래스 작업에서 SOTA 정확도(0.5874)와 AUC(0.8746)를 달달성했으며, 3-클래스 작업에서도 이전 최고치보다 각각 약 1.77%, 2.41% 높은 성과를 보였습니다.
- TCGA-BRCA: 정확도 0.9330, AUC 0.9559를 달성하여 이전 최고 정확도 대비 1.97%, AUC 대비 1.01% 향상되었습니다.
생존 예측 (Survival Prediction):
- 6개의 TCGA 서브셋(BLCA, BRCA, KIRC, KIRP, LUAD, STAD)에서 평가되었습니다. KHiM-Mamba는 평균 C-Index 0.7394를 기록하여, 이전 최고 모델인 2DMamba를 2.31% 앞질렀습니다.
- Kaplan-Meier 곡선은 고위험군과 저위험군 사이의 훨씬 더 뚜렷한 분리를 보여주었으며, 눈에 띄게 낮은 로그-순위(log-rank) $p$ -값(예: LUAD의 경우 $2.63 \times 10^{-6}$ )을 나타냈습니다.
도메인 일반화 (교차 센터 - Domain Generalization):
- CAMELYON16으로 학습하고 CAMELYON17(5개 독립 센터)로 테스트했습니다. KHiM-Mamba는 평균 AUC 0.9184를 달성하여, 이전 최고 모델인 TransMIL보다 2.27% 높은 성능을 보이며 염색 및 스캐너 변이에 대한 견고함을 입증했습니다.
퓨샷 학습 (Few-Shot Learning):
- TCGA-NSCLC의 8-shot 및 16-shot 설정에서 KHiM-Mamba는 평균 정확도 83.11%, AUC **93.03%**를 달성하여, 다른 MIL 및 VLM-MIL 방식들을 크게 상회했습니다. 이는 언어 변조로부터 오는 강력한 귀납적 편향(inductive bias)을 나타냅니다.
효율성:
- 약 2.2M 개의 파라미터를 가진 선형 복잡도를 유지하며, 긴 시퀀스(5,000개 패치)에서도 다른 Mamba 기반 방식(예: PAM, 2DMamba)과 대등하거나 더 빠른 추론 시간을 보여줍니다.

의의 및 주장

본 논문은 KHiM-Mamba가 순수 시각적 특징 집합을 넘어섬으로써 계산 병리학의 근본적인 간극을 해결한다고 주장합니다. 진단적 지식으로 은닉 상태의 진화를 명시적으로 조절함으로써, 모델은 증거 축적을 병리학적으로 의미 있는 영역으로 유도하고 무관한 배경 노이즈를 효과적으로 필터링합니다.

저자들은 이 접근 방식이 기존의 스캐닝 전략(예: 그룹화, 계층적 스캐닝)과 **직교(orthogonal)**하며, 이들과 결 조합될 수 있음을 강조합니다. 이 연구의 의의는 **언어 조건부 변조(language-conditioned modulation)**가 핵심 상태 공간 메커니즘을 변화시키는 것이 일반화, 도메인 변화에 대한 견고성, 데이터가 부족한 시나리오에서의 성능을 향상시키는 강력한 귀납적 편향임을 입증했다는 데 있습니다. 이는 의료 영상 분야에서 지식 가이드형 상태 공간 모델을 설계하기 위한 새로운 관점을 제시합니다.