💻 computer science

Predicting Functions, Not Features: KANs with Function-Space Joint-Embedding Predictive Learning for Medical Image Segmentation

이 논문은 마스킹된 자기지도 학습 방식으로 개별 에지 함수(edge function)의 구조화된 다중 반경 시그니처를 예측함으로써 콜모고로프-아르노프 네트워크(KAN)를 의료 영상 분할을 위해 강화하는 새로운 프레임워크인 FS-JEPA(Function-Space Joint-Embedding Predictive Learning)를 소개하며, 이를 통해 다섯 가지 벤치마크에서 최첨단 성능을 달성하였다.

원저자: Yungeng Liu, Xuanzi Fang, Yuge Zhang, Shuqi Ren, Haijin Zeng, Yongyong Chen

게시일 2026-08-13

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Yungeng Liu, Xuanzi Fang, Yuge Zhang, Shuqi Ren, Haijin Zeng, Yongyong Chen

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

당신이 로봇에게 X선 영상을 보고 심장 그림을 그리도록 가르치려 한다고 상상해 보세요. 이것은 인공지능의 한 분야인 **의료 영상 분할(medical image segmentation)**의 세계입니다. 컴퓨터가 장기, 종양 또는 조직의 정밀한 윤곽선을 그리는 법을 배우는 과정이죠. 이는 매우 까다로운 작업입니다. 의료 영상은 종종 흐릿하고, 건강한 조직과 병든 조직 사이의 경계가 모호하며, 의사를 도울 수 있도록 로봇이 믿을 수 없을 만큼 정확해야 하기 때문입니다.

이를 위해 현대의 AI는 "신경망(neural networks)"을 사용하는데, 이는 거대한 수학적 연결망과 같습니다. 오랫동안 이러한 네트워크는 고정된 규칙(항상 켜져 있거나 꺼져 있는 전등 스위치 같은 규칙)을 사용하여 정보를 처리했습니다. 하지만 더 새롭고 화려한 유형의 네트워크인 **콜모고로프-아르노프 네트워크(Kolmogorov–Arnold Network, KAN)**가 등장했습니다. KAN은 모든 연결에 "학습 가능한 함수"를 사용합니다. 이를 다음과 같이 생각할 수 있습니다. 일반적인 네트워크에서 모든 전선은 딱딱한 파이프와 같습니다. 하지만 KAN에서 모든 전선은 네트워크가 필요에 따라 늘어나거나, 직선이 되거나, 곡선이 될 수 있도록 조절할 수 있는 신축성 있고 형태가 변하는 고무줄과 같습니다. 여기서 연구자들이 던지는 큰 질문은 이것입니다. "어떻게 하면 이 고무줄들이 완벽한 방식으로 늘어나도록 가르칠 수 있을까?"

여기서 새로운 연구가 등장합니다. 연구자들은 KAN이 강력하긴 하지만, 표준적인 학습 방식이 마치 오케스트라 전체의 최종 소리만으로 연주자를 판단하는 것과 같다는 것을 발견했습니다. 즉, 개별 바이올린 연주자들의 소리를 전혀 듣지 않는 것이죠. 그들은 "고무줄(함수)"이 혼합되기 전, 그들이 어떻게 행동하고 있는지에 대한 구체적인 피드백을 충분히 받지 못하고 있다는 점을 깨달았습니다. 이를 해결하기 위해 그들은 FS-JEPA라는 새로운 학습 기법을 발명했습니다. 단순히 최종 결과만을 보는 대신, 각 고무줄의 행동 "형태"를 미리 예측하게 만든 것입니다. 이 방법을 다섯 가지 서로 다른 의료 영상 데이터셋에 테스트한 결과, 이 방식이 이전의 KAN 기반 방식들보다 훨씬 더 선명하고 정확한 윤곽선을 그려내며, 평균 정확도 점수를 2.25 퍼센트 포인트 향상시킨다는 것을 보여주었습니다.

