🤖 machine learning

Protein contacts are already in the attention: a single-forward-pass alternative to the Categorical Jacobian

이 논문은 약 19회의 순전파(forward pass)를 필요로 하는 계산 비용이 높은 Categorical Jacobian을 통해 추출되었던 단백질 접촉 신호가, 최소한의 레이블된 데이터만으로도 누설이 제거된 벤치마크에서 단 한 번의 순전파만으로 우수하거나 대등한 성능을 달 수 있도록 식별 가능한 소수의 어텐션 헤드에 이미 집중되어 있음을 입증한다.

원저자: Rome Thorstenson

게시일 2026-06-23

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Rome Thorstenson

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

핵심 아이디어: 단백질 "두뇌" 속에서 지름길 찾기

단백질 언어 모델(PLM)을 수백만 개의 단백질 서열을 읽은 매우 똑똑한 학생이라고 상상해 보세요. 이 덕분에 이 학생은 기하학을 명시적으로 배우지 않았음에도 불구하고, 단백질이 어떻게 3차원 구조로 접히는지 학습했습니다.

과학자들은 이 학생의 "마음을 읽어", 단백질의 어느 부분들이 서로 붙어서(접촉하여) 그 3차원 구조를 형성하는지 알아내려고 노력해 왔습니다.

기존 방식 (The "Categorical Jacobian"):
이전의 가장 좋은 방법은 매우 철저하지만 기진맥진할 정도로 힘든 탐정과 같았습니다. 두 부분이 서로 친구인지(붙어 있는지) 알아내기 위해 탐정은 다음과 같이 행동했습니다:

단백질을 가져옵니다.
서열의 첫 번째 글자를 가능한 다른 모든 글자로 바꿉니다 (19번).
학생의 답변이 어떻게 변하는지 관찰합니다.
이 과정을 단백질의 모든 글자에 대해 반복합니다.

만약 단백질이 300개의 글자로 이루어져 있다면, 이 탐정은 단 하나의 답을 얻기 위해 학생에게 테스트를 5,700번 (300개 글자 × 19번의 변화) 수행하게 해야 합니다. 정확하긴 하지만, 믿을 수 없을 정도로 느리고 계산 비용이 많이 듭니다.

새로운 방식 (The "Attention Fusion"):
이 논문은 학생이 단백질을 처음 읽을 때 이미 자신의 노트에 답을 적어 두었다고 주장합니다. 이 노트들을 "어텐션 맵(attention maps)"이라고 부릅니다.

저자들은 만약 우리가 학생의 노트를 살펴본다면, 단백질의 어느 부분이 서로 붙어 있는지에 대한 정보가 이미 그곳에, 단 몇 개의 특정 "헤드(heads)"(학생의 뇌 속에 있는 특정 공책이나 페이지라고 생각하세요)에 집중되어 있다는 것을 발견했습니다.

이 방식은 5,700번의 테스트를 실행하는 대신, 다음과 같이 수행합니다:

학생에게 단백질을 단 한 번 통과시킵니다.
학생이 "서로 붙어 있음"에 대한 메모를 적어둔 특정 페이지(어텐션 헤드)들을 찾아냅니다.
그 페이지들을 하나로 합쳐 평균을 내어 답을 얻습니다.

결과: 새로운 방식은 기존의 탐정 방식만큼 정확하며(때로는 더 뛰어납니다), 단 **한 번의 통과(single pass)**만으로 끝납니다. 이는 전체 공식을 처음부터 다시 유도하는 대신, 교과서 뒷면에 있는 '치트 시트(요약 정리본)'를 찾는 것과 같습니다.

주요 연구 결과 설명

1. "누수(Leakage)" 문제 (부정행위 테스트)

저자들은 이전 연구들이 일종의 "부정행위"를 했을 수도 있다는 점을 주목했습니다. 모델을 평가하는 데 사용된 테스트용 단백질들이 가끔 학생이 학습 중에 이미 암기했던 단백질들과 유사했기 때문입니다. 이는 마치 학생이 숙제에서 이미 봤던 질문들로 기말고사를 치르는 것과 같습니다.

이를 해결하기 위해 저자들은 "누수가 없는(leakage-clean)" 테스트를 만들었습니다. 그들은 학생이 공부를 마친 이후에 발견된 단백질들만을 사용했습니다.

