TAG-lite: a pre-training task-affinity workflow for multi-task ADMET learning with disjoint datasets

기술 요약: 다중 작업 ADMET 학습을 위한 사전 훈련 태스크 친화도 워크플로우인 TAG-lite

문제 정의

다중 작업 학습(MTL)은 공유된 구조를 활용하여 ADMET(흡수, 분포, 대사, 배설 및 독성) 예측을 개선하는 경로를 제공한다. 그러나 MTL의 효능은 태스크 그룹화에 크게 의존한다. 충돌하는 엔드포인트를 함께 학습시키면 그래디언트 충돌이 발생하여 부정적 전이(negative transfer)를 유발한다. 태스크 호환성을 결정하는 기존 방식은 일반적으로 두 가지 유형의 척계에 의존한다: 전이 기반 방식(공유 파라미터를 업데이트한 후의 손실 변화 측정)과 정렬 기반 방식(그래디언트 간의 코사인 유사도 측정).

ADMET 영역에서의 결정적인 한계는 이러한 방법들이 일반적으로 태스크가 동일한 입력 화합물을 공유한다고 가정한다는 점이다. 실제로 ADMET 데이터셋은 서로 다른 엔드포인트가 별개의 화합물 세트에서 측정되는 경우가 많아 서로 분리되어 있는 경우가 흔하다. 기존 연구들은 생물학적 범주화나 쌍별 공동 학습(pairwise co-training)을 통해 이를 우회하려고 시도했으나, 후자의 경우 조합적으로 불가능하다(27개 엔드포인트에 대해 351회의 실행 필요). 또한, 기존의 친화도 척도는 공통된 화합물 배치가 존재하지 않을 때 구조적으로 적용할 수 없다. 더욱이, 이전 문헌에서는 화합물 중첩이 약 30% 미만으로 떨어지면 그래디언트 정렬이 정보를 제공하지 못하게 된다고 제안한다. 이는 많은 ADMET 쌍이 존재하는 영역이다.

방법론: TAG-lite

저자들은 분리된 입력 설정에 특화되어 설계된 태스크 친화도 그룹화의 최소 적응형 모델인 TAG-lite를 제안한다. 이 워크플로우는 5단계로 구성된다:

1. 공유 인코더 학습 (Shared-Encoder Training): 단일 다중 작업 모델(AttentiveFP 인코더 사용)을 모든 27개 엔드포인트에 대해 동일한 손실 가중치를 사용하여 동시에 학습시킨다. 이 학습은 표현(representation)을 안정화하기 위한 "웜업(warm-up)" 단계와 인코더를 고정하는 "측정(measurement)" 단계로 나뉜다.
2. 그래디언트 추출 (Gradient Extraction): 측정 단계 동안, 공유 인코더 파라미터로부터 태스크별 그래디언트를 추출한다. 결정적으로, 배치는 각 태스크에 대해 독립적으로 샘플링된다. 이들은 절차적 인덱스는 공유하지만 완전히 다른 화합물을 포함한다.
3. 친화도 추정 (Affinity Estimation): 태스크 $i$와 $j$ 사이의 친화도는 공유 파라미터 공간에서의 그래디언트 코사인 유사도로 계산된다:
   $$a_{ij} = \frac{1}{B} \sum_{b} \frac{g_i(b) \cdot g_j(b)}{\|g_i(b)\| \|g_j(b)\|}$$
   이 접근 방식은 실현된 손실 변화가 아닌 방향적 호환성을 측정하므로, 공유 화합물 없이도 친화도 추정을 가능하게 한다.
4. 클러스터링 및 그룹화 (Clustering and Grouping): 결과물인 $27 \times 27$ 친화도 행렬을 거리 행렬($D = 1 - A$)로 변환한다. 계층적 클러스터링을 수행하며, 중간 매개체를 통해 이질적인 태스크를 병합하는 "체이닝(chaining)" 현상에 대한 견고성을 위해 평균, 단일 또는 완전 연결법 대신 **Ward 연결법(Ward linkage)**을 선택한다. 이를 통해 태스크를 그룹(예: $k=8$)으로 분할한다.
5. 스크리닝 규칙 (Screening Rule): 어떤 그룹이 공동 학습에 적합한지 결정하기 위해 두 가지 조건의 스크리닝 규칙을 적용한다:
   • 조건 1: 그룹 내 평균 친화도가 양수여야 하며, Youden-최적 임계값 $\tau^* \approx 0.151$보다 높아야 한다.
   • 조건 2: 그룹 내 평균 화합물 중첩(Jaccard 지수)이 $\ge 0.5\%$여야 한다.
     위 기준을 통과하지 못한 그룹은 공동 학습에서 제외된다.

주요 기여

본 논문은 다음과 같은 6가지 구체적인 기여를 한다:

