The Galaxy's Guide to the Tokenizer: A Benchmark for Scientific Foundation Models

기술 요약: 토크나이저에 대한 은하계 가이드

문제 정의
트랜스포머 기반의 파운데이션 모델이 천문학적 데이터로부터 과학적 통찰을 추출하는 주요 패러다임으로 자리 잡았으나, 토크나이저 전략이 학습된 표현(representation)에 미치는 구체적인 영향은 여전히 제대로 이해되지 않고 있습니다. 천문학과 같이 관측 데이터가 알려진 물리적 과정에 의해 지배되며 독립적으로 측정된 정답(예: 적색편이, 항성 질량)을 동반하는 과학적 영역에서는, 원시 관측치를 토큰 시퀀스로 변환하는 방식이 모델의 과학적 추론 능력 형성에 어떻게 기여하는지 평가하는 것이 매우 중요합니다. 기존 문헌은 주로 아키텍처의 선택에 초점을 맞추고 있어, 재구성 충실도(reconstruction fidelity)가 잠재 공간 내 물리 정보의 접근성과 어떻게 상관관계를 갖는지에 대한 이해의 공백을 남겨두고 있습니다.

방법론
저자들은 천문학적 데이터에 대해 자기회귀적으로 학습된 디코더 전용 트랜스포머인 AstroPT 백본을 사용하여, 통제된 통합 프레임워크 내에서 네 가지 서로 다른 토크나이저 전략을 체계적으로 비교합니다. 본 연구는 DESI Legacy Survey (DR8)에서 추출한 640,000개의 은하 이미지를 사용하며, 각 이미지는 256×256 픽셀의 포스트에이지 스탬프(postage stamp)로 처리되었습니다.

평가된 네 가지 토크나이저 전략은 다음과 같습니다:

1. Affine: 이미지 패치에서 임베딩 공간으로의 단순 선형 투영을 사용하는 최소한의 베이스라인입니다.
2. AIM (MLP 기반): 원래의 AstroPT 프레임워크를 따라, 선형 투영을 비선형성을 도입하기 위한 다층 퍼셉트론(MLP)으로 대체합니다.
3. JetFormer: 연속적인 토큰을 생성하기 위해 가역적 정규화 흐름(invertible normalizing flow)을 사용하는 흐름 기반 접근 방식입니다. 이는 별도의 사전 학습된 인코더 없이 이미지 가능도(likelihood)를 최대화하도록 엔드 투 엔드로 학습됩니다.
4. VQ-VAE: 이미지 패치를 학습된 코드북의 시각적 토큰으로 양자화하는 이산적 접근 방식입니다. 인코더는 분포 내(in-distribution) 은하 이미지로 사전 학습된 후, 메인 AstroPT 학습 동안 고정됩니다.

평가 프로토콜
본 연구는 두 가지 주요 평가 축을 채택합니다:

• 물리적 특성 예측: 추출된 임베딩으로부터 선형 및 MLP 프로브(probe)를 학습시켜 167,000개의 홀드아웃 은하의 물리적 특성을 예측합니다. 이러한 특성에는 광도 등급, 색지수, 광학 및 분광 적색편이, 비정상적 별 형성률(sSFR), 항성 질량 및 형태학적 지표가 포함됩니다.
• 재구성 충실도: 5,000개의 무작위 테스트 세트 은하를 대상으로 SSIM과 PSNR을 사용하여 이미지 재구성 품질을 평가합니다.
• 정보 함유량 검증: 정보의 손실과 접근 불가능한 정보를 구분하기 위해, 저자들은 ImageNet으로 사전 학습된 ResNet-50 특징을 사용하여 재구성된 이미지 자체를 프로브함으로써, 픽셀 데이터의 정보 함유량을 임베딩 기하학으로부터 분리하여 조사합니다.

주요 결과
결과는 **재구성 충실도와 물리적 추론을 위한 표현의 품질 사이의 유의미한 디커플링(decoupling)**을 보여줍니다:

• JetFormer: 가장 높은 재구성 충실도(평균 PSNR ~31.11 dB, SS로 SSIM ~0.76)를 달로하며, 미세한 공간적 세부 사항과 낮은 표면 휘도 성분을 보존합니다. 그러나 그 임베딩은 물리적 특성 예측을 위한 선형적 접근성이 가장 낮습니다. 재구성된 이미지를 프로브한 결과, JetFormer가 거의 모든 원래 정보를 유지하고 있음이 확인되었으나, 이 정보가 잠재 공간 내에 "흩어져" 있어 단순한 프로브로는 쉽게 복구할 수 없음을 보여주었습니다.
• VQ-VAE: 물리적 특성 예측, 특히 선형 프로브에서 가장 강력한 성능을 나타냅니다. 이산적 병목(discrete bottleneck)은 물리적 정보를 보다 선형적으로 접근 가능한 방식으로 조직합니다. 그러나 재구성 충실도는 더 낮으며(평균 PSNR ~23.57 dB), 형태학이나 절대 광도와 같은 미세한 세부 사항을 체계적으로 손실하지만 색지수는 보존합니다.
• Affine 및 AIM: 두 방법은 거의 모든 특성에서 거의 동일하게 수행되었으며, 이는 AIM의 MLP 헤드가 트랜스포머 백본이 충분히 표현력이 있는 경우 단순 선형 투영보다 이득이 거의 없음을 시사합니다. 이들은 JetFormer보다 국소적 형태학 정보를 더 잘 보존하지만, 회귀 작업에서는 일반적으로 VQ-VAE보다 성능이 떨어집니다.
• 트레이드오프: 본 연구는 근본적인 트레이드오프를 식별합니다: 부드럽고 조밀하며 정보가 완전한 표현(JetFormer)은 높은 충실도의 재구성을 선호하는 반면, 딱딱하고 패치 기반인 토크나이징 전략(VQ-VAE)은 물리적 추론에 더 직접적으로 해석 가능하고 접근 가능한 표현을 촉진합니다.

의의 및 주장
본 논문은 토크나이징이 중립적인 전처리 단계가 아니라, 무엇이 학습되고 어떻게 조직되는지를 결정하는 결정적인 설계 선택이라고 주장합니다. 저자들은 재구성 품질이 과학적 추론을 위한 표현의 유용성을 나타내는 신뢰할 수 있는 대리 지표(proxy)가 아님을 주장합니다.

독립적으로 측정된 물리량에 근거하여 평가를 수행함으로써, 본 연구는 과학적 데이터가 파운데이션 모델을 위한 객관적이고 순환 논리에 빠지지 않는 벤치마크를 제공할 수 있음을 입증합니다. 저자들은 토크나이저의 선택이 특정 다운스트림 작업에 의해 주도되어야 한다고 제언합니다: 광도 또는 분광 추론에는 VQ-VAE가 유리하며, 이미지 생성 또는 재구성 작업에는 JetFormer가 더 우수합니다. 이 연구는 물리적 정답(ground truth)이 존재하는 다른 도메인에서도 유사한 프레임워크가 가치가 있을 것임을 언급하며, 해석 가능한 벤치마크를 구축하는 데 있어 천문학의 잠재력을 강조합니다.

한계점
저자들은 각 토크나이저를 위한 서로 다른 학습 절차(예: VQ-VAE의 사전 학습 vs JetFormer의 엔드 투 엔드 학습)의 필요성 때문에 비교가 완전히 절제(ablative)된 것은 아니라고 인정합니다. 또한, 선형 및 MLP 프로브에 대한 의존도가 표현에 존재하는 정보를 과소평가할 수 있으나, 재구성된 이미지를 프로브함으로써 극단적인 경우를 완화하고자 했습니다. 본 연구는 데이터 크기나 모델 용량에 따른 스케일링에 대해서는 탐구하지 않았습니다.