무슨 일이 일어났나? "부정행위"의 이점을 제거하자, 기존의 탐정 방식(Categorical Jacobian)은 성능이 크게 떨어졌습니다. 반면, 새로운 "치트 시트" 방식(Attention Fusion)은 강력함을 유지했습니다.
시사점: 기존 방식은 학생이 특정 테스트 질문을 기억하고 있는 것에 부분적으로 의존했습니다. 새로운 방식은 실제로 구조를 더 잘 이해하고 있습니다.

2. "분산된(Diffuse)" vs "집중된(Concentrated)" 노트

저자들은 모델마다 노트를 작성하는 방식이 다르다는 것을 발견했습니다:

집중형 모델 (Concentrated Models): 어떤 모델은 답을 단 하나의 특정 노트에 적습니다. 그 하나의 노트만 본다면 답을 얻을 수 있습니다.
분산형 모델 (Diffuse Models): 다른 모델들은 답을 여러 노트에 흩뿌려 놓습니다. 단 하나의 노트에는 전체 그림이 담겨 있지 않지만, 이들을 모두 평균 내면 답이 나타납니다.

이 논문은 몇 가지 예시를 살펴봄으로써 여러분이 어떤 유형의 모델을 가지고 있는지 파악할 수 있으며, 그 후에 하나의 노트를 볼지 아니면 열 개의 노트를 평균 낼지를 결정할 수 있음을 보여줍니다.

3. "인과적(Causal)" 막다른 길

저자들은 단백질을 문장처럼 왼쪽에서 오른쪽으로, 뒤를 돌아보지 않고 읽는 다른 유형의 모델("Causal" 모델)을 테스트했습니다.

결과: 기존의 탐정 방식과 새로운 치트 시트 방식 모두 이 모델들에서는 완전히 실패했습니다.
시사점: 이는 단백질이 3차원적으로 어떻게 붙어 있는지를 이해하기 위해서는, 모델이 단순히 단어를 하나씩 읽는 것이 아니라 단백질 전체를 한꺼번에 볼 수 있는 능력(양방향성, bidirectional)이 필요하다는 것을 시사합니다.

4. 속도와 비용

속도의 차이는 엄청납습니다.

기존 방식: 만약 1,000개의 단백질을 분석하고 싶다면, 기존 방식은 단백질당 수천 번의 시뮬레이션을 실행해야 하므로 수백 달러의 컴퓨터 계산 비용이 들 수 있습니다.
새로운 방식: 단 한 번의 실행만 필요하기 때문에 비용이 아주 적게 듭니다. 이는 해변의 모래알을 일일이 손으로 세는 것과 위성 사진을 찍어 간단한 알고리즘으로 전체 양을 추정하는 것의 차이와 같습니다.

요약

이 논문은 단백질 모델들이 이미 단백질이 어떻게 접히는지 알고 있으며, 단지 그 지식을 내부적인 "어텐션" 메커니즘 속에 숨겨두고 있다고 주장합니다. 이러한 특정 "노트"를 직접 읽는 법을 배움으로써, 우리는 수천 번의 값비싼 시뮬레이션을 실행하지 않고도 이전보다 훨씬 빠르고 정확하게 단백질 구조를 예측할 수 있습니다. 이 방식은 모델이 단백히 전체를 한꺼번에 볼 수 있을 때 가장 잘 작동하며, 모델이 본 적 없는 완전히 새로운 데이터로 테스트했을 때도 그 성능이 유지됩니다.

기술 요약: 단백질 접촉 정보는 이미 어텐션(Attention) 내에 존재한다

문제 정의

방대한 양의 레이블 없는 서열 코퍼스로 학습된 단백질 언어 모델(PLMs)은 풍부한 단백질 구조 표현을 학습한다. 핵심적인 질문은 이러한 구조적 지식이 감독 학습(supervised training) 없이도 추출될 수 있는지 여부이다. Zhang 등[2024]은 **표준 범주적 자코비안(Standard Categorical Jacobian, CJ)**을 해결책으로 제안했다. 이는 모든 잔기(residue)를 모든 대안 아미노산으로 섭동(perturbation)시키고, 모델의 마스크드 언어 모델(MLM) 로짓(logits)에서의 반응을 측정함으로써 접촉 맵(contact map)을 재구성하는 방식이다. CJ는 특정 벤치마크에서 높은 정밀도(P@L ≈ 0.80)를 달enc하지만, 길이 $L$ 인 단백질에 대해 약 $19L$ 번의 순전파(forward pass)가 필요하여 계산 비용이 매우 높다. 또한, CJ는 위치별 범주형 헤드(per-position categorical head, 양방향 마스크드-LM)를 가진 아키텍처에 구조적으로 종속되어 있어, 인코더-디코더 또는 인과적(causal) 모델로의 적용이 제한된다.

