Reverse Distillation: Consistently Scaling Protein Language Model Representations

이 논문은 자연어 처리나 컴퓨터 비전과 달리 단백질 언어 모델 (PLM) 의 확장성이 부진한 문제를 해결하기 위해, 작은 모델의 표현을 큰 모델의 직교 부분 공간으로 분해하여 '마트료시카' 구조를 갖는 역 증류 (Reverse Distillation) 프레임워크를 제안하고, 이를 통해 단백질 벤치마크에서 더 큰 모델이 일관되게 우수한 성능을 발휘하도록 함을 보여줍니다.

핵심

1. 문제: "왜 큰 배가 더 느릴까?" (역설적인 현상)

일반적으로 인공지능은 크기가 커질수록 똑똑해집니다. 하지만 단백질을 다루는 AI 모델에서는 이상한 일이 일어납니다.

• 현상: 중간 크기의 모델 (예: 6.5 억 개 파라미터) 이 가장 성능이 좋고, 가장 거대한 모델 (150 억 개 파라미터) 은 오히려 성능이 떨어집니다.
• 비유: 마치 거대한 도서관을 상상해 보세요.
  • 작은 도서관: 책이 적어서 중요한 정보 (기본적인 문법, 자주 쓰는 단어) 를 빠르게 찾아냅니다.
  • 거대한 도서관: 책이 너무 많아서, 중요한 정보 사이에 쓸데없는 잡동사니나 너무 구체적인 사소한 정보들이 섞여 있습니다.
  • 결과: 도서관이 커질수록 '찾고 싶은 정보'를 골라내는 데 시간이 걸리고, 오히려 엉뚱한 정보에 혼란을 느껴 정답을 못 맞추게 됩니다.

이 논문은 **"큰 모델이 가진 엄청난 정보량이 오히려 방해가 되고 있다"**고 진단합니다.

2. 해결책: "리버스 증류 (Reverse Distillation)"란 무엇인가?

기존의 '지식 증류 (Knowledge Distillation)'는 큰 모델을 작게 압축하는 것이었습니다. 하지만 이 논문은 그 반대로, 작은 모델을 기준으로 큰 모델을 분해하는 방식을 제안합니다.

비유: "요리 레시피의 정리"

• 작은 모델 (기본 요리사): 가장 기본적이고 보편적인 재료 (소금, 설탕, 기본 향신료) 만 다룹니다. 모든 요리에 공통적으로 필요한 '기본 맛'을 잘 알고 있습니다.
• 큰 모델 (천재 요리사): 기본 맛뿐만 아니라, 아주 드문 특수 재료나 복잡한 조미법까지 다룹니다. 하지만 이 드문 재료들이 기본 맛과 뒤섞여 있어서, 요리사가 "이 요리에 정말 필요한 게 뭐지?"라고 헷갈려 합니다.

리버스 증류의 과정:

1. 기본 맛 분리: 작은 요리사 (작은 모델) 가 아는 '기본 맛'을 먼저 추출합니다.
2. 나머지 분리: 큰 요리사 (큰 모델) 가 가진 정보에서 '기본 맛'을 뺀 나머지, 즉 **유일하게 큰 요리사만 아는 '특별한 재료'**만 따로 떼어냅니다.
3. 결합: 이 두 가지를 깔끔하게 나란히 배치합니다.

이렇게 하면 큰 모델의 정보량이 줄어들지 않으면서도, 혼란 없이 필요한 정보만 깔끔하게 정리됩니다.

3. 핵심 기술: "마트료시카 인형" 같은 구조

이 방법의 가장 멋진 점은 '중첩 (Nested)' 구조를 만든다는 것입니다.

• 비유: 마트료시카 인형을 생각해 보세요.
  • 가장 작은 인형 (작은 모델) 을 꺼내면 그 자체로 완성된 인형입니다.
  • 그 안에 들어있는 다음 인형 (중간 모델) 을 꺼내면, 작은 인형이 그대로 포함되면서 새로운 옷 (추가 정보) 이 더해진 형태입니다.
  • 가장 큰 인형 (큰 모델) 을 꺼내면, 작은 인형과 중간 인형이 모두 포함되면서 가장 화려한 장식이 추가된 형태입니다.

이 논문이 만든 **'리버스 증류된 모델'**은 바로 이런 구조입니다.

• 모델의 처음 320 개 숫자만 봐도 작은 모델과 똑같은 성능을 냅니다.
• 처음 1280 개 숫자까지 보면 중간 모델과 똑같은 성능을 냅니다.
• 전체 숫자를 보면 가장 큰 모델의 성능을 냅니다.

즉, 필요한 만큼만 정보를 꺼내 쓸 수 있어 매우 효율적이고, 모델이 클수록 성능이 무조건 좋아지는 예측 가능한 성장을 가능하게 합니다.

4. 실제 성과: "거대 모델이 다시 제자리로"