문제점: "오케스트라" 대 "솔로 연주자"

구체적인 메커니즘을 살펴보겠습니다. 표준적인 KAN에서 네트워크는 층(layer)으로 구성됩니다. 각 층은 입력을 받아 이 특별한 "고무줄" 함수들을 통과시켜 출력을 생성합니다. 문제는 네트워크가 보통 최종 출력, 즉 "오케스트라"의 노래를 보고 학습된다는 점입니다. 노래가 좋으면 네트워크는 금메달을 받습니다. 노래가 나쁘면 빨간 불이 들어옵니다.

하지만 여기에는 함정이 있습니다. 많은 다양한 조합의 고무줄들이 동일한 최종 노래를 만들어낼 수 있습니다. 어떤 고무줄은 너무 과하게 늘어났지만, 다른 고무줄이 딱 알맞게 눌려서 이를 상쇄했을 수도 있습니다. 네트워크는 금메달을 받지만, 개별 고무줄은 왜 자신이 그렇게 늘어났는지 배우지 못합니다. 그들은 명확한 목표 없이 추측만 할 뿐입니다. 이는 지휘자가 바이올리니스트에게 "음악이 좋네요"라고 말하면서도, 정작 그 연주자의 특정 음이 음정이 맞았는지 틀렸는지는 알려주지 않는 것과 같습니다.

해결책: 고무줄의 "형태"를 예측하기

이 논문의 저자들(Yungeng Liu와 Xuanzi Fang 주도)은 게임의 규칙을 바꾸기로 했습니다. 그들은 이렇게 물었습니다. "만약 우리가 각 고무줄이 최종 노래에 섞이기 전에, 그 행동을 예측하도록 AI를 가르칠 수 있다면 어떨까?"

그들은 **함수 공간 결합 임베딩 예측 학습(Function-Space Joint-Embedding Predictive Learning, FS-JEPA)**이라는 시스템을 만들었습니다. 이것이 어떻게 작동하는지 재미있는 비유를 들어 설명하겠습니다.

당신에게 마법의 고무줄(KAN 엣지 함수)이 있다고 상상해 보세요. 기존 방식에서는 고무줄이 단 하나의 특정 지점으로 늘어났을 때만 관찰했습니다. 하지만 고무밥은 그 지점에서 아주 조금만 왼쪽이나 오른쪽으로 움직여도 매우 다르게 행동할 수 있습니다. 어떤 곳에서는 뻣뻣할 수도 있고, 어떤 곳에서는 매우 잘 늘어날 수도 있습니다.

새로운 방식인 FS-JEP는 단 하나의 점만을 보지 않습니다. 대신 AI에게 **"다중 반경 시그니처(multi-radius signature)"**를 예측하도록 요구합니다. 이것은 고무줄의 중심뿐만 아니라 주변의 여섯 가지 지점(마치 작은 헤일로 모양의 측정값처럼)에서 형태를 찍어내는 것과 같습니다. 이 "시그니처"는 AI에게 고무줄이 국소적으로 어떻게 행동하고 있는지를 정확히 알려줍니다.

학습 과정은 "형태 맞히기" 게임처럼 진행됩니다:

마스크 처리된 온라인 브랜치(The Masked Online Branch): AI는 고-무줄의 행동 중 일부가 가려진(masked) 버전을 봅니다. 그리고 전체 "시그니처"(여섯 지점에서의 형태)를 추측해야 합니다.
타겟 브랜치(The Target Branch): "선생님" 역할을 하는 AI(차분하고 꾸준하게 움직이는 멘토와 같은 존재)는 가려지지 않은 전체 고무줄을 보고 올바른 시그니처를 생성합니다.
일치(The Match): 학생 AI는 선생님의 시그니처와 일치시키려고 노력합니다. 형태를 맞히면 학습이 이루어집니다.

