Federated Learning for ICD Classification with Lightweight Models and Pretrained Embeddings

기술 요약: 경량 모델과 사전 학습 임베딩을 활용한 ICD 분류를 위한 연방 학습

문제 제기
임상 기록을 국제 질병 분류 (ICD) 코드로 자동 분류하는 것은 역학 연구 및 의료 행정에서 중요하지만, 수동 작업이 필요하기 때문에 병목 현상으로 남아 있습니다. 대규모 언어 모델 (LLM) 은 이 작업을 자동화하는 데 유망한 것으로 나타났지만, 민감한 환자 데이터의 중앙 집중화를 방지하는 개인정보 보호 규정 (예: GDPR) 과 윤리적 우려로 인해 의료 현장에서의 실제 배포는 방해받고 있습니다. 전통적인 연방 학습 (FL) 접근법은 분산된 노드 간에 방대한 트랜스포머 모델을 미세 조정하는 경우가 많아 계산 비용이 많이 들고, 통신 부하가 크며, 자원이 제한된 병원 환경에서 배포하기 어렵습니다. 또한, 의료 분야의 기존 연방 학습 연구는 사망률 예측과 같은 단순한 작업에 초점을 맞추는 경향이 있어 복잡하고 다중 레이블인 ICD 코딩보다는 덜 다루어집니다.

방법론
저자들은 프라이버시 및 확장성 문제를 해결하기 위해 의미 표현과 분류를 분리하는 모듈형 2 단계 파이프라인을 제안합니다:

1. 오프라인 임베딩 추출: 기관 간 LLM 을 미세 조정하는 대신, 본 연구는 Massive Text Embedding Benchmark 리더보드에서 선정된 6 개의 최신 오픈소스 임베딩 모델을 활용하여 원시 임상 퇴원 요약문을 고정 차원의 조밀한 벡터 (768~1536 차원) 로 인코딩합니다. 이 과정은 각 노드에서 로컬 및 오프라인으로 수행되므로 원시 텍스트가 시설을 떠나는 일이 없습니다. 선정된 모델로는 Qwen3-Embed-0.6B, gte-Qwen2-1.5B-instruct, gte-multilingual-base, snowflake-arctic-embed-l/m-v2.0, nomic-embed-text-v1.5가 포함됩니다.
2. 경량 분류: 생성된 임베딩은 단순하고 가벼운 다층 퍼셉트론 (MLP) 분류기의 입력으로 사용됩니다. 세 가지 아키텍처를 평가했습니다: 표준 MLP, 배치 정규화 및 드롭아웃을 갖춘 세 개의 스택된 레이어를 가진 DeepMLP, 그리고 잔여 연결을 포함하는 DeepResMLP 입니다.
3. 연방 학습 설정: 본 연구는 Flower 프레임워크를 사용하여 20 개의 가상 클라이언트 간 수평 연방 학습 환경을 시뮬레이션했습니다. 데이터는 ICD-9 및 ICD-10 코드로 레이블이 지정된 퇴원 요약문인 MIMIC-IV 데이터셋에서 분할되었습니다. 학습 전략은 30 회의 통신 라운드에 걸쳐 집계하는 FedAvg 를 사용합니다.
4. 실험 설계: 저자들은 중앙 집중식 학습 (베이스라인) 과 연방 학습을 비교하는 광범위한 실험을 수행했습니다. 두 가지 코딩 표준 (ICD-9 및 ICD-10) 하에서 6 개의 임베딩 모델과 3 개의 분류기 아키텍처에 대한 성능을 평가했습니다. 또한, 안정성을 평가하기 위해 10 개의 무작위 계층화 분할을 사용한 제거 연구 (ablation study) 를 수행했습니다.

주요 기여

