💻 computer science

RD-ViT: Recurrent-Depth Vision Transformer for Semantic Segmentation with Reduced Data Dependence Extending the Recurrent-Depth Transformer Architecture to Dense Prediction

본 논문은 적응형 계산 시간과 전문가 혼합과 같은 메커니즘을 통해 심장 MRI 벤치마크에서 최첨단 의미 분할 성능을 달성하면서도 단일 공유 블록으로 표준 심층 스택을 대체하여 매개변수 수와 데이터 의존성을 크게 줄이는 순환-깊이 비전 트랜스포머인 RD-ViT를 소개합니다.

원저자: Renjie He

게시일 2026-05-06

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Renjie He

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

의료 영상 (심장 MRI 나 치과 X 선과 같은) 을 보고 특정 부위 (심장 방이나 개별 치아 등) 주위에 완벽한 윤곽선을 그리도록 로봇을 가르친다고 상상해 보세요. 이 작업은 '분할 (segmentation)'이라고 합니다.

오랫동안 이 작업을 수행하는 데 가장 뛰어난 로봇은 **비전 트랜스포머 (Vision Transformers, ViT)**였습니다. 표준 ViT 를 생각하면, 10 명의 서로 다른 전문가가 줄지어 일하는 팀과 같습니다. 첫 번째 전문가가 이미지를 보고, 두 번째 전문가가 첫 번째의 결과를 보고, 세 번째 전문가가 두 번째의 결과를 보고, 이런 식으로 이어집니다. 각 전문가마다 고유한 뇌 (매개변수) 를 가지고 있으며 처음부터 훈련해야 합니다. 이는 놀라울 정도로 잘 작동하지만, 모든 전문가가 각자의 고유한 업무를 학습해야 하므로 방대한 양의 훈련 데이터 (수천 개의 예시) 가 필요합니다.

의료 분야에서는 수천 개의 레이블이 지정된 예시를 얻는 것이 어렵습니다. 의사는 바쁘고 이미지 레이블링에는 시간이 많이 걸립니다. 따라서 이 논문의 저자들은 다음과 같이 질문했습니다: 적은 데이터로도 똑같이 잘 학습하는 로봇을 만들 수 있을까요?

해결책: "반복 전문가" (RD-ViT)

저자들은 새로운 로봇인 RD-ViT를 만들었습니다. 10 명의 서로 다른 전문가를 고용하는 대신, 한 명의 뛰어난 전문가를 고용하여 이미지를 보고, 생각한 뒤, 다시 보고, 다시 생각하며 이 과정을 여러 번 반복하도록 요청했습니다.

이 "반복 전문가"가 이렇게 잘 작동하도록 만든 방법은 다음과 같습니다:

1. "루프" (순환 깊이)
서로 다른 사람들이 긴 줄을 이루는 대신, 한 사람이 이미지를 보고 자신의 메모를 본 뒤 다시 보는 방식입니다.

문제: 같은 사람에게 10 번 보라고만 하면, 루프에 갇히거나 혼란스러워질 수 있습니다.
해결책 (LTI-Stable Injection): 저자들은 이 전문가에게 특별한 "안정성 규칙"을 부여했습니다. 이는 자동차의 서스펜션과 같습니다. 전문가가 이미지를 몇 번 보더라도 이 규칙은 그들이 미치거나 갇히지 않도록 보장합니다. 그들은 항상 명확한 답변에 도달합니다. 이를 통해 로봇은 깨지지 않고 필요한 만큼 "생각"할 수 있습니다.

2. "스마트 중단" (적응형 계산 시간)
의료 이미지의 모든 부분이 동일한 난이도를 갖는 것은 아닙니다. 심장 주변의 빈 공간은 무시하기 쉽지만, 심장 근육의 얇고 구불구불한 가장자리는 추적하기 매우 어렵습니다.

비유: 시험을 보는 학생을 상상해 보세요. 쉬운 질문 (예: "하늘의 색은 무엇인가?") 에 대해서는 즉시 답하지만, 어려운 질문 (예: "이 복잡한 적분을 계산하라") 에 대해서는 추가 시간을 들여 생각합니다.
작동 방식: RD-ViT 는 이미지의 작은 조각마다 "중단 버튼"을 가지고 있습니다. 조각이 쉽다면 (예: 배경), 로봇은 1~2 번의 루프 후 생각을 멈춥니다. 조각이 어렵다면 (예: 심장 경계), 로봇은 올바르게 만들기 위해 3, 4, 또는 5 번 더 루프를 돌립니다. 이는 에너지를 절약하고 중요한 부분에서 로봇을 더 똑똑하게 만듭니다.

