Interpreting Genomic Language Models using Sparse Autoencoders

Motivation

• Known: SAEs는 vision 및 natural language model에서 polysemanticity와 superposition 문제를 해결하여 해석가능한 feature를 추출하는 효과적인 도구로 자리잡았으며, InterPLM 연구를 통해 protein language model(ESM-2)에도 성공적으로 적용된 바 있다. gLM인 Evo2 또한 Goodfire AI와의 협업을 통해 SAE가 DNA, RNA, protein, organism 수준의 생물학적 속성과 대응되는 feature를 합성할 수 있음이 초기에 확인되었다.
• Gap: 기존 Evo2-SAE 연구는 제한된 genomic feature 집합과 단일 model layer만을 다루어 분석 범위가 협소했고, gLM에 대한 SAE 기반 해석가능성을 다양한 layer, 다양한 genomic concept, 그리고 SAE 학습 데이터 속성(진화적 근접성, context length)과의 관계 측면에서 체계적으로 평가한 연구는 부재했다. 또한 genomics는 긴 서열 길이와 불완전한 기능 주석(ground truth)으로 인해 vision/NLP나 단백질 도메인보다 해석적으로 더 도전적인 환경임에도 이를 다루는 표준화된 벤치마크가 없었다.
• Why: gLM은 뉴클레오타이드 돌연변이 효과 예측, 조절 요소 식별 등 다양한 genomic prediction task에서 강력한 성능을 보이지만 내부 메커니즘이 불투명하여 생의학·임상 응용에서 신뢰성 있는 검증이 어렵다. 이 연구는 SAE를 통해 gLM 내부 표현을 생물학적으로 의미 있는 개념과 정렬시킴으로써 모델 해석성을 높이고, 질병 연구(예: SLE) 등 실질적 응용 가능성을 제시한다는 점에서 중요하다.
• Approach: Evo2를 사례 연구 gLM으로 삼아 다양한 layer에서 추출한 embedding에 SAE를 학습시키고, 인간 genome 주석(유전자 구조, 조절 영역, 반복 서열, 보존 영역 등)으로 구성한 벤치마크와 genome 특화 해석가능성 지표를 통해 SAE feature와 raw embedding의 해석가능성을 정량적으로 비교한다.

Achievement

Figure 2. Comparison of SAE features and Evo2 embeddings

1. SAE 기반 해석가능성 벤치마크 구축: promoter, exon, intron, UTR, coding sequence, cis-regulatory element, enhancer, transcription factor binding site, CpG island, repeat element, 진화적 보존 영역 등을 포함하는 다양한 human genomic annotation 데이터셋과 genome-tailored 해석가능성 지표를 마련했다.
2. SAE feature의 우수한 해석가능성 입증: 55개 genomic concept 평가 중 42개(76%)에서 SAE feature가 raw Evo2 embedding보다 더 높은 해석가능성을 보였고, 26개는 F1 score 0.7을 초과했으며, 특히 intermediate layer에서 최상의 성능을 보였다.
3. SAE 학습 데이터 속성의 영향 규명: SAE의 해석가능성이 학습 genome의 진화적 근접성과 sequence context length 같은 학습 데이터 속성에 의존함을 실증적으로 확인했다.
4. Feature atlas 개발: SAE weight를 이용한 baseline보다 우수한 graph 기반 표현 학습 알고리즘을 제안하여, 의미적으로 관련된 genomic feature들을 시각화·군집화할 수 있는 feature atlas를 구축했다.
5. 질병 응용 사례 제시: Systemic Lupus Erythematosus(SLE) 관련 변이가 해석 가능한 SAE feature를 어떻게 교란하는지 분석하여 질병 중심 genomic 탐색에서 SAE 프레임워크의 활용 가능성을 시연했다.

How

Figure 3. Example of probe predictions of exons on a selected human gene. SAEs were trained on E. Coli, C. Elegans, Mous

• GENCODE v49 등에서 유래한 gene structure, 조절 영역, 반복 요소, 진화적 보존 영역 주석으로 구성된 벤치마크 데이터셋 구축
• Evo2(Hyena 아키텍처 기반, 최대 1백만 token context)의 여러 intermediate layer에서 embedding을 추출하고, 이를 이용해 encoder-decoder 구조(ˆx = Wdec(σ(Wencx + benc)) + bdec)의 SAE를 post-hoc 학습
• SAE feature activation signal과 raw embedding을 각각 probe하여 genomic annotation과의 선형 정렬(F1 score 등)을 비교 평가
• SAE 학습에 사용된 genome의 진화적 근접성 및 context length를 달리하여 해석가능성 변화를 분석
• SAE feature의 activation pattern을 기반으로 feature를 공유 공간에 임베딩하는 graph 기반 representation learning 알고리즘을 개발하여 SAE weight 기반 baseline과 성능 비교
• SLE 관련 유전 변이를 SAE feature 공간에 매핑하여 질병 특이적 feature perturbation 분석 수행

Originality

• gLM 해석가능성 연구에서 SAE를 다양한 layer와 다양한 genomic concept에 걸쳐 체계적으로 적용한 최초의 포괄적 연구로, 기존 Evo2-SAE 연구(Goodfire AI 협업)의 제한된 범위를 크게 확장함
• 인간 genome 주석 기반의 표준화된 벤치마크와 genome 특화 해석가능성 지표를 새롭게 제안하여 gLM 해석가능성 평가의 기준을 마련함
• SAE 해석가능성이 학습 데이터의 진화적 근접성 및 context length에 의존한다는 새로운 경험적 통찰을 제공함
• SAE feature를 activation pattern 기반 graph로 표현하여 feature atlas를 구축하는 새로운 representation learning 방법을 제안, 기존 model weight 기반 접근을 능가함
• 질병(SLE) 중심의 실질적 응용 사례를 통해 해석가능성 연구를 넘어 임상적 활용 가능성을 시연함

Limitation & Further Study

• Evo2라는 단일 gLM만을 주요 사례로 다루어, 다른 아키텍처(예: Nucleotide Transformer 등)에 대한 일반화 가능성이 명확하지 않음
• genomics 특유의 불완전한 기능 주석으로 인해 벤치마크가 알려진 개념에 편향되어 있으며, 미지의 기능적 요소나 새로운 생물학적 개념 발견에는 한계가 있을 수 있음
• SLE 사례 분석이 예시적(illustrative) 수준에 그쳐 실제 임상적 검증이나 다른 질병으로의 확장성은 추가 연구가 필요함
• SAE 학습 데이터 속성(진화적 근접성, context length)에 대한 분석이 제한된 조합에서만 이루어져, 보다 폭넓은 hyperparameter 및 데이터 조합에 대한 체계적 탐구가 후속 연구로 요구됨
• feature atlas의 graph 기반 표현 방법이 대규모 SAE dictionary에 대해 확장성(scalability)을 갖는지에 대한 논의가 부족함

Evaluation

총평: gLM 해석가능성이라는 상대적으로 미개척 영역에서 SAE를 체계적으로 적용하고 검증한 견고한 실증 연구로, 벤치마크와 방법론적 기여가 향후 genomic AI 해석가능성 연구의 중요한 토대가 될 것으로 평가된다.