결정적으로, AI는 단순히 형태를 추측하는 것이 아니라 자신이 어떤 고무줄을 보고 있는지도 추적합니다. 고무줄이 수천 개나 있기 때문에, 시스템은 "좌표"를 사용하여 학생이 선생님이 보여주는 것과 동일한 고무줄의 형태를 추측하고 있는지 확인합니다. 이는 학습이 정밀하게 이루어지고 서로 뒤섞이지 않도록 보장합니다.

결과: 더 선명한 윤곽선, 더 높은 점수

연구진은 유방 병변의 초음파 영상, 갑상선 스캔, 선 조직의 조직학적 이미지 등 다섯 가지 의료 영상 데이터셋을 대상으로 이 새로운 방법을 테스트했습니다. 그들은 FS-JEPA 방식을 기존의 가장 우수한 KAN 기반 모델 및 다른 표준 의료 AI 모델들과 비교했습니다.

결과는 명확했습니다:

더 높은 정확도: FS-JEPA 방식은 다섯 가지 데이터셋 전체에서 평균 **Dice 점수 83.37%**와 **IoU 75.04%**를 달 기록하며 가장 높은 성적을 거두었습니다.
경쟁 압도: 이 방식은 가장 강력한 경쟁 모델인 UUEKAN보다 Dice 점수에서 +2.25 퍼센트 포인트 앞섰습니다.
추가 비용 없음: 가장 좋은 점은 이 "예측 학습"이 오직 훈련 중에만 일어난다는 것입니다. 일단 AI 훈련이 끝나면, 추가적인 "추측" 부분들은 제거됩니다. 최종 모델은 원래 모델과 크기와 속도가 동일하지만, 훈련 과정에서 더 좋은 습관을 배웠기 때문에 더 똑똑합니다.

이것이 왜 중요한가

이 논문은 학습의 목표를 최종적인 "오케스트라"의 소리에서 개별 악기의 구체적인 "형태"로 옮김으로써, 이러한 유연한 네트워크를 훨씬 더 효과적으로 훈련할 수 있다는 점을 시사합니다. 저자들은 단순히 하나의 점을 예측하는 것만으로는 부족하며, AI가 함수를 진정으로 이해하기 위해서는 "국소적 변화(다중 반경 시그니처)"를 볼 수 있어야 한다는 것을 발견했습니다.

실험을 통해 그들은 단순히 최종 특징을 예측하거나 고무줄의 단일 응답을 예측하는 등의 다른 아이디어들을 검토했습니다. 그러한 방법들은 효과가 없었으며, 이는 함수의 국소적 행동을 포착하는 것이 KAN의 잠재력을 끌어올리는 핵심임을 시사합니다.

궁극적으로 이 연구는 의료 AI 모델을 더 크게 만들거나 복잡하게 만들지 않고도 더 나은 결과를 얻을 수 있음을 보여줍니다. 대신, 우리가 그들을 가르치는 방법을 바꿈으로써—즉, 그들 자신의 내부 함수에 대한 상세한 형태를 예측하도록 함으로써—우리는 그들이 의학에서 가장 중요한 것들의 경계를 훨씬 더 선명하고 정확하게 그리도록 도울 수 있습니다.

기술 요약: 특징이 아닌 함수를 예측하기 – 함수 공간 결합 임베딩 예측 학습을 통한 KAN

문제 정의

의료 영상 분할은 약한 대비, 모호한 경계, 큰 외형 변화와 같은 문제에 직면해 있습니다. U-Net 변형들과 같은 인코더-디코더 아키텍처는 다중 스케일 융합을 통해 이를 해결하지만, 스칼라 가중치와 고정된 활성화 함수에 의존하므로 비선형 변환이 암시적으로만 정의됩니다.

Kolmogorov–Arnold 네트워크(KAN)는 학습 가능한 단변량 함수( $\phi_{j,i}$ )로 에지(edge)를 매개변수화함으로써 명시적인 함수 표현을 제공합니다. 그러나 기존의 KAN 기반 분할 모델(예: U-KAN)은 에지 집계(aggregation) 이후에 정의된 목적 함수를 통해서만 에지 함수를 최적화합니다. 이는 다음과 같은 두 가지 주요 한계를 생성합니다:

명시적 사전 집계 타겟의 부재: 감독(supervision)이 들어오는 에지 응답의 합을 제약하기 때문에, 서로 다른 에지 함수 구성이 동일한 노드 출력을 생성할 수 있습니다. 결과적으로 개별 에지 함수는 집계 전에는 명시적인 학습 타겟을 갖지 못합니다.
국소적 동작 특성화의 불충분: 순전파 과정에서 함수는 단 하나의 입력 의존적 앵커(anchor)에서만 평가됩니다. 동일한 응답을 보이는 함수라도 주변부에서는 서로 다른 국소적 동작을 보일 수 있지만, 이러한 변동성은 다운스트림 태스크에 의해 암시적으로만 제약됩니다.

기존의 결합 임베딩 예측 학습(JEPA) 방식은 특징, 토큰 또는 이미지 임베딩 상에서 작동하지만, 학습 가능한 신경 함수를 예측 타겟으로 취급하지는 않습니다. KAN의 함수 공간으로 JEPA를 확장하는 것은 쉽지 않은데, 단일 에지 응답만으로는 국소적 동작을 특징지을 수 없으며, 샘플링된 에지들이 브랜치 간에 명시적으로 정렬되어 있어야 하기 때문입니다.

방법론: FS-JEPA

저자들은 KAN의 사전 집계 함수 공간으로 예측 학습을 전환하는 프레임워크인 **함수 공간 결합 임베딩 예측 학습(FS-JEPA)**을 제안합니다.

핵심 프레임워크

이 아키텍처는 표준 U자형 KAN 분할 네트워크와 훈련 전용 예측 브랜치로 구성됩니다.

타겟 브랜치(Target Branch): 전체 컨텍스트를 가진 지수 이동 평균(EMA) 브랜치가 타겟 시그니처를 생성합니다.
온라인 브랜치(Online Branch): 부분적으로 마스킹된 토큰 시퀀스를 수신하여 타겟 시그니처를 예측합니다.
에지 대응(Edge Correspondence): 두 브랜치는 동일한 샘플링된 에지 인덱스를 공유하여 예측이 동일한 입력-출력 에지 함수에 대응하도록 보장합니다.

주요 구성 요소

1. 에지 샘플링 및 선택
모든 $O(BN d_{in} d_{out})$ 에지의 응답을 계산하는 과도한 비용을 피하기 위해, 프레임워크는 $K$ 개의 입력-출력 에지 쌍을 무작위로 샘플링합니다. 선택된 에지 연산자( $\Phi_S$ )는 이 특정 에지들을 평가하기 위해 스플라인 파라미터와 매듭 그리드(knot grids)를 수집하며, 이를 통해 보조 활성화 비용을 $O(BNK)$로 줄입니다.

2. 다중 반경 에지 함수 시그니처
단일 격리된 응답을 예측하는 대신, 이 방법은 각 샘플링된 에지 함수를 다중 반경 시그니처로 표현합니다.

에지의 앵커 입력 $u$ 에 대해, 로컬 오프셋 집합 $\mathcal{D} = \{-0.10, -0.05, -0.025, 0.025, 0.05, 0.10\}$ 이 정의됩니다.
시그니처 $s_e$ 는 오프셋 $\Delta$ 에 대한 함수 평가값들의 벡터입니다: $s_e = [\phi_e(u+\Delta)]_{\Delta \in \mathcal{D}}$ .
이 구조화된 표현은 단일 지점 평가가 포착할 수 없는 국소적 함수 변동을 포착합니다.

3. 결정론적 에지 대응
확률적 샘플링 문제를 해결하기 위해, 샘플링된 에지 $(i, j)$ 의 소스 및 타겟 노드 위치를 결정론적 좌표 $(p_i, p_j)$ 로 인코딩합니다. 이 좌표들은 함수 시그니처와 결합되어 시그니처 예측기에 조건으로 제공되며, 학습 가능한 파라미터 없이 일관된 대응을 보장합니다.