3. "전문가 팀" (전문가 혼합)
로봇이 같은 "루프"를 반복해서 사용하지만, 저자들은 **전문가 혼합 (Mixture of Experts, MoE)**이라는 반전을 추가했습니다.

비유: 한 명의 전문가가 머릿속에 8 명의 보이지 않는 조력자 팀을 가지고 있다고 상상해 보세요. 로봇이 "우심실 (심장의 일부)" 패치를 보면 자동으로 "조력자 #1"을 호출합니다. "심근 (심장 근육)"을 보면 "조력자 #2"를 호출합니다.
놀라운 점: 저자들은 로봇에게 어떤 부위에 어떤 조력자를 사용할지 알려주지 않았습니다. 로봇이 스스로 알아낸 것입니다! "우심실" 조력자는 그 모양에 대한 전문가가 되었고, "좌심실" 조력자는 자신의 모양에 대한 전문가가 되었습니다. 이는 추가 지시 없이 로봇이 일을 잘하려고 노력함으로써 발생합니다.

그들은 무엇을 발견했나요?

저자들은 이 새로운 로봇을 두 가지 매우 다른 의료 작업에서 테스트했습니다:

1. 심장 (ACDC 데이터셋)

도전 과제: 훈련 이미지가 매우 적었습니다 (환자 100 명만).
결과:
- 2D 슬라이스(한 번에 하나의 평면 이미지 보기) 에서 RD-ViT 는 데이터가 매우 적음에도 불구하고 표준 10 명 전문가 팀보다 실제로 더 뛰어났습니다. 더 빠르게 학습하고 실수가 적었습니다.
- 3D 볼륨(심장 전체 블록 보기) 에서 표준 팀이 약간 더 좋았지만, "전문가 팀 (MoE)"을 추가하지 않는 한 그랬습니다. MoE 를 사용하면 RD-ViT 는 표준 팀의 정확도 **99.4%**를 달성하면서도 메모리 (매개변수) 를 절반 미만으로 사용했습니다.
- 보너스: 로봇을 8 번 루프하도록 훈련시켰다면, 실제 테스트 중에는 16 번 루프하도록 요청할 수 있으며, 재훈련 없이도 성능이 더 좋아지거나 (또는 동일하게 유지) 된다는 것을 발견했습니다. 이는 교과서를 바꾸지 않고도 더 오래 공부하여 더 좋은 성적을 얻을 수 있는 학생과 같습니다.

2. 치아 (ToothFairy2 데이터셋)

도전 과제: 3D 치과 스캔에 같은 로봇을 적용하여 32 개의 서로 다른 치아를 식별하고 올바르게 번호를 매기는 작업 (예: "치아 15" 대 "치아 25") 을 시도했습니다. 이는 어느 치아가 어느 것인지 알기 위해 전체 턱을 봐야 하므로 어렵습니다.
결과: 로봇은 치아를 성공적으로 식별하고 89.1% 의 정확도로 올바른 번호를 할당했습니다.
작동 이유: 로봇의 "전역 주의 (global attention)"는 한 번에 전체 턱을 볼 수 있게 했습니다. 오른쪽 상단의 치아와 왼쪽 상단의 치아가 다르다는 것을 구별할 수 있었는데, 이는 작은 국소 조각만 보는 이전의 단순한 로봇들이 struggle 하던 부분입니다.

결론

이 논문은 훌륭한 의료 분할을 수행하기 위해 거대한 다양한 AI 모델 군대가 필요하지 않음을 보여줍니다. 대신, 루프를 돌고, 적응적으로 생각하며, 내부 전문가를 갖춘 하나의 스마트한 모델을 사용할 수 있습니다.

데이터가 부족할 때 (특히 2D 에서) 더 빠르게 학습합니다.
다른 레이어 간에 "뇌"를 공유함으로써 공간을 절약합니다.
지시 없이 자신의 전략을 스스로 파악합니다 (예: 어떤 전문가가 어떤 신체 부위를 처리할지).
이미지의 어려운 부분에서는 더 깊이 생각하고 쉬운 부분에서는 덜 생각할 수 있습니다.

저자들은 모든 코드와 결과를 공개하여 이 "순환 (recurrent)" 접근 방식이 AI 에게 인체를 보게 하는 강력하고 효율적인 방법임을 입증했습니다.

기술 요약: RD-ViT: 데이터 의존성이 감소된 의미론적 분할을 위한 재귀-깊이 비전 트랜스포머

문제 제기