이 방법을 적용한 결과 (ESM-2 라는 모델 계열에서 실험):

• 기존: 150 억 파라미터 모델이 650 억 파라미터 모델보다 성능이 나빴습니다.
• 리버스 증류 후: 150 억 파라미터 모델이 650 억 파라미터 모델보다 일관되게 더 좋은 성능을 냈습니다.
• 의미: 큰 모델이 가진 '잠재력'이 잡음 (노이즈) 때문에 숨어 있었는데, 이 방법으로 잡음을 제거하고 본래의 힘을 끌어낸 것입니다.

5. 요약: 왜 이 연구가 중요한가?

이 연구는 "크기만 키우는 것"이 답이 아님을 보여줍니다. 대신 **"작은 모델의 지식을 바탕으로 큰 모델의 정보를 정리 (분해) 하는 것"**이 더 중요하다는 것을 증명했습니다.

• 일상적인 교훈: 단순히 많은 정보를 쌓아두는 것 (큰 도서관) 보다, 중요한 정보를 체계적으로 분류하고 정리하는 것 (리버스 증류) 이 훨씬 더 똑똑한 결과를 가져옵니다.
• 미래: 이 기술은 단백질 연구뿐만 아니라, AI 모델이 커질수록 생기는 모든 문제를 해결하는 열쇠가 될 수 있습니다.

한 줄 요약:

"거대한 AI 모델이 혼란스러워하는 이유는 정보가 너무 많아서가 아니라, 정보가 엉켜있어서입니다. 작은 모델을 기준으로 정보를 깔끔하게 분리해 주니, 거대 모델이 다시 제 실력을 발휘하게 되었습니다."

기술 요약

1. 문제 정의 (Problem)

자연어 처리 (NLP) 나 컴퓨터 비전 분야에서는 모델의 크기가 커질수록 성능이 예측 가능하게 향상되는 '스케일링 법칙 (Scaling Laws)'이 잘 성립합니다. 그러나 **단백질 언어 모델 (PLMs, 예: ESM-2)**에서는 이와 반대로 비직관적인 스케일링 행동이 관찰됩니다.

• 성능 저하 현상: 많은 작업 (특히 기능 예측 및 변이 효과 예측) 에서 동일한 계열의 모델 중 중간 크기의 모델이 가장 큰 모델보다 더 좋은 성능을 보입니다. 예를 들어, ESM-2 계열의 경우 6.5 억 (650M) ~30 억 (3B) 파라미터 모델이 정점을 찍고, 150 억 (15B) 파라미터 모델은 오히려 성능이 저하됩니다.
• 임베딩의 불연속성: 서로 다른 크기의 모델에서 생성된 임베딩은 서로 연결되어 있지 않습니다. 자연어 처리의 '마트료시카 (Matryoshka)' 임베딩처럼 큰 모델의 임베딩 앞부분을 잘라내어 작은 모델의 임베딩으로 사용할 수 없으며, 차원을 축소하면 성능이 급격히 떨어집니다.
• 근본 원인: 작은 모델은 용량 제약으로 인해 보편적이고 빈번한 생물학적 특징 (2 차 구조, 소수성 패턴 등) 을 학습하는 반면, 큰 모델은 희귀하고 고차원적인 특징 (가족별 패턴, 에피스타시스 등) 을 추가로 학습하려 합니다. 그러나 이 추가적인 용량이 기존 특징과 얽히게 되어 (entangled), 하류 작업에서 작업 관련 신호를 분리해 내기 어렵게 만들고 노이즈를 증가시킵니다.

2. 방법론 (Methodology)

저자들은 **Reverse Distillation (역 증류)**이라는 새로운 프레임워크를 제안하여 큰 모델의 표현을 작은 모델의 표현을 기반으로 분해합니다.

• 핵심 아이디어: 작은 모델의 표현을 '기저 (Basis)'로 삼고, 큰 모델의 표현에서 이 기저에 직교하는 (orthogonal) 잔차 (residual) 정보를 추출합니다. 이를 통해 서로 다른 규모의 모델 간 특징 간섭을 방지하고, 큰 모델이 가진 고유한 정보를 분리해냅니다.
• 수학적 정의:
  • 작은 모델 $M_r$과 큰 모델 $M_p$ ($|M_r| < |M_p|$) 가 있을 때, 큰 모델의 임베딩 $H_p$를 $H_p \approx [H_r, H_{res}]$로 분해합니다.
  • 여기서 $H_r$은 작은 모델의 임베딩과 동일하며, $H_{res}$는 큰 모델만이 제공하는 직교 잔차 정보입니다.
  • 이 분해는 $M_r$의 표현을 완전히 포함하는 모든 $k_p$ 차원 표현 중 재구성 오차 (MSE) 를 최소화하는 최적의 방법임을 증명합니다 (Eckart-Young 정리에 기반).