4. 결합 최적화
프레임워크는 분할 손실( $L_{seg}$ )과 함수 공간 예측 손실( $L_{sig}$ )을 결합하여 최적화합니다.

예측 손실: $L_{sig} = \text{SmoothL1}(\hat{s}_e, s^+_e)$ 로 정의되며, 여기서 $\hat{s}_e$ 는 온라인 브랜치의 예측이고 $s^+_e$ 는 EMA 브랜치의 타겟입니다.
훈련: 총 손실은 $L_{total} = L_{seg} + \lambda_{deep}L_{deep} + \lambda_{sig}\sum L^{(b)}_{sig}$ 입니다.
추론: 모든 예측 구성 요소(온라인 브랜치, EMA 브래치, 시그니처 예측기)는 제거되며, 원래의 분할 경로만 남습니다.

주요 기여

함수 공간 JEPA: 결합 임베딩 예측 학습을 KAN의 네이티브 사전 집계 함수 공간으로 공식화하여, 예측 타겟을 특징 임베딩에서 구조화된 에지 함수 시그니처로 전환했습니다.
다중 반경 시그니처 표현: KAN 에지 함수의 국소적 함수 동작을 포착하는 표현을 도입하여, 함수 공간 학습을 위한 효과적인 예측 타겟을 제공했습니다.
훈련 전용 예측 프레임워크: 추론 시 추가적인 파라미터나 계산 비용 없이, 무작위 초기화 상태에서 에지 대응을 유지하며 KAN 에지 함수 시그니처와 분할 표현을 공동으로 학습하는 방법을 제시했습니다.

실험 결과

이 방법은 다섯 가지 의료 영상 분할 벤치마크(BUSI, DDTI, TN3K, CVC-ClinicDB, GlaS)에서 평가되었습니다.

성능: FS-JEPA는 다섯 가지 데이터셋 전체에서 가장 높은 평균 Dice 점수(83.37%)와 IoU(75.04%)를 달 기록했습니다.
비교: 가장 강력한 경쟁 KAN 기반 모델인 UUEKAN보다 Dice 기준 +2.25 퍼센트 포인트 높은 성능을 보였습니다.
절제 연구(Ablation):
- FS-JEPA 없이 훈련된 매치된 스캐폴드와 비교했을 때, 제안된 방법은 평균 Dice를 0.85 퍼센트 포인트 향상시켰습니다.
- 예측 타겟에 대한 절제 연구 결과, 다중 반경 에지 함수 시그니처(0.7926 Dice)를 예측하는 것이 단일 에지 응답(0.7771 Dice)이나 토큰/노드 특징을 예측하는 것보다 유의미하게 우수함을 보여주었습니다.
- 메커니즘 분석을 통해 FS-JE-PA가 표준 분할 손실이 놓치는 에지 수준의 변동에 민가하게 반응하도록 KAN 에지 표현을 직접 재형성하는 명시적인 사전 집계 학습 신호를 제공함을 확인했습니다.

의의 및 주장

본 논문은 FS-JEPA가 기존의 특징 임베딩을 넘어 구조화된 학습 가능 신경 함수의 표현으로 예측 학습을 확장하는 것의 효과를 입증한다고 주장합니다. 다중 반경 함수 시그니처를 예측함으로써, 이 방법은 KAN 에지 함수에 대한 직접적인 표현 학습을 가능하게 합니다. 저자들은 성능 향상이 모델 용량의 증가가 아닌 함수 공간 예측 감독에서 기인한다는 점을 강조하며, 추론 시의 파라미터 수가 베이스라인 스캐폴드와 동일함을 명시합니다. 이 프레임-워크는 기존 K-based 비전 모델의 근본적인 한계를 해결하며, 개별 KAN 에지의 함수적 동작을 명시적으로 제약할 수 있는 방법을 제시합니다.