비전 트랜스포머 (ViT) 는 의미론적 분할과 같은 밀집 예측 작업에서 최첨단 성능을 달성했습니다. 그러나 표준 ViT 는 아키텍처 설계로 인해 높은 데이터 의존성을 겪습니다. 깊은 스택의 각 레이어는 독립적으로 학습되어야 하는 고유한 매개변수를 가지기 때문입니다. 이로 인해 매개변수 수가 깊이 ( $O(L \cdot d^2)$ ) 에 따라 선형적으로 증가하여, 훈련 데이터가 부족할 때 ViT 가 과적합되기 쉽습니다. 이 한계는 라벨링된 데이터셋이 종종 작고 (예: ACDC 심장 MRI 벤치마크는 100 개의 환자 볼륨만 포함), 라벨링 비용이 많이 들고 시간이 오래 걸리는 의료 영상 분야에서 특히 중요합니다. 본 논문은 계산 효율성을 유지하면서 제한된 훈련 데이터로 높은 분할 정확도를 달성할 수 있는 아키텍처의 필요성에 대응합니다.

방법론

저자들은 OpenMythos 프로젝트의 재귀-깊이 트랜스포머 (RDT) 개념을 밀집 예측 작업에 적용한 RD-ViT(Recurrent-Depth Vision Transformer) 를 제안합니다. 고유한 트랜스포머 블록의 깊은 스택 대신, RD-ViT 는 단일 공유 트랜스포머 블록을 $T$ 번 반복 적용합니다. 아키텍처는 네 가지 주요 구성 요소로 이루어집니다:

프렐류드 & 코다: 각각 초기 특징 추출과 최종 정제를 위한 표준 ViT 블록입니다.
재귀 루프: 단일 ViT 블록을 반복적으로 적용하는 핵심 혁신입니다. 이 루프에는 몇 가지 특수 메커니즘이 통합되어 있습니다:
- LTI-안정 주입: 긴 시간 범위의 재귀 중 수렴을 보장하기 위해 상태 전이 행렬 $A$ 를 $A = \exp(-\exp(\log d_t + \log A))$ 로 매개변수화합니다. 이 구성은 모든 고유값이 엄격히 $(0, 1)$ 내에 있도록 보장하여 재귀가 축소 사상 (contraction mapping) 이 되도록 합니다. 이를 통해 훈련 시보다 더 많은 루프로 추론하는 안정적인 깊이 외삽이 가능해집니다.
- 적응형 계산 시간 (ACT): 학습된 정지 메커니즘이 각 반복 단계에서 각 패치 토큰에 정지 확률을 할당합니다. "쉬운" 영역 (예: 균질한 배경) 의 토큰은 일찍 정지하는 반면, "어려운" 영역 (예: 경계) 의 토큰은 더 많은 반복을 거칩니다. 이는 공간 계산 맵을 생성하여 자원을 동적으로 할당합니다.
- 깊이별 LoRA: 전체 랭크 가중치를 복제하지 않고 공유 블록이 다양한 깊이에서 다르게 동작하도록 허용하기 위해, 루프 인덱싱 스케일링이 적용된 저랭크 적응 (LoRA) 어댑터를 적용합니다.
- 혼합 전문가 (MoE): 선택적 희소 MoE 레이어가 표준 피드포워드 네트워크 (FFN) 를 대체합니다. 라우터는 토큰을 특정 전문가 서브네트워크로 유도하여 공유 블록 내에서 조건부 계산과 클래스별 전문화를 가능하게 합니다.
분할 헤드: 전치 합성곱을 사용하여 패치 수준의 특징을 픽셀/볼륨 수준 예측으로 업샘플링합니다.

이 아키텍처는 차원에 구애받지 않으며, 2D(Conv2d 패치 임베딩을 통해) 와 3D(Conv3d 를 통해) 입력 모두를 지원합니다.

주요 기여

아키텍처 적응: LTI-안정 주입, ACT, 깊이별 LoRA, MoE 를 통합하여 밀집 예측을 위한 완전한 2D/3D 분할 아키텍처를 구현했습니다.
데이터 효율성 검증: ACDC 심장 MRI 데이터셋에 대한 실증적 검증을 통해 RD-ViT 가 표준 ViT 보다 훨씬 적은 매개변수와 데이터로, 특히 2D 설정에서 더 나은 성능을 발휘하거나 동등한 성능을 달성할 수 있음을 보여주었습니다.
발현되는 전문화: 명시적인 라우팅 감독 없이 MoE 전문가들이 서로 다른 해부학적 구조 (예: 우심실, 심근, 좌심실) 를 위해 자발적으로 전문화되는 증거를 제시하여, 다중 클래스 분할을 위한 구조적 귀납적 편향을 제공했습니다.
깊이 외삽: LTI-안정성 보장 덕분에 훈련 시 사용된 것보다 더 많은 루프 반복 ( $T$ ) 으로 안전하게 추론할 수 있음을 입증하여 성능 저하 없이 수행되었습니다.
일반화: 동일한 아키텍처를 근본적으로 다른 작업인 CBCT 볼륨에서의 3D 치아 인스턴스 분할 및 번호 매기기에 성공적으로 적용하여, 장거리 공간 추론을 위한 글로벌 어텐션을 활용했습니다.