최근의 평가들은 데이터 누수(data leakage)에 대한 우려를 제기했다. 즉, 프리트레이닝 코퍼스가 평가 세트와 중복되어 성능 지표가 부풀려졌을 가능성이 있다는 것이다. 따라서 계산 효율적이고, 다양한 아키텍처에 적용 가능하며, 프리트레이닝 누수에 강건한 방법론이 필요하다.

방법론

저자들은 범주적 자코비안에 의해 재구성된 접촉 신호가 이미 모델의 어텐션 맵(attention maps) 내에 명시적으로 존재하며, 식별 가능한 작은 어텐션 헤드 집합에 집중되어 있다고 제안한다.

1. Top-K 어텐션 퓨전 (Top-K Attention Fusion)

입력을 섭동시키는 대신, 제안된 방법은 PLM의 어텐션 메커니즘에서 직접 접촉 신호를 추출한다:

추출(Extraction): 각 어텐션 헤드에 대해, 어텐션 맵을 대칭화하고, 대각 성분을 제거(zero-diagonalized)하며, 평균 곱 보정(Average Product Correction, APC)을 사용하여 교정한다.
선택(Selection): 소수의 레이블된 선택 세트(단 10개의 단백질)에서 접촉 예측 능력(top-L/2 장거리 정밀도로 측정)에 따라 헤드를 순위 매긴다.
퓨전(Fusion): 선택된 상위 $K$ 개의 헤드를 단순한 비가중 산술 평균을 사용하여 결합한다. 이는 추론 시 단 한 번의 순전파만을 요구한다.
레이블 효율성(Label Efficiency): 상위 $K$ 개 헤드의 선택만이 유일한 감독 단계이며, 이를 위해 최소한의 레이블된 단백질 예산(선택을 위한 10개, 감독 회귀 모델과의 매칭된 예산 대비 50개)이 필요하다.

2. 표현-CJ (Representation-CJ)

MLM 헤드가 없는 아키텍처(예: 인코더-디코더 또는 인과적 모델)에 대해 CJ와 비교하기 위해, 저자들은 **표현-CJ(Representation-CJ)**를 도입한다. 이는 로짓 반응을 은닉 상태(hidden-state) 반응으로 대체하여 자코비안 구성을 일반화한다:
$J_{repr}[i, a', j] = \| h_{i \to a'}[\ell, j] - h_{unpert}[\ell, j] \|_2$
여기서 $h$ 는 잔기 $i$ 가 아미노산 $a'$ 로 섭동되었을 때 위치 $j$ 의 $\ell$ 번째 레이어에서의 은닉 상태를 나타낸다. 이는 동일한 계산 비용( $19L + 1$ 순전파)을 유지하면서도 은닉 상태를 노출하는 모든 아키텍처에 적용 가능하다.

3. 평가 프로토콜

본 연구는 결과가 프리트레이닝 중복의 산물이 아님을 보장하기 위해 누수-클린(leakage-clean) 평가 전략을 채택한다:

CAMEO-PTA25: 모델의 프리트레이닝 컷오프(2022-09 이후) 이후에 등록된 단백질 체인들로 구성되며, 현재의 UniRef50에 대해 50% identity로 필터링된 주요 평가 세트이다.
분리된 선택(Disjoint Selection): 정보 누수를 방지하기 위해 헤드 선택과 평가는 데이터의 서로 분리된 하위 집합에서 수행된다.
베이스라인(Baselines): 제안된 방법은 표준 범주적 자코비안(logit-CJ), 최적의 단일 어텐션 헤드, 그리고 감독 L1-로지스틱 회귀 베이스라인과 비교된다.

주요 기여