• 프라이버시 보호 아키텍처: 본 논문은 원시 임상 텍스트가 로컬 노드를 결코 떠나지 않는 실현 가능한 파이프라인을 보여줍니다. 임베딩 모델을 고정하고 연방 방식으로 가벼운 분류 헤드만 학습함으로써, 이 접근법은 엄격한 데이터 주권 요구 사항을 준수합니다.
• 표현과 분류의 분리: 본 연구는 ICD 코딩을 위해 방대한 모델의 엔드 - 투 - 엔드 미세 조정이 필수적이라는 관념에 도전합니다. 고품질의 고정된 임베딩이 충분한 의미적 풍부함을 제공할 수 있으며, 간단한 분류기 (예: MLP) 가 경쟁력 있는 성능을 달성할 수 있다는 가설을 제시합니다.
• 종합적 벤치마킹: 이 작업은 중앙 집중식 및 연방 설정 하에서 6 개의 서로 다른 임베딩 모델과 3 개의 분류기 깊이에 대한 엄격한 평가를 제공하며, 임베딩 품질 대 분류기 복잡성이 다중 레이블 성능에 미치는 영향을 상세히 분석합니다.
• 확장성 분석: 제안된 아키텍처는 최신 엔드 - 투 - 엔드 트랜스포머 (예: 약 139M 파라미터를 가진 PLM-ICD) 에 비해 훨씬 적은 파라미터 (가장 큰 분류기의 경우 710 만 개 미만) 를 사용하는 모델을 사용하여 표준 병원 하드웨어에서의 배포를 가능하게 합니다.

결과

• 임베딩 품질의 우세: 실험 결과, 분류기의 복잡성보다 임베딩 모델의 선택이 예측 성능의 주요 결정 요인임이 드러났습니다. 예를 들어, Qwen3-Embed-0.6B와 간단한 MLP 의 조합은 ICD-10 에서 가장 높은 마이크로 F1 점수 (0.3942) 를 달성한 반면, gte-Qwen2-1.5B-instruct와 DeepResMLP 는 ICD-9 에서 가장 좋은 성능을 보였습니다.
• 분류기 복잡성: 더 깊은 아키텍처 (DeepMLP, DeepResMLP) 는 간단한 MLP 보다 일관되게 우월한 성능을 보이지 않았습니다. 많은 구성에서 추가된 복잡성은 체감하는 수익을 감소시키거나 오히려 성능 저하를 초래하여, 의미적 부담이 사전 학습된 임베딩에 의해 효과적으로 처리됨을 시사합니다.
• 연방 학습 대 중앙 집중식 학습: 연방 학습 결과는 중앙 집중식 학습 베이스라인과 밀접하게 일치했습니다. 예를 들어, Qwen3-Embed-0.6B + MLP 구성은 ICD-10 에서 연방 학습 시 0.3944 의 마이크로 F1 을 달성한 반면, 중앙 집중식 학습에서는 0.3942 를 기록하여 표준화된 임베딩을 사용할 때 데이터 분산이 성능을 크게 저하시키지 않음을 나타냈습니다.
• 빈번한 코드에 대한 성능: 상위 10 개 및 상위 50 개 가장 빈번한 ICD 코드에서 모델은 훨씬 높은 F1 점수 (상위 10 개의 경우 종종 0.60 초과) 를 달성하여 일반적인 진단에 대한 강력한 변별력을 보여주었습니다.
• 안정성: 10 개의 무작위 분할에 대한 제거 연구는 성능의 낮은 분산 (F1 점수의 표준 편차 종종 0.003 미만) 을 보여주어 이 접근법의 견고성을 확인했습니다.

의의 및 주장
저자들은 이 작업이 "확장 가능하고 프라이버시를 고려한 의료 코딩 시스템을 향한 실현 가능한 길"을 제공한다고 주장합니다. 그들은 민감한 데이터를 중앙 집중화해야 하는 비용이 많이 들고 경직된 엔드 - 투 - 엔드 아키텍처에 의존하지 않고도 고성능 의료 코딩 시스템을 구축할 수 있다고 논합니다. 본 연구는 의료 분야에서 "풍부하지만 고립된 데이터의 역설"을 로컬에서 고정된 표준화된 임베딩을 활용하고 경량 헤드만 협력적으로 학습함으로써 해결할 수 있음을 시사합니다.

이 논문은 한계에 대해 겸손한 어조를 유지합니다. 정적이며 미세 조정되지 않은 임베딩에 대한 의존성은 희귀하거나 매우 도메인 특이적인 코드에서 성능을 제한할 수 있음을 인정합니다. 또한, 연방 학습 설정은 실제 병원 네트워크의 이질성을 완전히 반영하지 않을 수 있는 이상적인 조건 (동질적인 데이터 분포, 동기식 업데이트) 을 가정합니다. 저자들은 SOTA 모델의 파라미터의 일부로 경쟁력 있는 결과를 달성하지만, 향후 연구는 견고성을 더욱 향상시키기 위해 적응형 연방 전략, 공동 미세 조정, 도메인 특이적 의료 지식의 통합을 탐구해야 한다고 결론지었습니다.