실험 결과

심장 분할 (ACDC 데이터셋)

3D 볼륨:
- RD-ViT Tiny + MoE(300 만 개 매개변수) 는 0.812 의 Dice 점수를 달성하여 표준 ViT 기준 (0.817, 570 만 개 매개변수) 의 **99.4%**에 도달했으며, 매개변수는 **53%**만 사용했습니다.
- Ablation: MoE 가 가장 큰 향상 (+0.8 백분율 포인트, RD-ViT Tiny 대비) 을 제공했습니다. ACT 와 LoRA 는 작지만 일관된 개선 (+0.3 및 +0.2 백분율 포인트 각각) 을 제공했습니다.
- 깊이 외삽: 추론 루프를 $T=8$ (훈련) 에서 $T=16$ 으로 증가시켰을 때 성능이 안정적으로 유지되어 LTI 주입의 안정성을 확인했습니다.
2D 슬라이스 수준:
- RD-ViT 는 모든 데이터 비율에서 표준 ViT 를 능가했습니다. 훈련 데이터 10% 에서 RD-ViT 는 0.774의 Dice 를 달성한 반면 표준 ViT 는 0.762였습니다. 100% 데이터에서는 0.882대 0.872였습니다.
- 저자들은 2D 슬라이스에서 더 풍부한 공간적 맥락이 반복적 정제를 3D 저해상도 설정보다 더 효과적으로 만들었다고 설명합니다.

치아 분할 (ToothFairy2 데이터셋)

3D 치아 인스턴스 분할 및 FDI 번호 매기기에 적용 (33 개 클래스).
단일 단계 RD-ViT Tiny + MoE 모델 (540 만 개 매개변수) 은 다음을 달성했습니다:
- 0.797 Seed Dice(치아 대 배경).
- 89.1% FDI 분류 정확도.
결과는 트랜스포머의 글로벌 자기 어텐치가 치아 번호 매기기에 필요한 장거리 턱-사분면 맥락을 효과적으로 포착하며, 이는 지역 CNN 수용 필드가 종종 어려움을 겪는 작업임을 시사합니다.

계산 효율성

ACT 동역학: 평균 ponder 시간 (유효 반복 횟수) 이 훈련 중 약 2.6 에서 약 1.4 로 감소하여 모델이 쉬운 영역을 더 빠르게 처리하도록 학습했음을 나타냅니다.
계산 할당: ACT 정지 맵은 심장 경계에서 더 높은 계산 할당을, 균질한 영역에서는 더 낮은 할당을 보여주었습니다.

중요성 및 주장

본 논문은 RD-ViT 가 데이터가 제한된 의료 분야에서 매개변수 효율적인 분할을 위한 실행 가능한 아키텍처 전략을 제공한다고 주장합니다. 가중치 공유 (재귀) 와 조건부 계산 (MoE) 을 결합함으로써, 모델은 표준 ViT 의 약 절반 정도의 매개변수로 거의 최첨단 성능을 달성합니다.

저자들은 LTI-안정 주입이 깊이 외삽을 가능하게 하는 데 결정적이며, 재훈련 없이 속도와 정확도 간의 트레이드오프를 조정할 수 있는 "무료 추론 시간 노브"를 제공한다고 강조합니다. 또한, MoE 전문가들의 자발적 전문화는 재귀-깊이 아키텍처가 복잡한 분할 문제를 자연스럽게 클래스별 하위 작업으로 분해할 수 있음을 시사합니다.

이 연구는 2D 와 3D 성능 간의 이분법을 강조합니다. RD-ViT 는 풍부한 공간적 맥락으로 인해 2D 에서 뛰어나지만, 3D 성능은 현재 재귀 아키텍처 자체보다는 모델 용량에 의해 더 제한받고 있습니다. 그러나 치과 영상으로의 성공적인 이전은 아키텍처의 도메인 중립적 잠재력을 보여줍니다.

본 논문은 RD-ViT 가 모든 3D 시나리오에서 고도로 최적화된 데이터 집약적 파이프라인을 아직 능가하지는 못하지만, 데이터 의존성을 크게 줄이고 효율적이고 해석 가능하며 적응 가능한 의료 영상 분할을 위한 새로운 패러다임을 제공한다고 결론지었습니다. 모든 코드, 노트북 및 결과는 공개되었습니다.