Gaussian processes on ray-guided transformed uniform grids for fast, flexible, and auto-differentiable adaptive source reconstruction in lens modelling

기술 요약: 광선 유도 변환 균등 격자(RTU Grid) 상의 가우시안 프로세스

문제 정의
강한 중력 렌즈 현상(SGL)은 우주론 및 암흑 물질 분포를 제약하는 데 중요한 도구이지만, 견고한 추론을 위해서는 정확한 소스 재구성이 필수적이다. SGL의 근본적인 과제는 소스 평면(source plane) 전반에 걸쳐 정보의 밀도가 불균일하다는 점이다. 예를 들어, 여러 개의 이미지(예: 4중상)를 생성하는 영역은 적은 수의 이미지를 생성하는 영역보다 샘б링 속도가 높으며, 확대율(magnification) 또한 연속적으로 변화한다. 델로네 테셀레이션(Delaunay tessellation)이나 보로노이 빈닝(Voronoi binning)과 같은 전통적인 적응형 메싱 기법은 필요한 곳에 해상도를 배치함으로써 이를 해결하려 한다. 그러나 이러한 방법들은 종종 불연속적인 연산(예: 질량 모델이 변함에 따라 메쉬의 가장자리가 정점을 교체하는 현상)에 의존하여, 현대적 추론 프레임워크에 필요한 미분 가능성을 깨뜨리거나 정규화 선택을 제한할 수 있다. 반대로, 표준 가우시안 프로세스(GP) 사전 확률은 유연한 정규화를 제공하지만, 계산 효율성을 위해 FFT(고속 푸리에 변환)를 활용하려면 정규 격자가 필요하다는 점에서 적응성과 미분 가능성 사이의 긴장을 유발한다.

방법론
저자들은 광선 유도 변환 균등(RTU) 격자를 사용하는 새로운 소스 재구성 접근 방식을 제안한다. 이 방법은 렌즈 데이터가 요구하는 적응형 해상도를 유지하면서도, FFT 기반 GP의 계산 효율성을 결합하는 것을 목표로 한다.

1. RTU 격자 구축:

   • 소스는 변환된 좌표계 $\mathbf{y} = \mathbf{g}(\mathbf{x})$ 상의 균등 격자에서 모델링된다.
   • 변환 $\mathbf{g}$는 이미지 평면에서 소스 평면으로 역추적된 광선(ray)들의 경험적 누적 분포 함수(eCDF)로부터 유도된다.
   • 구체적으로, 저자들은 각 좌표축으로 광선의 위치를 투영하고, 이를 정렬하여 1차원 eCDF를 형성한 뒤, 이 분포에 큐빅 스플라인(cubic spline)을 적합시킨다. 이는 변환이 최소 2회 연속 미분 가능함을 보장한다.
   • eCDF는 신호 대 잡음비 임계값에 의해 정의된 아크 마스크(arc mask) 내에 있는 광선만을 사용하여 구축되므로, 관측된 영역에만 해상도가 할당된다.
   • 이 변환은 소스 평면 좌표를 확장하여, 새로운 $\mathbf{y}$ 좌표계에서는 광선의 밀도가 대략 균일하게 되도록 만든다.

2. 소스 모델링:

   • 소스의 밝기는 이 변환된 균등 격자 상의 가우시안 프로세스로 정의된다.
   • GP는 푸리에 공간에서 백색 잡음을 생성하고, 여기에 제곱근(특정 유연한 Matérn 파워 스펙트럼)을 곱한 뒤 역 FFT를 적용하여 구현된다.
   • 양수 조건은 softplus 함수를 통해 강제되며, 최종 소스 밝기는 더 날카로운 특징을 허용하기 위해 파워 로(power law)에 의해 스케일링된다.
   • 물리적 소스 평면으로의 매핑은 $s_{source}(\mathbf{x}) = s^*_{source}(\mathbf{g}(\mathbf{x}))$를 통해 수행된다.