효율성 및 성능: Top-K 어텐션 퓨전 방법은 단 한 번의 순전파만으로 접촉 신호를 복구하며, 테스트된 모든 양방향 모델(ESM-2 패밀리, ESM-1b, ProtT5-XL)에 대해 범주적 자코비안보다 뛰어난 성능을 보인다.
누수-클린 검증: 누수-클린 CAMEO 분할 데이터에서, 퓨전 방법은 CJ보다 유의미하게 높은 성능을 보인다(예: ESM-2-650M에서 9 퍼센트 포인트 차이). 결정적으로, 이 클린 데이터에서 범주적 자코비안은 더 이상 최적의 단일 어텐션 헤드보다 우세하지 않으며, 이는 이전의 인-디스트리뷰션(in-distribution) 우위가 부분적으로는 암기(memorization)에 의한 것이었음을 시사한다.
아키텍처 일반화: 저자들은 어텐션 퓨전을 통해 인코더-디코더 모델(ProtT5-XL)에도 접촉 신호가 존재함을 입증하였으며, 비-MLM 모델로 자코비안 기반 분석을 확장하기 위해 Repr-CJ를 도입하였다.
메커니즘적 통찰 (확산성, Diffuseness): 퓨전할 최적의 헤드 수( $K$ $K$ )는 모델 내 접촉 신호의 "확산성"에 따라 달라진다.
- 집중된 클러스터(Concentrated Clusters): 일부 모델(예: ESM-2-8M, AMPLIFY-350M)은 신호를 하나의 지배적인 헤드에 집중시키며, 더 많은 헤드를 평균하면 신호가 희석된다.
- 확산된 클러스터(Diffuse Clusters): 더 큰 모델(예: ESM-2-650M, ProtT5-XL)은 신호를 여러 헤드에 분산시키며, 합의(consensus)를 복구하기 위해서는 평균화가 필요하다.
인과적 모델의 한계: 퓨전 및 CJ 방법 모두 인과적 언어 모델(ProGen2-large 및 -xlarge)에서 붕괴되는 모습을 보이는데, 이는 대칭적인 쌍 구조(symmetric pair structures)를 인코딩하기 위해 양방향 프리트레이닝이 필수적임을 시사한다.

결과

성능 격차: 누수-클린 데이터에서, 퓨션은 ESM-2-650M에 대해 CJ보다 약 9 퍼센트 포인트 앞선다. 두 방법 모두 인-디스트리뷰션(Zhang eval-200)에서 누수-클린(CAMEO-PTA25) 벤치마크로 이동할 때 절대 정밀도가 30~36 퍼센트 포인트 하락하는데, 이는 기존의 비감독 접촉 예측 결과가 프리트레이닝 중복에 의해 상당히 부풀려졌음을 나타낸다.
감독 vs 비감독: 매칭된 레이블 예산(50개 단백질)에서, 상위 $K$ 개 헤드의 비가중 평균은 동일한 헤드에 대한 감독 L1-로지스틱 회귀 적합(fit)과 통계적으로 구별할 수 없을 정도로 유사한 성능을 보인다. 이는 성능의 주요 동력이 파라메트릭 형태가 아니라 헤드의 선택임을 의미한다.
계산 경제성: 퓨전은 단백질당 한 번의 순전파를 요구하는 반면, CJ는 약 $19L$ 번을 요구한다. 계산 비용의 손익분기점은 20개의 테스트 단백질 미만이며, 따라서 퓨전은 대규모 애플리케이션에 훨씬 더 효율적이다.
Repr-CJ 검증: Repr-CJ는 두 방법이 모두 정의되는 아키텍처에서 logit-CJ와 밀접하게 상관관계(Pearson $r > 0.90$ )를 가지며, 이를 통해 faithful한 일반화임을 입증한다.

의의 및 주장

본 논문은 범주적 자코비안이 새로운 접촉 정보를 추출하는 것이 아니라, 모델의 어텐션 패턴에 이미 표현되어 있는 구조적 쌍 정보를 재구성하는 것이라고 주장한다. 주요 발견은 이 신호가 소수의 어텐션 헤드를 사용하는 단순하고 파라미터가 필요 없는 평균화를 통해 접근 가능하다는 것이며, 이는 섭동 기반 방법보다 계산적으로 우월한 대안을 제공한다.

저자들은 "누수-클린" 평가가 인-디스트리뷰션과 실제 일반화 성능 사이의 상당한 격차를 드러낸다는 점을 강조하며, 향arian 비감독 접촉 예측 벤치마크는 기본적으로 프리트레이닝을 고려한 분할(pretraining-aware splits)을 채택해야 한다고 제언한다. 또한, 어텐션 헤드에 구조적 정보가 할당되는 방식은 아키텍처의 필연성이 아니라 프리트레이닝 과정에서 학습된 속성임을 밝히며, 이는 동일한 모델 패밀리 내에서도(예: ESM-2 스케일) 다양하게 나타난다.

마지막으로, 두 방법 모두 인과적 LM에서 붕로되는 현상은 대칭적인 쌍 구조를 어텐션에 인코딩하기 위해 양방향 프리트레이닝이 결정적인 요소라는 증거로 제시되나, 저자들은 테스트된 인과적 모델의 범위가 제한적이라는 점을 들어 이를 상관관계에 기반한 발견이라고 언급한다.