1. 알고리즘 워크플로우: 쌍별 방식에 필요한 수백 번의 실행 대신 단 한 번의 학습 실행만으로 태스크 그룹화를 위한 분리된 입력 그래디언트-친화도 알고리즘인 TAG-lite를 도입한다.
2. 예측 타당성: 공동 학습 전의 친화도 추정치가 AUC 0.745의 곡선 아래 면적(AUC)을 통해 전이의 방향(이득 vs 손해)을 예측할 수 있음을 입증한다.
3. 중첩 특성화: 351개 엔드포인트 쌍에 걸쳐 중첩을 직접 측정하여, 헤비 테일(heavy-tailed) 분포(중앙값 0.99%, 그러나 CYP 기질과 같은 일부 카테고리는 97% 이상의 중첩 공유)를 밝혀냄으로써 ADMET 데이터가 균일하게 희소하다는 견해에 도전한다.
4. 클러스터링 선택: 체이닝을 피하기 위해 ADMET 데이터에는 Ward 연결법이 필수적임을 입증하며, 평균 및 단일 연결법은 유효한 분할을 생성하는 데 실패함을 보여준다.
5. 이중 조건 스크리닝: 친화도와 화합물 중첩을 결합한 스크리닝 규칙을 제안하며, 이들이 상호 보완적인 정보를 제공함을 보여준다.
6. 스크리닝과 수술의 상호보완성: 사전 훈련 스크리닝(그룹 선택용 TAG-lite)과 최적화 중 그래디언트 수술(예: PCGrad)이 서로 보완적임을 보여준다. 즉, 스크리닝은 호환 가능한 그룹을 선택하고, 수술은 해당 그룹 내의 성능을 개선한다.

결과

• 분리성 (Disjointness): 351개 쌍의 Jaccard 중첩 중앙값은 0.99%였으며, 83.2%의 쌍이 5% 미만이었다. 그러나 특정 카테고리(예: CYP 기질)는 높은 중첩(>60%)을 보여, 저중첩 및 고중첩 영역이 혼재되어 있음을 나타냈다.
• 친화도 예측:
  • 친화도는 전이의 방향을 성공적으로 예측했다 (AUC = 0.745).
  • 크기 제한 (Magnitude Limitation): 저중첩 영역에서 전이의 크기를 예측하는 능력은 붕괴되었다. 중첩이 5% 미만일 때, 친화도와 이득 사이의 상관관계는 0과 구별할 수 없었다 ($r = -0.186$). 5% 이상일 때는 상관관계가 강했다 ($r = +0.700$).
  • 판별력: 저중첩 영역에서도 친화도가 크기 예측 능력은 상실했음에도 불구하고, 유익한 전이와 해로운 전이를 판별하는 능력(방향성)은 여전히 높게 유지되었다 (AUC $\approx$ 0.75).
• 그룹 성능:
  • 스크리닝 규칙을 사용하여 두 개의 그룹이 공동 학습 대상으로 승인되었다. 동일 가중치 적용 시, 이 그룹 내 9개 태스크 중 7개가 개선되었다.
  • 승인된 그룹에 PCGrad(그래디언트 수술)를 적용하면 동일 가중치 방식의 혼합된 결과와 달리, 해당 그룹 내 모든 태스크에서 균일한 개선이 나타났다.
  • HydFE 사례: 긍정적인 그룹 내 친화도를 가졌음에도 부정적 전이가 발생한 특정 실패 사례(HydFE)가 나타났다. 이는 HydFE의 작은 데이터셋 크기가 그룹 내 더 큰 파트너들에 의해 압도되었기 때문이며, 이는 그래디언트 크기를 무시하는 코사인 유사도의 한계를 보여준다.
• 이전 주장과의 비교: 본 결과는 그래디언트 정렬이 30% 미만의 중첩에서 정보력이 없다는 Zhang과 Cheng의 주장을 정교화한다. 저자들은 중첩이 낮으면 크기 추정은 제한되지만, 방향 판별은 여전히 견고하다는 점을 발견했다.

의의 및 주장

본 논문은 TAG-lite가 전통적인 전이 기반 방식이 실패하는 분리된 입력 ADMET 데이터셋에서 태스크 그룹화를 위한 실용적이고 계산 효율적인 프로토콜을 제공한다고 주장한다. 주요 의의는 다음과 같다:

• 운영적 통합: 공유 입력이 없는 환경에서 독립적으로 샘플링된 그래디언트를 공통 친화도 행렬로 변환하여 그래디언트 정렬을 성공적으로 적용했다.
• 방향성 스크리닝: 화합물 중첩이 극도로 낮을 때(<5%)에도 친화도가 유익한 태스크 조합을 식별하기 위한 유효한 "Go/No-Go" 결정 도구임을 확립했다.
• "불가능성" 주장의 정교화: 저중첩이 친화도를 무의미하게 만든다는 주장을 해체했다. 저자들은 저중첩이 효과 크기 추정을 제한할 수는 있지만, 방향 분류를 제거하지는 않는다고 주장한다.
• 상호보완성: 사전 훈련 스크리닝(TAG-lite)과 훈련 중 최적화(그래디언트 수술)가 MTL의 서로 다른 측면을 다루며 함께 사용되어야 함을 입증했다.

저자들은 다음과 같은 점을 언급하며 주장에 대해 신중한 태도를 유지한다:

• 임계값($\tau^* \approx 0.16$, 중첩 $\ge 0.5\%$)은 인샘플(in-sample) 작동 지점이며 외부 검증이 필요하다.
• 이 방법은 전이 방향을 예측할 뿐, 보편적인 효과 크기 예측기는 아니다.
• 전역 평균 친화도는 태스크 적합성을 판단하는 데 부적절한 지표이며, 적합성은 특정 파트너에 상대적이다.
• 한계점으로는 크기 가중 친화도 변형의 부재, "안티-친화도(anti-affinity)" 대조 실험의 부재, 그리고 타 아키텍처에서의 재현성 확인 필요성이 있다.

결론적으로, TAG-лоite는 단 한 번의 학습 실행으로부터 태스크 그룹을 도출할 수 있게 하여, 방대한 쌍별 공동 학습 대신 확장 가능한 대안을 제공하며, 분리된 ADMET 데이터셋에서 부정적 전이를 피하기 위한 통계적으로 검증된 메커니즘을 제공한다.