3. 추론 프레임워크:

   • 이 접근 방식은 JAX로 작성된 자동 미분 가능 포워드 모델링 프레임워크인 Herculens 내에 구현되었다.
   • 추론은 Numpyro를 사용한 확률적 변분 추론(SVI)을 통해 수행되며, 증거 하한(ELBO)을 최적화한다.
   • 이 방법은 FFT 속도를 활용하면서도 정규화를 위한 임의의 파워 스펙트럼 선택을 지원하여 GP의 유연성을 유지한다.

주요 기여

• 미분 가능한 적응성: RTU 격자는 델로네/보로노이 메쉬의 불연속성 문제를 극복하면서도 FFT 기반 GP의 속도를 유지하는, 연속적으로 미분 가능한 적응형 해상도 체계를 제공한다.
• 효율성 경험칙: 저자들은 광선 밀도에 기반한 소스 픽셀 수($N_{src}$) 선택을 위한 휴리스틱을 제안하며, 쌍선형 보간(bilinear interpolation)의 자유도를 맞추기 위해 픽셀당 약 4개의 광선을 사용할 것을 권장한다.
• 구현: 이 방법은 Herculens 프레임워크에 완전히 통합되어 있어, 최첨단 샘플링 알고리즘 및 자동 미분 워크플로우와 호환된다.

결과
저자들은 세 가지 모의 데이터셋(JWST NIRCAM 품질 시뮬레이션)과 실제 렌즈 시스템인 SDSS J0946+1006에 대해 RTU 격자를 테스트했다.

• 적합 품질 및 픽셀 효율성: RTU 격자는 균등 격자와 대등한 이미지 평면 적합 품질($\chi^2_{img}$)을 달성하면서도, 노이즈 수준의 잔차에 도달하기 위해 필요한 픽셀 수를 차원당 약 절반 수준(총 픽셀 수 기준 약 4배 적음)만 사용한다.
• ELBO 성능: 동일한 픽셀 수에 대해, RTU 격자는 일반적으로 균등 격보다 높은 증거 하한(ELBO)을 나타내며, 이는 변분 근사 하에서 데이터에 더 잘 부합함을 의미한다.
• 서브스트럭처(Substructure) 민감도: RTU 격의 사용은 서브스트럭처가 있는 모델과 없는 모델 간의 ELBO 차이에 미미한 영향만을 미친다. 이 방법은 격자 아티팩트로 인해 서브스트럭처에 대한 선호도가 인위적으로 부풀려지지 않으면서도, 서브스트럭처 유무(예: Mock 1 및 SDSS J0946+1006)를 성공적으로 구분해낸다.
• 소스 재구성 정확도: 기저 질량 모델이 정확할 때, RTU 격자는 균등 격 방식보다 더 적은 픽셀로 더 정확한 소스 재구성(더 낮은 $\Delta^2_{src}$)을 생성한다.

의의 및 주장
본 논문은 RTU 격자 방식이 실재성, 유연성, 자동 미분 가능성, 그리고 계산 효율성 사이의 절충안을 해결한다고 주장한다. 소스 픽셀이 더 균일한 수의 광선을 포함하도록 보장함으로써, 이 방법은 현저히 감소된 계산 비용으로 고충실도 소스 재구성을 가능하게 한다. 저자들은 이 접근 방식이 Euclid 탐사와 같이 대규모 렌즈 샘플을 분석해야 하는 경우, 코스몰로지 및 암흑 물질 제약의 편향을 최소화하기 위한 견고하고 유연한 소스 모델링을 제공하는 데 특히 유익하다고 강조한다. 이 방법은 기존의 적응형 메싱 기법에 대한 실질적인 개선책으로서, 서브스트럭처 탐지 능력을 저해하지 않으면서도 빠르고 유연하며 자동 미분이 가능한 대안을 제시한다.