• 알고리즘 프로세스:
  1. Phase 1: 작은 모델과 큰 모델에 동일한 시퀀스를 입력하여 임베딩을 생성합니다.
  2. Phase 2: 작은 모델 임베딩을 큰 모델 임베딩으로 예측하는 선형 매핑 ($W^*$) 을 학습합니다 (주성분 회귀, PCR 사용).
  3. Phase 3: 예측 오차 (잔차) 에 SVD(특이값 분해) 를 적용하여 주요 직교 성분을 추출합니다.
  4. Chaining: 이 과정을 모델 계열의 모든 규모에 걸쳐 순차적으로 적용하여, 작은 모델부터 큰 모델까지 계층적인 구조를 만듭니다.
• 결과물 (Matryoshka-style Embeddings): 생성된 임베딩은 '마트료시카' 구조를 가집니다. 큰 모델의 임베딩 앞부분 $k$ 차원은 작은 모델의 임베딩과 정확히 일치하며, 뒤쪽 차원은 큰 모델의 고유 정보를 담고 있습니다.

3. 주요 기여 (Key Contributions)

1. 계층적 분해 (Hierarchical Decomposition): PLM 계열을 각 규모가 직교 정보를 추가하는 계층 구조로 변환하는 방법을 제시했습니다.
2. 단조적 개선 (Monotonic Improvement): 역 증류된 임베딩은 크기가 커질수록 성능이 일관되게 향상되는 스케일링 법칙을 회복시킵니다.
3. 기반 모델 대비 성능 향상: 동일한 임베딩 차원 (예: 1280 차원) 을 가진 경우에도, 역 증류된 모델 (예: rd.650M) 이 기존 ESM-2 650M 모델보다 더 좋은 성능을 보입니다.
4. 일반화 가능성: 특정 모델 계열 (ESM-2) 에 국한되지 않고, 스케일링 문제가 존재하는 모든 모델 계열에 적용 가능한 프레임워크입니다.

4. 실험 결과 (Results)

ProteinGym 벤치마크 (Deep Mutational Scanning, DMS):

• 스케일링 회복: 역 증류된 모델 (rd.650M, rd.3B, rd.15B) 은 각각의 기반 모델보다 일관되게 높은 성능을 보였습니다. 특히 **rd.15B(150 억 파라미터)**가 모든 테스트된 모델 중 가장 강력한 성능을 기록했습니다.
• 비교 우위: 기존 ESM-2 15B 모델은 3B 모델보다 성능이 낮았으나, 역 증류를 적용한 rd.15B 는 rd.3B 보다 높은 성능을 보이며 스케일링 법칙을 성공적으로 복원했습니다.
• 다양한 작업: 2 차 구조 예측 (SSP Q3/Q8), 금속 이온 결합 (MIB), 국소화 (LOC) 등 다양한 단백질 속성 예측 작업에서도 역 증류 모델이 기반 모델을 능가했습니다.

해석 가능성 분석 (Sparse Autoencoders, SAE):

• 역 증류된 임베딩 (rd.35M) 으로 학습된 SAE 는 기반 모델 (ESM-2 35M) 보다 더 많은 GO(Gene Ontology) 용어를 포착했습니다.
• 역 증류된 특징들은 더 구체적이고 기능적으로 관련성이 높은 생물학적 정보를 담고 있는 것으로 확인되었습니다.

추론 시간 (Inference Time):

• 여러 모델을 연쇄적으로 호출해야 하므로 오버헤드가 발생할 것으로 예상되었으나, 작은 모델들의 추론 속도가 빨라 전체적인 시간 증가폭은 제한적이었습니다 (예: rd.15B 는 기본 모델 대비 약 1.7 배의 시간 소요).

5. 의의 및 결론 (Significance)

이 논문은 PLM 의 스케일링 한계가 모델의 표현 능력 부족이 아니라, **표현 용량의 비효율적인 사용 (특징들의 얽힘)**에 기인함을 시사합니다.

• 모델 재학습 불필요: 기존 모델을 다시 학습시키지 않고, 선형 분해만으로 큰 모델의 잠재된 정보를 추출하여 성능을 극대화할 수 있음을 증명했습니다.
• 새로운 패러다임: "큰 모델이 언제 도움이 되는가?"라는 질문에서 "어떻게 규모 간 기여도를 체계적으로 결합할 것인가?"라는 질문으로 연구 방향을 전환시켰습니다.
• 실용적 가치: 역 증류된 임베딩은 계산 효율성 (필요한 차원만 사용) 과 재사용성 (Matryoshka 구조) 을 제공하며, 생물학적 기초 모델뿐만 아니라 다른 도메인의 모델 스케일링 문제 해결에도 적용 가능한 통찰을 제공합니다.

결론적으로, Reverse Distillation은 단백질 언어 모델의 비효율적인 스케일링 문제를 해결하고, 더 큰 모델이 항상 더 나은 성능을 내도록 만드는 체계적인 프레임워크를 제시한 획기적인 연구입니다.