11institutetext: Simon Fraser University
11email<EMAIL_ADDRESS>

# Representing Anatomical Trees by Denoising Diffusion of Implicit Neural
Fields

Ashish Sinha  Ghassan Hamarneh

###### Abstract

Anatomical trees play a central role in clinical diagnosis and treatment
planning. However, accurately representing anatomical trees is challenging due
to their varying and complex topology and geometry. Traditional methods for
representing tree structures, captured using medical imaging, while invaluable
for visualizing vascular and bronchial networks, exhibit drawbacks in terms of
limited resolution, flexibility, and efficiency. Recently, implicit neural
representations (INRs) have emerged as a powerful tool for representing shapes
accurately and efficiently. We propose a novel approach for representing
anatomical trees using INR, while also capturing the distribution of a set of
trees via denoising diffusion in the space of INRs. We accurately capture the
intricate geometries and topologies of anatomical trees at any desired
resolution. Through extensive qualitative and quantitative evaluation, we
demonstrate high-fidelity tree reconstruction with arbitrary resolution yet
compact storage, and versatility across anatomical sites and tree
complexities. The code will be available
here111https://github.com/sinashish/TreeDiffusion.

###### Keywords:

A

natomical Tree Vasculature Bronchial TreeRepresentation LearningNeural
FieldsImplicit Neural Representations Denoising Diffusion.

##  1 Introduction

(a) Per-instance optimization

(b) Isosurface extraction from INR

(c) Learn tree-space via denoising diffusion

(d) Generate novel tree structures

Figure 2: Overview: (a) An INR is optimized for each sample in the training
dataset, and then flattened to a 1D vector. (b) During inference, the INR is
recovered from the flattened vector, followed by MC to extract the underlying
signal. (c) The diffusion transformer takes the flattened vectors as input to
model the diffusion process. After training, novel INRs can be sampled and
used for downstream tasks via (b). (d) Novel tree structures visualized as
mesh.

### Motivation for Studying Anatomical Trees.

Accurate extraction and analysis of anatomical trees, including vasculature
and airways, is crucial for predicting fatal diseases like ischemic heart
disease [36] and aiding clinical tasks like surgical planning [30],
computational fluid dynamics (CFD) [34], and disease prognosis [1]. However,
the application of extracted trees for such purposes may be hindered by their
explicit representation as meshes or volumetric grids.

### Previous Works on Tree Representation.

Prior works have utilized discrete and explicit representations, such as
medial axis [25], minimal paths [14], grammar [15], voxel grids [42], and
meshes [38], to represent anatomical trees. However, all these explicit
representations have their own peculiar shortcomings, requiring complex ad-hoc
machinery [12, 28] to obtain a smooth surface essential for vascular modeling
[1], and/or large memory footprint [19], up to O$(n^{3})$ for voxels, to
preserve fine details. Although rule-based representations, such as L-system,
offer some degree of control over the produced structures [27], the complexity
of rules and parameters and the computational cost of generating and rendering
the shapes limit their use in current medical applications. Therefore, there
is a need for a continuous shape representation that is independent of spatial
resolution, is memory efficient, and can easily be integrated into deep
learning pipelines.

### Implicit Neural Representation (INR).

Recently, INRs (or neural fields, implicit fields) have been proposed to
address the shortcomings of explicit representations, which use multi-layer
perceptrons (MLP) to fit a continuous function that _implicitly_ represents a
signal of interest [21, 23, 29] as level sets.

This approach not only learns a representation of arbitrary shapes with a
small number of parameters while still achieving a high surface accuracy, but
also supports reconstructions at any arbitrary resolution. To this end, we
propose to use INRs for representing anatomical trees and demonstrate their
efficiency and effectiveness in modeling simple and complex, vasculature and
airways, in 2D and 3D, on both synthetic and real data.

### Generative Modeling of Tree Distributions.

While implicit representations and level sets have been widely used for
segmentation [2, 18] and 3D modeling [11, 13] of vessel structures, there are
limited works that address the modeling of the “shape space” of tree
structures. Previous works that attempt to analyze and synthesize tree
structures are either _model-based_ or _data-driven_. _Model-based_ approaches
[8, 9, 33, 39] record the tree topology along with branch attributes [7, 41],
_e.g._ , length and angle. However, for trees with diverse topologies (_e.g._
, different number of branches), a fixed-size representation no longer
suffices, and altering the model size across trees impairs statistical
analysis. _Data-driven_ methods use generative modeling to represent the
distribution of training data and sample the distribution to synthesize novel
instances. GANs and VAEs were used to model the distribution of vessels [6,
37] but did not leverage the descriptive power of INRs as [1, 23] do.
Diffusion models [10] were employed in [4] to learn the distribution on latent
VAE codes of shapes but required large training sets for faithful encoding.
For a more faithful representation, INRs were used in [5], though neither [4,
5] modeled tubular structures or trees.

### Our contributions

To address the shortcomings of previous works, we are the first to make the
following contributions: (1) We train INRs to achieve a faithful
representation of complex anatomical trees and demonstrate their usage in
segmenting trees from medical images; (2) we employ diffusion models on the
space of INR-represented trees for learning tree distributions and generating
plausible novel trees with complex topology; (3) we showcase the versatility
of our approach in representing trees of varying dimensionality, complexity,
and anatomy; and (4) We quantitatively assess our method’s compactness
representation and reconstruction accuracy, at arbitrarily high resolution.

Figure 3: Pipeline: Top: Given a 3D mesh, sampled points and GT occupancies
(inside/outside), an INR $(\theta)$ is optimized to fit to the shape. Bottom:
Optimized INRs are flattened to a 1D vector, fed to the transformer-based
diffusion model $\mathcal{D}(\phi)$ and optimized to predict noise. The novel
INRs can then be sampled from the trained $\mathcal{D}(\phi)$ in the reverse
process.

##  2 Methodology

### Overview.

We propose a diffusion-based generative approach comprising two steps. First,
we optimize an INR realized as an MLP by overfitting its parameters to each
tree in a dataset so that each INR faithfully represents each tree as a neural
occupancy field (Fig. 2(a)). Second, the optimized INRs are flattened into 1D
vectors and used to train a denoising diffusion model (DDM) that captures the
statistical distribution of the trees (Fig. 2(c)). The DDM is used to
synthesize new INRs, via a reverse diffusion process (Fig. 2(b)), which
represents the occupancy field (or a mesh, if marching cubes [17] is applied)
of the tree (Fig. 2(d)).

### Stage 1: Learning per-sample INR.

We represent each tree, in a dataset $\\{S_{i}\\}_{i=1}^{N}$ of $N$ trees, as
a neural occupancy field, or INR, $f(x;\theta):\mathbb{R}^{3}\to[0,1]$,
modelled as an MLP and parameterized by $\theta\in\mathbb{R}^{D}$. Note that,
unlike prior works [23, 40] that involve parameter sharing across the entire
dataset, we optimize an INR separately for each tree sample.

Given $o_{i}(x)$, the ground truth occupancy of $x$ for $S_{i}$, we optimize
$\theta_{i}$ as:

$\operatorname*{arg\,min}_{\theta_{i}}\sum_{x_{j}\in\mathbb{R}^{3}}\|f_{\theta_{i}}(x_{j})-o_{i}(x_{j})\|_{2}.$
(1)

Similar to [21, 22], our MLP architecture is a fully-connected network with
$L$ layers, of hidden size $D$, and ReLU activation functions. To learn high-
fidelity neural representations of anatomical trees, we follow [21] to first
instantiate all training instances $S_{i}$ in Lipschitz domain
$\Omega\in[-1,1]$ and then adaptively sample points and ground-truth
occupancies both inside and outside the 3D surface of $S_{i}$, to supervise
the optimization of each MLP until convergence. LABEL:fig:optimization
illustrates this optimization process.

### Stage 2: Diffusion on INRs.

We train a diffusion model on the space of INRs, _i.e._ , the weights and
biases $\\{\theta_{i}\\}_{i=1}^{N}$. Our transformer-based [35] diffusion
model $\mathcal{D}(\phi)$ is inspired by [5, 24], and takes the flattened 1D
vectors of $\theta_{i}$ as input. However, before passing the input to
$\mathcal{D}(\phi)$, each $\theta_{i}$ undergoes a _layer-by-layer_
decomposition into $k$ tokens by MLPs [24]. This is essential due to the
potential variation in the dimensionality of different layers in $\theta_{i}$
and ensures correspondence of any element $j$ across all trees
$\theta_{i},\forall i,j$. During forward diffusion, we apply noise $\eta_{t}$
at timestep $t$ to each vector $\theta_{i}$ to obtain a noisy vector
$\theta_{i}^{\star}$.

Following [5, 24], the sinusoidal embedding of $t$ along with
$\theta_{i}^{\star}$ is fed to a linear projection layer, whose output is
concatenated with a learnable positional encoding vector to obtain
$\mathcal{G}$. The resultant $\mathcal{G}$ is then passed to the transformer,
which outputs denoised INRs $\hat{\theta_{i}}$. Parameters $\phi$ of
$\mathcal{D}$ are optimized using:

$\operatorname*{arg\,min}_{\phi}\frac{1}{N}\sum_{i=1}^{N}\|\hat{\theta_{i}}-\theta_{i}\|_{2}.$
(2)

After training, we employ DDIM [31] sampling, _i.e._ , a reverse diffusion
process, to sample new INRs by feeding-in noise as $\hat{\theta}_{i}$ to
$\mathcal{D}(\phi)$, and gradually denoising it. Figure 3 shows an
illustration of the overall diffusion process.

##  3 Datasets, Results, and Implementation Details

### Datasets.

We use the following datasets to evaluate our approach: (1) VascuSynth: 120
synthetic 3D vascular trees, ranging from 1 branch to complex trees with
multiple branches and bifurcations [9]; (2) IntRA: 103 3D meshes of
intracranial vasculature extracted from MRA [38]; (3) BraTS: 40 3D MRI brain
scans [20]; (4) HaN-Seg: 20 segmentations of Head&Neck CTA scans [26]; (5)
DRIVE: 20 2D retinal fundus color images with vessel segmentation masks [32];
(6) EXACT: One CT scan of a bronchial tree [16]; (7) WBMRA: One whole body
MRA; and (8) CoW: One circle of Willis mesh. Our evaluation covers the
following aspects.

Figure 4: Compression _vs_ Reconstruction Accuracy: (Top) We
decimate/downscale the mesh/volume to occupy $\approx$ same memory space as
INR, and report the reconstruction error using Chamfer distance (CD) and edge
length loss (Edge) for meshes and INRs, and $L_{1}$ and $L_{2}$ for volumes.
Note the higher error for meshes and volumes w.r.t INRs for the same storage.
(Bottom) Illustration of GT and downsampled volumes. Notice the disconnected
components.

### Fidelity and Compactness.

Following [21], we use Chamfer distance (CD) to evaluate the fidelity of INRs
in representing anatomical trees and summarize it in Tables 3 and 1. Figure 4
shows that we can achieve higher reconstruction accuracy with a smaller memory
footprint. For example, as seen in Table 3 we need 68 MB for volumes and 12 MB
for meshes compared to INRs that only occupy 0.75 MB and 0.63 MB,
respectively, offering $\approx 90\times$ and $\approx 19\times$ compression,
respectively, with minimal loss of reconstruction accuracy. Similarly, Table 1
shows that $10k$ parameters of INR are enough to encode the geometry of a
$128^{3}$ volume of IntRA, resulting in $\approx 220\times$ compression.

Table 1: Quantitative comparison of reconstruction accuracy and compression
ratio on the VascuSynth (V) [9] and IntRA (I) [38] for different
configurations of our MLP architecture. CD denotes chamfer distance
($10^{-3}$) between the reconstructed mesh and the ground truth, whereas
compression denotes the ratio between size of INR and size of the original
input modality, _i.e._ , raw volume. $\uparrow$ denotes higher is better and
$\downarrow$ lower is better.

| $L=1$ |  |  | $L=3$ |  |  | $L=5$ |  |
---|---|---|---|---|---|---|---|---|---
|  |  |  Compression ($\uparrow$) |  CD ($\downarrow$) |  |  |  Compression ($\uparrow$) |  CD ($\downarrow$) |  |  |  Compression ($\uparrow$) |  CD ($\downarrow$)
$D$ |  $\\#$Params (M) |  $\text{Size}_{{inr}}$ (MB) |  V |  I |  V |  I |  $\\#$Params (M) |  $\text{Size}_{{inr}}$ (MB) |  V |  I |  V |  I |  $\\#$Params (M) |  $\text{Size}_{{inr}}$ (MB) |  V |  I |  V |  I
64 | 0.01 |  0.03 |  27.05 |  223.57 |  12.95 |  29.84 | 0.03 |  0.10 |  9.47 |  78.25 |  8.65 |  15.48 | 0.04 |  0.16 |  5.73 |  47.37 |  8.74 |  21.90
128 | 0.03 |  0.13 |  7.32 |  60.50 |  10.36 |  24.21 | 0.10 |  0.38 |  2.49 |  20.55 |  7.09 |  8.33 | 0.17 |  0.63 |  1.50 |  12.37 |  7.99 |  17.85
256 | 0.13 |  0.51 |  1.88 |  15.53 |  9.24 |  18.19 | 0.40 |  1.51 |  0.61 |  5.06 |  6.46 |  7.48 | 0.66 |  2.51 |  0.38 |  3.12 |  7.07 |  11.33
512 | 0.53 |  2.01 |  0.47 |  3.86 |  10.24 |  24.66 | 1.58 |  6.02 |  0.08 |  1.29 |  7.40 |  19.99 | 2.63 |  10.03 |  0.09 |  0.77 |  6.96 |  12.50
1024 | 2.10 |  8.03 |  0.04 |  0.96 |  11.37 |  40.06 | 6.30 |  24.04 |  0.04 |  0.32 |  7.37 |  15.87 | 10.50 |  42.48 |  0.02 |  0.19 |  7.64 |  17.26

### Versatility.

We assess the effectiveness of INRs in representing anatomical trees from
various medical image modalities. Figures 5(a) and 3 qualitatively and
quantitatively show the adaptability of INRs in modeling DRIVE vessel masks,
3D meshes, CTA and MRA volumes from VascuSynth, IntRA, EXACT and CoW. Figure
5(b) depicts how the same MLP architecture can represent trees of varying
complexity.

(a) Various anatomical sites and imaging modalities

(b) Various complexities

Figure 5: Versatility: Visualization of different synthetic and real
anatomical trees represented as an INR from various medical imaging modalities
and organs. We normalize all shapes to $[-1,1]$ and images to $[0,1]$ to
report the reconstruction error using MSE and CD ($\times 10^{-3}$) between
ground truth and the underlying signal extracted from INR.

### Arbitrary Resolution.

Achieving high-quality mesh and volume reconstruction is crucial for the
accurate representation of vasculature, _e.g._ , for fluid flow simulation.
Figure 6 illustrates how INRs can reconstruct anatomical trees from both IntRA
and VascuSynth at various resolutions without the need for adjusting the model
size or retraining.

(a) (b)

Figure 6: Arbitrary Resolution: (a) Comparison of 2x, 4x, and 8x zoom on an
IntRA sample represented as a volumetric grid (top) and INR (bottom). The
resolution for each sub-figure is shown on top as volume3/INR3. Notice the
smoothness of the surface even at 8x zoom. (b) Zoomed-in regions of a
VascuSynth mesh reconstructed from INRs and ground truth at different mesh
resolutions displayed using faces and edges.

(a)

(b)

(c)

(d)

Figure 7: Tree Statistics computed on VascuSynth. (a) Distance between INRs of
trees with similar bifurcations is low and increases as the complexity
increases; (b) t-SNE plot of the space of trees as INRs; (c-d) Histograms of
tortuosity and total length for ground truth and generated trees.

Modality |  Rel. Error ($\%$) | Input Size (MB) | INR Size (MB)
---|---|---|---
DRIVE (RGB) [32] | $0.018$ | $0.37_{\small\pm 0.0055}$ | $0.066_{\downarrow\times 5.60}$
DRIVE (Mask) [32] | $1.204$ | $0.02_{\small\pm 0.0013}$ | $0.003_{\downarrow\times 6.60}$
BraTS [20] | $0.039$ | $68.11_{\small\pm 0.00}$ | $0.753_{\downarrow\times 90.45}$
HAN-Seg [26] | $5.627$ | $12.1_{\small\pm 1.55}$ | $0.630_{\downarrow\times 19.20}$

Table 2: Quantitative results on tree structures present in different medical
imaging modalities represented using INRs. We report the relative percentage
error ($\%$) between the reconstructed signal and ground truth.

Metric | Value
---|---
MMD $\downarrow$ | $13.36_{\small\pm 8.37}$
COV $\uparrow$ | $0.46_{\small\pm 0.11}$
$1-$NNA $(\%)$ $\downarrow$ | $87.49_{\small\pm 8.99}$

Table 3: Quantitative results for novel tree generation on VascuSynth.
$\uparrow$: higher is better; $\downarrow$: lower is better. For $1$-NNA,
$50\%$ is ideal.

### Space of INR-based Trees.

Once trees are represented via INRs, we study the space of INRs. The t-SNE
plot in Figure 7(a) illustrates how trees with similar number of bifurcations,
_i.e._ topological complexities, tend to cluster in the 2D t-SNE space.
However, we observe cases where trees with the same number of branches are
distant, which may be due to how t-SNE projects high-dimensional samples to 2D
thus distort inter-sample distances. Therefore, to further investigate, we
compute the matrix of $L_{2}$ distances between INRs of VascuSynth trees with
$i$ and $j$ bifurcations, where $i,j\in\\{1,\dots,12\\}$ (Fig. 7(b)). The
reported distances are the average across $10$ trees per bifurcation count. We
notice low ($\approx 0$) values along the diagonal, indicating that trees with
the same branch count have similar INRs. As the difference in branch count
increases, the distance increases, visible by the color gradient in Fig. 7(b).
Furthermore, for a constant difference in branch count (_e.g._ ,
$|3-4|=1=|8-9|$), the dissimilarity is more pronounced with larger branch
counts, despite the numerical difference being the same.

### Tree Synthesis.

Training a DDM on INR-based trees results in a model that captures the tree
distribution and can be sampled using DDIM [31] to generate novel trees.
First, we show qualitative results of novel INR-based trees generated by our
model in Figure 9. Next, since the evaluation of unconditional generation of
tree-like structures can be challenging due to lack of direct correspondence
to ground truth data, we follow [4, 5] to quantitatively assess the samples
generated using diffusion. Specifically, we use minimum matching distance
(MMD), coverage (COV), and 1-nearest neighbor accuracy (1-NNA) to measure
quality, diversity, and plausibility, respectively. Moreover, since the
generations are random, we run the evaluation 3 times and report the mean
values in Table 3. Additionally, we follow [6, 37] and report vessel-based
metrics in Figs. 7(c) and 7(d), where we see high similarities in tortuosity
and centerline length histograms between ground truth training data and
synthesized trees.

### INR-based Image Segmentation.

We present two proof-of-concept experiments on how INR representation is
leveraged to perform vessel tree segmentation. Figure 8 depicts the evolution
of INR-represented segmentation masks as they gradually fit to target
vasculature in both CoW and WBMRA.

Circle of Willis
Whole Body MRA

Figure 8: Segmentation. The evolution of INR towards segmenting two vessel
trees. Analogous to Mumford-Shah based segmentation [3], we use a piecewise-
constant version of the INR (via simple thresholding) as it is being optimized
to faithfully represent the original image.

Figure 9: Tree Synthesis. Some novel INRs sampled from the diffusion model,
visualized as meshes.

### Implementation Details.

We used PyTorch and ADAM optimizer with learning rate $\alpha=10^{-3}$ for all
experiments. To fit INRs to trees, we optimize the MLP architecture in [21]
for max $5k$ iterations on 12GB NVIDIA GeForce GTX TITAN X GPU, taking $\sim$1
minute/tree. We train the diffusion model (128-D hidden size, 4 layers, and 4
self-attention modules, as in [24]) on 48GB NVIDIA A40 GPU for $6k$ epochs,
with batch size 8, $0.9\times\alpha$ decay per 200 epoch, $1k$ diffusion time-
steps, linear noise scheduler $\in[10^{-4},10^{-2}]$, taking $\approx 3$ days.

##  4 Conclusion

We presented the first work to use implicit neural fields for faithful
representation of topologically-complex anatomical trees and learnt their
distribution via training a diffusion model in the space of neural fields. We
demonstrated quantitatively and qualitatively the advantages of our method:
versatility (_e.g._ , 2D/3D; vascular/airway; simple/complex topology), lower
memory footprint while achieving highly accurate reconstructions at arbitrary
high resolutions; synthesis of plausible trees; and application to
segmentation of medical images. Further, our representation is amenable to
integration into deep learning pipelines and can be easily transformed into
other shape representations. Future work may include integrating our method
into more advanced deep segmentation pipelines while encoding semantic
annotations of tree parts.

## References

  * [1] Alblas, D., et al.: Going off-grid: Continuous implicit neural representations for 3D vascular modeling. In: International Workshop on Statistical Atlases and Computational Models of the Heart. pp. 79–90. Springer (2022)
  * [2] van Bemmel, C.M., et al.: Level-set based carotid artery segmentation for stenosis grading. In: MICCAI. pp. 36–43. Springer (2002)
  * [3] Chan, T.F., et al.: Active contours without edges. IEEE Transactions on image processing 10(2), 266–277 (2001)
  * [4] Chou, G., et al.: Diffusion-SDF: Conditional generative modeling of signed distance functions. In: CVPR. pp. 2262–2272 (2023)
  * [5] Erkoç, Z., et al.: HyperDiffusion: Generating implicit neural fields with weight-space diffusion. arXiv preprint arXiv:2303.17015 (2023)
  * [6] Feldman, P., et al.: VesselVAE: Recursive variational autoencoders for 3D blood vessel synthesis. In: MICCAI. pp. 67–76. Springer (2023)
  * [7] Feragen, A., et al.: Toward a theory of statistical tree-shape analysis. IEEE TPAMI 35(8), 2008–2021 (2012)
  * [8] Galarreta-Valverde, M.A., et al.: Three-dimensional synthetic blood vessel generation using stochastic L-systems. In: Medical Imaging 2013: Image Processing. vol. 8669, pp. 414–419. SPIE (2013)
  * [9] Hamarneh, G., et al.: VascuSynth: Simulating vascular trees for generating volumetric image data with ground-truth segmentation and tree analysis. Computerized medical imaging and graphics 34(8), 605–616 (2010)
  * [10] Ho, J., et al.: Denoising diffusion probabilistic models. NeurIPS 33, 6840–6851 (2020)
  * [11] Hong, Q., et al.: High-quality vascular modeling and modification with implicit extrusion surfaces for blood flow computations. Computer Methods and Programs in Biomedicine 196, 105598 (2020)
  * [12] Hu, S.M., et al.: Subdivision-based mesh convolution networks. ACM (TOG) 41(3), 1–16 (2022)
  * [13] Kretschmer, J., et al.: Interactive patient-specific vascular modeling with sweep surfaces. IEEE TVCG 19(12), 2828–2837 (2013)
  * [14] Li, H., et al.: Vessels as 4-D curves: Global minimal 4-D paths to extract 3-D tubular surfaces and centerlines. IEEE TMI 26(9), 1213–1223 (2007)
  * [15] Lindenmayer, A., et al.: Mathematical models for cellular interactions in development I. filaments with one-sided inputs. Journal of theoretical biology 18(3), 280–299 (1968)
  * [16] Lo, P., et al.: Extraction of airways from CT (EXACT’09). IEEE TMI 31(11), 2093–2107 (2012)
  * [17] Lorensen, W.E., et al.: Marching cubes: A high resolution 3d surface construction algorithm. In: Seminal graphics: pioneering efforts that shaped the field, pp. 347–353 (1998)
  * [18] Lorigo, L.M., et al.: CURVES: Curve evolution for vessel segmentation. MedIA 5(3), 195–206 (2001)
  * [19] Maturana, D., et al.: VoxNet: A 3D convolutional neural network for real-time object recognition. In: 2015 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). pp. 922–928. IEEE (2015)
  * [20] Menze, B.H., et al.: The multimodal brain tumor image segmentation benchmark (BRATS). IEEE TMI 34(10), 1993–2024 (2014)
  * [21] Mescheder, L., et al.: Occupancy networks: Learning 3D reconstruction in function space. In: CVPR. pp. 4460–4470 (2019)
  * [22] Niemeyer, M., et al.: Differentiable volumetric rendering: Learning implicit 3D representations without 3D supervision. In: CVPR. pp. 3504–3515 (2020)
  * [23] Park, J.J., et al.: DeepSDF: Learning continuous signed distance functions for shape representation. In: CVPR. pp. 165–174 (2019)
  * [24] Peebles, W., et al.: Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209.12892 (2022)
  * [25] Pizaine, G., et al.: Implicit medial representation for vessel segmentation. In: Medical Imaging 2011: Image Processing. vol. 7962, pp. 1184–1190. SPIE (2011)
  * [26] Podobnik, G., et al.: HaN-Seg: The head and neck organ-at-risk CT and MR segmentation dataset. Medical physics 50(3), 1917–1927 (2023)
  * [27] Prusinkiewicz, P., et al.: Synthetic topiary. In: Annual Conference Series. ACM SIGGRAPH, Addison Wesley. vol. 10 (1994)
  * [28] Qi, C.R., et al.: Volumetric and multi-view CNNs for object classification on 3D data. In: CVPR. pp. 5648–5656 (2016)
  * [29] Sitzmann, V., et al.: Implicit neural representations with periodic activation functions. NeurIPS 33, 7462–7473 (2020)
  * [30] Sobocinski, J., et al.: The benefits of EVAR planning using a 3D workstation. European Journal of Vascular and Endovascular Surgery 46(4), 418–423 (2013)
  * [31] Song, J., et al.: Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 (2020)
  * [32] Staal, J., et al.: Ridge-based vessel segmentation in color images of the retina. IEEE TMI 23(4), 501–509 (2004)
  * [33] Talou, G.D.M., et al.: Adaptive constrained constructive optimisation for complex vascularisation processes. Scientific Reports 11(1), 6180 (2021)
  * [34] Tran, K., et al.: Patient-specific computational flow modelling for assessing hemodynamic changes following fenestrated endovascular aneurysm repair. JVS-vascular science 2, 53–69 (2021)
  * [35] Vaswani, A., et al.: Attention is all you need. NeurIPS 30 (2017)
  * [36] WHO, et al.: Global health estimates: Leading causes of death (2019)
  * [37] Wolterink, J.M., et al.: Blood vessel geometry synthesis using generative adversarial networks. arXiv preprint arXiv:1804.04381 (2018)
  * [38] Yang, X., et al.: IntrA: 3D intracranial aneurysm dataset for deep learning. In: CVPR (2020)
  * [39] Zamir, M., et al.: Arterial branching within the confines of fractal L-system formalism. The Journal of general physiology 118(3), 267–276 (2001)
  * [40] Zhang, B., et al.: 3DShape2VecSet: A 3D shape representation for neural fields and generative diffusion models. arXiv preprint arXiv:2301.11445 (2023)
  * [41] Zhao, M., et al.: Leveraging tree statistics for extracting anatomical trees from 3D medical images. In: CRV. pp. 131–138. IEEE (2017)
  * [42] Zhao, M., et al.: Tree-LSTM: using LSTM to encode memory in anatomical tree prediction from 3D images. In: MICCAI Workshop (MLMI). pp. 637–645. Springer (2019)


# AESTRA: Deep Learning for Precise Radial Velocity Estimation in the Presence
of Stellar Activity

Yan Liang Department of Astrophysical Sciences, Princeton University,
Princeton, NJ 08544, USA  Joshua N. Winn Department of Astrophysical Sciences,
Princeton University, Princeton, NJ 08544, USA  Peter Melchior Department of
Astrophysical Sciences, Princeton University, Princeton, NJ 08544, USA  Center
for Statistics & Machine Learning, Princeton University, Princeton, NJ 08544,
USA

###### Abstract

Stellar activity interferes with precise radial velocity measurements and
limits our ability to detect and characterize planets, particularly Earth-like
planets. We introduce AESTRA (Auto-Encoding STellar Radial-velocity and
Activity), a deep learning method for precise radial velocity measurements. It
combines a spectrum auto-encoder, which learns to create realistic models of
the star’s rest-frame spectrum, and a radial-velocity estimator, which learns
to identify true Doppler shifts in the presence of spurious shifts due to
line-profile variations. Being self-supervised, AESTRA does not need “ground
truth” radial velocities for training, making it applicable to exoplanet host
stars for which the truth is unknown. In tests involving 1,000 simulated
spectra, AESTRA can detect planetary signals as low as 0.1 $\rm{m\,s}^{-1}$
even in the presence of 3 $\rm{m\,s}^{-1}$ of activity-induced noise and 0.3
$\rm{m\,s}^{-1}$ of photon noise per spectrum.

High resolution spectroscopy(2096) — Stellar activity(1580) — Radial
velocity(1332)

††software: spender (Melchior et al., 2023), SOAP (Dumusque et al., 2014),
emcee (Foreman-Mackey et al., 2013), astropy (Robitaille et al., 2013)

##  1 Introduction

The search for Earth-like planets in the habitable zones of Sun-like stars is
one of the premier challenges of modern astronomy. Despite significant
advancements in instrumentation and data analysis techniques, such as the
achievement of $\sim$0.3 $\rm{m\,s}^{-1}$ precision in recent ESPRESSO
observations of Proxima Centauri (Mascareño et al., 2020), the detection of
Earth-like planets will require further advances. The next generation of
Extremely Precise Radial Velocity (EPRV) instruments is designed to achieve an
instrumental precision of $\sim$0.1 $\rm{m\,s}^{-1}$ (Blackman et al., 2020;
Hall et al., 2018). However, the emergent spectrum of a star is inherently
variable, leading to apparent variations in the star’s radial velocity that
can mimic or mask planetary signals. This “radial-velocity jitter” arises from
magnetic activity, granulation, and acoustic oscillations, which have
different characteristic amplitudes and timescales (Haywood et al., 2014;
Mascareño et al., 2020; Milbourne et al., 2019). Given that even the quietest
stars exhibit RV jitter of approximately 2 $\rm{m\,s}^{-1}$ (Brems et al.,
2019), it is essential to mitigate the influence of stellar activity signals
on EPRV measurements if we are to have any chance of detecting Earth-like
exoplanets.

Numerous data analysis techniques have been proposed to determine true RVs
from stellar spectra in the presence of activity. A common approach is to
extract RVs using traditional techniques, such as cross-correlation with a
spectral template, and then de-correlate the apparent RV time series with
various activity indices that are derived independently from the spectra
(Milbourne et al., 2019; Giguere et al., 2016). Another approach, sometimes
combined with the previous approach, is to model the RV time series using
Gaussian Process (GP) regression to represent the effects of activity (Haywood
et al., 2014; Mascareño et al., 2020). There are also data-driven techniques,
such as performing shift-invariant operations on the cross-correlation
function (CCF) to isolate activity-induced signals (Collier Cameron et al.,
2021), and performing full spectral modeling to address contamination due to
variability in telluric absorption lines (Bedell et al., 2019). Supervised
machine learning techniques have been applied to recognize spurious RV signals
associated with changes in shape of the CCF (de Beurs et al., 2022).

On the observational side, there has been a growing trend toward more
intensive and longer-term monitoring of stars. It is no longer unusual to see
hundreds of RVs gathered for a single object, and we will probably begin
seeing thousands of RVs in the coming decade (Crass et al., 2021). For
example, the Terra Hunting Experiment will use a high-resolution echelle
spectrograph and an automated 2.5 m telescope to measure RVs of a few dozen
stars on every clear night for at least 10 years (Thompson et al., 2016; Hall
et al., 2018). Such intensive monitoring may open up new possibilities for
stellar activity mitigation by providing larger samples of activity-related
perturbations.

In this paper, we present AESTRA (Auto-Encoding STellar Radial-velocity and
Activity), a deep learning method for distinguishing between true Doppler
shifts and activity-induced spectral perturbations, given a large sample of
stellar spectra covering a variety of the star’s activity states and Doppler
shifts. The architecture of AESTRA combines a neural network for radial-
velocity estimation, and a spectrum auto-encoder called spender (Melchior et
al., 2023; Liang et al., 2023) for activity modeling. AESTRA is designed to
disentangle pure Doppler shifts from apparent Doppler shifts due to spectral
perturbations, without any prior knowledge of the star’s spectrum or orbital
motion.

The spectrum auto-encoder architecture spender was originally built for
analyzing and representing galaxy spectra in the presence of differing
cosmological redshifts. It features an attentive autoencoder that scans over
the wavelength range and compresses the information in each spectrum into a
small number of parameters, which are ideally independent of redshift (Liang
et al., 2023), and a decoder capable of transforming the parameters back into
an accurate reconstruction of the spectrum in the rest frame (i.e., without
any redshift). Spender proved capable of compressing $\sim$500,000 galaxy
spectra from the Sloan Digital Sky Survey (Ahumada et al., 2020) into a low-
dimensional latent space while maintaining high fidelity in the reconstructed
spectra (Melchior et al., 2023; Liang et al., 2023). Cosmological redshifts
are many orders of magnitude larger than the Doppler shifts associated with
exoplanets; for a Earth-like planet around a Sun-like star, $z\sim 10^{-10}$.
Nevertheless, we wondered whether the ability of spender to identify and
isolate intrinsic spectral features while disregarding Doppler shifts and
other artifacts would be applicable to exoplanet detection.

The remainder of the paper is structured as follows. Section 2 describes the
architecture of AESTRA and the loss functions that are used for training.
Section 3 presents the application of AESTRA to simulated datasets. Finally,
Section 4 summarizes our findings, discusses limitations and potential
extensions of AESTRA, and discusses potential applications to solar
spectroscopy and exoplanet surveys.

##  2 Methods

Figure 1:  AESTRA combines a convolutional RV estimator, whose output is
denoted $v_{\rm encode}$, and an attentive auto-encoder that represents
activity-induced spectral features with a latent vector $\mathbf{s}$. The
decoder network expands $\mathbf{s}$ into a rest-frame spectrum
$\mathbf{y}_{\rm act}$, which is shifted into the observed frame to produce
the reconstructed spectrum $\mathbf{y}_{\rm obs}^{\prime}$. The architecture
is trained with a fidelity loss function that ensures $\mathbf{y}_{\rm
obs}\approx\mathbf{y}_{\rm obs}^{\prime}$ and an RV-consistency loss function
that seeks to recover $v_{\rm offset}$ from an artificially Doppler-shifted
“augment” spectrum $\mathbf{y}_{\rm aug}$. The combined loss function is
designed to disentangle activity-induced spectral features and Doppler shift-
induced features.

Figure 1 illustrates the architecture of AESTRA. The input consists of a
collection of hundreds or more of spectra of a single star, obtained over a
range of different activity states and different phases of the orbital motion
of any planets. The input is processed in parallel by a neutral network RV
estimator and a spectrum autoencoder. The RV estimator learns the effect of a
Doppler shift: a pure stretching of the wavelength scale, regardless of any
other details of the spectrum. Training of the RV estimator network does not
require a spectral template or line list, or indeed any prior knowledge about
the star. Instead, we introduce artificial Doppler shifts into the input
spectra, and optimize the RV estimator network to recover these shifts. Within
the RV estimator are filters (convolutional kernels) that sweep across the
residual spectrum. Ideally, during training, the parameters of these filters
are tuned to recognize the patterns that result from Doppler shifts, while
also learning to ignore photon noise. Thus, this part of AESTRA offers an
alternative approach to traditional RV extraction methods involving template
matching and cross-correlation functions.

This RV estimator, trained to recognize pure wavelength stretching, can
accurately determine velocity differences between stellar spectra exhibiting
the same activity state. However, when measuring RV differences between
spectra of different activity states, the RV estimator is subject to biases
due to spectral distortions that mimic Doppler shifts. In the absence of any
“ground truth” about the star and its motion, we rely on a key assumption: the
star’s orbital motion is uncorrelated with its intrinsic variability during
the observation period. Assuming this condition holds, we can mitigate the
activity-induced RV bias by decorrelating the RV estimator output against
activity indicators. These indicators should capture the star’s activity state
while remaining unaffected by Doppler shifts.

In this paper, we introduce an auto-encoder network to generate these activity
indicators through spectrum modeling. The encoder compresses the information
in each stellar spectrum into a small number of “latent parameters”, and the
decoder transforms a set of latent parameters into an accurate reconstruction
of a rest-frame stellar spectrum. The architecture is designed to assign
latent parameters based on spectral perturbations that are independent of the
Doppler shift of the spectrum. In this respect, the latent parameters are
generalizations of traditional activity indicators such as the bisector span
of the CCF or the equivalent width of certain emission lines. This de-trending
technique is explored in more depth in Section 2.3

###  2.1 Architecture

We begin with a collection of normalized one-dimensional input spectra. Each
spectrum is represented by a vector $\mathbf{y}_{\rm
obs}\mathbf{\in}\mathbb{R}^{L}$, giving the intensity at each of $L$
wavelengths $\mathbf{\lambda}_{\rm obs}$. We do not model the Earth’s
barycentric motion, thereby implicitly assuming that the barycentric
corrections have been applied without error. To prepare the input for AESTRA,
we first establish a template spectrum, $\mathbf{b}_{\rm
obs}\in\mathbb{R}^{L}$ that captures the star’s average spectral features.
This template is subtracted from each observed spectrum to create a residual
spectrum $\mathbf{r}_{\rm obs}=(\mathbf{y}_{\rm obs}-\mathbf{b}_{\rm obs})$.
The purpose of this step is to isolate the spectral time variations and reduce
the dynamic range of the input. In our experiments we somewhat arbitrarily
chose one of the simulated spectra with the smallest applied perturbations to
be the template. The exact choice of the template does not matter as long as
it is representative of the typical spectrum. The only effect of small changes
to the template is to apply small time-independent offsets to all of the
residual spectra, which should not interfere with training.

After template subtraction, the residual spectra are processed through two
pipelines: the encoder-decoder network and the RV estimator network.

The goal of the encoder-decoder network is to summarize spectral perturbations
due to stellar activity using a small number of parameters. The residual
spectrum is compressed by a parameterized encoder $f_{\theta}(\cdot)$ into a
latent vector $\mathbf{s}\in\mathbb{R}^{S}$. Here, $\theta$ denotes the
adjustable parameters in the encoder.

$\displaystyle\mathbf{s}=f_{\theta}(\mathbf{r}_{\rm obs}).$ (1)

Following the approach of Melchior et al. (2023), the encoder
$f_{\theta}(\cdot)$ consists of three convolutional layers, an attention
layer, and a three-layer Multi-Layer Perceptron (MLP). A more in-depth
description of these machine learning components can be found in Smith & Geach
(2023). The attention layer is integrated to help the encoder focus on
significant spectral features amidst varying positions and noise. In practice,
the attention layer works by dividing the extracted feature matrix into two
sub-matrices. The first represents the identified features across different
wavelengths, while the second indicates the relative significance of these
features (commonly called “attention weights”). By multiplying these matrices
element wise and then summing over wavelengths, we derive “attended” features
that emphasize the most relevant spectral information.

Conceptually, we may regard the components of the latent vector as
generalizations of traditional activity indicators, such as line widths or
bisector spans. We chose a dimension $S=3$ for the latent space as a balance
between model complexity, which favors large $S$, and interpretability, which
favors small $S$. This relatively low-dimensional latent space simplifies the
interpretation of the latent vector as activity indicators. As an illustration
of the meaning of the latent vector, Figure 2 shows the effect of perturbing
each component, one at a time, using the model to be described in Section 3.3.
In this case, $s_{1}$ appears to be related mainly to line widths, and $s_{2}$
and $s_{3}$ to line asymmetries.

The decoder function $g_{\phi}(\cdot)$ with adjustable parameters $\phi$ is
another three-layer MLP which uses $\mathbf{s}$ to generate an activity
spectrum $\mathbf{y}_{\rm act}$ in the stellar intrinsic rest frame,
represented as a vector of length $L^{\prime}$ with a slightly extended
wavelength coverage ($L^{\prime}>L$):

$\displaystyle\mathbf{y}_{\rm act}=g_{\phi}(\mathbf{s}).$ (2)

Figure 2: Visualization of the effect of perturbing the latent vector based on
a representative spectrum to be described in Section 3. The black curve shows
the activity spectrum $\mathbf{y}_{\rm act}$ generated from a particular
choice of the latent vector $\mathbf{s}$. The colored lines show the resulting
activity spectrum after perturbing each component of the latent vector:
perturbing $s_{1}$ (purple) appears to introduce line width changes, while
perturbations in $s_{2}$ (blue) and $s_{3}$ (orange) produce features related
to line asymmetries.

To encourage the decoder to focus on the activity-induced perturbations, we
add a trainable rest-frame template spectrum ${\mathbf{b}}_{\rm
rest}\in\mathbb{R}^{L^{\prime}}$ to the decoded activity spectrum to produce
the complete rest-frame model:

$\displaystyle\mathbf{y}_{\rm rest}=\mathbf{y}_{\rm act}+\mathbf{b}_{\rm
rest}.$ (3)

The rest-frame template, $\mathbf{b}_{\rm rest}$, represents a hypothetical
spectrum with zero activity ($\mathbf{y}_{\rm act}=0$). The rest-frame
template is distinct from the observed-frame template described earlier. One
may consider $\mathbf{b}_{\rm obs}$ to be an initial guess and
$\mathbf{b}_{\rm rest}$ as an optimized model for the baseline stellar
spectrum.

The RV estimation network $h_{\psi}(\cdot)$, visualized in Figure 1, is a
modified version of the encoder network. A key difference between the encoder
and the RV estimation network is that the encoder includes an attention layer
(Vaswani et al., 2017) – which is designed to help the model identify the most
important patterns of variability irrespective of wavelength shifts – and the
the RV estimator lacks such a layer.

The RV estimator begins with two convolution layers that apply convolutional
filters to the input data that are trained to respond maximally to Doppler
shifts. After each convolution layer, we apply a Parametric Rectified Linear
Unit (PReLU) operation (He et al., 2015), which introduces non-linearity to
the network and helps to enhance the recognized features. Next, max pooling is
applied, down-sampling the features by retaining only the dominant features
from each segment of the residual spectrum. This leaves us with a $64\times
20$ matrix representing 64 output channels and 20 wavelength segments. A
softmax operation is applied in the wavelength direction. This operation helps
to quantify the relative importance to different segments. Then, the resulting
matrix is flattened into a vector with 1,280 elements. These elements are
passed to a MLP with (128, 64, 32) nodes. The MLP learns to estimate Doppler
shifts based on the feature vector summarized by the convolutional layers. The
estimated RV is denoted $v_{\rm encode}$:

$\displaystyle v_{\rm encode}=h_{\psi}(\mathbf{r}_{\rm obs}).$ (4)

Lastly, the rest frame model spectrum $\mathbf{y}_{\rm rest}$ is Doppler-
shifted using the value of $v_{\rm encode}$ determined by the RV estimator and
interpolated to match the observed wavelengths $\mathbf{\lambda}_{\rm
obs}\in\mathbb{R}^{L}$ using cubic spline interpolation. This step accounts
for the motion of the star relative to the observer and enables a direct end-
to-end comparison between the observed and model spectra.

###  2.2 Loss Function

To ensure high-quality spectral reconstruction, we employ the end-to-end
fidelity loss function of Melchior et al. (2023). This function quantifies the
agreement between the input and reconstructed spectra, averaged over a batch
of $N$ spectra each of dimension $L$:

$\displaystyle L_{\rm fid}$
$\displaystyle=\frac{1}{NL}\sum_{i}^{N}\mathbf{w}_{i}\odot(\mathbf{y}_{{\rm
obs},i}-\mathbf{y}_{{\rm obs},i}^{\prime})^{2}$ (5)

Here, $\mathbf{y}_{i,{\rm obs}}$ denotes the $i$-th input spectrum,
$\mathbf{y}_{i,{\rm obs}}^{\prime}$ is the reconstructed spectrum,
$\mathbf{w}_{i}$ is the corresponding weight vector. The operation $\odot$
represents element-wise multiplication, and the squaring of the difference is
also performed element wise.

To achieve consistent RV estimations, we introduce a novel loss term, $L_{\rm
RV}$. To encourage the RV estimator to focus exclusively on the signatures
caused by Doppler-shifting, we exploit data augmentation, a technique used in
machine learning to increase the amount and diversity of training data by
applying various transformations or manipulations to the original dataset. In
this case, we artificially apply a Doppler shift while preserving the spectral
shape. Data augmentation allows us to generate very large samples for training
the RV estimator:

$\begin{split}&L_{\rm
RV}=\frac{1}{N}\sum_{i}^{N}\left[\frac{1}{\sigma_{v}^{2}}(v_{{\rm
aug},i}-v_{{\rm obs},i}-v_{{\rm offset},i})^{2}\right]\\\ &v_{{\rm
obs},i}=h_{\psi}(\mathbf{r}_{{\rm obs},i})\quad v_{{\rm
aug},i}=h_{\psi}(\mathbf{r}_{{\rm aug},i})\end{split}$ (6)

As shown in Eq. 6, we use a weighted mean squared error (MSE) loss to compare
the predicted RV offset between the original and augmented spectra ($v_{{\rm
aug},i}-v_{{\rm obs},i}$) to the true injected offset ($v_{{\rm offset},i}$).
During training, $v_{{\rm offset},i}$ is randomly drawn from a uniform
distribution $v_{\rm offset}\sim\mathcal{U}(-3,3)$ $\rm{m\,s}^{-1}$ , and
$\sigma_{v}$ represents the Cramer-Rao theoretical limiting uncertainty of a
single RV measurement based only on photon noise (see, e.g. Bouchy et al.,
2001; Beatty & Gaudi, 2015). Because the RV estimator is trained on augmented
spectra with artificial Doppler shifts, there is no need to split the data
into training, validation, and test sets. During each iteration of training
(commonly referred to as an “epoch” in the machine learning literature), a new
batch of training data is generated on the fly, ensuring that the model is
consistently exposed to new Doppler shifts.

Lastly, we introduce a regularization term to constrain the decoder’s
flexibility and prevent over-fitting. As both the activity spectrum and rest-
frame baseline are trainable, adjusting $\mathbf{b}_{\rm rest}$ by adding a
constant can be compensated by subtracting the same constant from
$\mathbf{y}_{\rm act}$. To resolve this degeneracy between $\mathbf{y}_{\rm
act}$ and $\mathbf{b}_{\rm rest}$, we define a Ridge regularization loss as
follows:

$\displaystyle L_{\rm reg}$ $\displaystyle=\frac{k_{\rm
reg}}{NL}\sum_{i}^{N}\frac{\mathbf{y}_{{\rm act},i}^{2}}{\sigma_{y}^{2}}$ (7)

The adjustable hyper-parameters $k_{\rm reg}$ and $\sigma_{y}$ determine the
relative weight of the regularization loss. In practice, we set
$\sigma_{y}=0.1$ to achieve a good balance between model flexibility and
trainability. During optimization, we cyclically adjust the weight of
regularization loss $k_{\rm reg}$, increasing from 0 to 1 over a cycle of 1000
iterations, and then setting $k_{\rm reg}$ to 0 and restarting the cycle.

We adopted a two-phase strategy to train AESTRA models. In the first phase,
the RV estimator was trained independently until $L_{\rm RV}$ approached 1,
while leaving the encoder-decoder network and rest-frame baseline unchanged.
In this stage, the RV estimator can accurately measure the velocity offset
between spectra of the same activity state, but the RV estimates are still
subject to activity-induced bias. In the second phase, we jointly optimized
all components of AESTRA—the encoder-decoder network, the rest-frame baseline,
and the RV estimator—using the combined loss function $L_{\rm total}$:

$L_{\rm total}=L_{\rm fid}+L_{\rm RV}+L_{\rm reg}$ (8)

In summary, $L_{\rm total}$ is designed to optimize three key aspects during
training: high reconstruction quality, consistent RV estimation, and shift-
invariant encoding. The fidelity loss, $L_{\rm fid}$, ensures that the
reconstructed spectra closely match the input spectra. The RV loss, $L_{\rm
RV}$, promotes consistency in RV estimations by training the estimator to
focus on signatures induced by stretching, using augmented data. Lastly, the
regularization loss, $L_{\rm reg}$, constrains the decoder’s flexibility and
resolves the degeneracy between the activity spectrum and rest-frame baseline,
improving training efficiency.

Figure 3: Latent distance distribution for randomly selected data pairs
($\Delta\mathbf{s}_{\mathrm{rand}}$, blue), and for data-augment pairs
($\Delta\mathbf{s}_{\mathrm{aug}}$, red), in a dataset of 1,000 simulated
spectra influenced by random activity. Augmented spectra are obtained by
artificially Doppler-shifting the corresponding original spectra. For a
Doppler-shift invariant encoder, the expected value of the augmented latent
distances is zero: $\langle\Delta\mathbf{s}_{\mathrm{aug}}\rangle=0$.

As an aside, we initially thought it would be necessary to add a consistency
loss term to encourage shift-invariant encoding, following Liang et al.
(2023). This term aims to assign the same latent vector to spectra with
identical activity-induced perturbations but varying Doppler shifts by
minimizing the sigmoid-modulated difference between the latent positions of
the original and corresponding augmented spectra:

$\displaystyle L_{\rm c}$ $\displaystyle=\frac{1}{N}\sum_{i}^{N}{\rm
sigmoid}\left[\frac{1}{\sigma_{s}^{2}S}|\mathbf{s}_{i}-\mathbf{s}_{\mathrm{aug},i}|^{2}\right]-0.5$
(9)

However, we found that this additional term was not necessary for the
numerical experiments we have performed. To demonstrate, we used AESTRA
trained without a consistency loss term to encode augmented spectra with
artificial Doppler shifts. As shown in Figure 3, the latent distance between
the original and corresponding augmented spectra (red) exhibit a much lower
dispersion than the latent distances between randomly selected data pairs
(blue), indicating that the latent parameters that AESTRA assigns to Doppler-
shifted spectra depend only weakly on the Doppler shift. We note this fact
here because although the consistency loss term was not necessary, it does not
harm model performance and might be useful when training on real data, for
which spectral perturbations are undoubtedly more complex.

Figure 4:  Distribution of 1,000 simulated spectra in latent space, color-
coded according to $v_{\rm encode}$. A subtle color gradient from the center
to the outskirts indicates a weak correlation between encoded RVs and latent
position. Given that latent position can be interpreted as activity
indicators, this suggests the encoded RVs are largely activity independent. As
a final correction to improve the RV extraction, we perform Gaussian smoothing
on $v_{\rm encode}$ in latent space, allowing for the extraction and
subtraction of the trend.

###  2.3 De-trending RVs with Latent Vectors

The star’s RV estimates should ideally be independent of the spectral
perturbations. However, Figure 4 shows that the encoded RVs are mildly
correlated with latent positions. Causal connections between activity and
orbital motion are unlikely; although hot Jupiters and other close-orbiting
planets may influence a star’s rotation and photospheric activity through
tidal or magnetic interactions, such effects are negligible for habitable-zone
exoplanets, the focus of this study. Nevertheless, even if causal correlations
are absent, Doppler shifts and stellar perturbations may be correlated because
of similarities in timescales — if, for example, the planet’s orbital period
is comparable to the star’s rotation period or the period of a long-term
magnetic activity cycle. To disentangle these effects, data should be
collected in sufficient quantities and over all relevant timescales.

But we suspect another effect is responsible for the mild trend observed in
Figure 4. The RV estimator is trained to determine accurate RV offsets from
pure spectral stretching, leaving the RV zero-points unspecified. The encoded
RVs may be affected by stellar activity and thus contain activity-dependent RV
zero-points $v_{0}$ that can bias our RV estimates. However, we have a trained
activity autoencoder, capable of producing high-quality reconstructions with
varying stellar activity. We thus consider the latents $\mathbf{s}$ as
generalized activity indices and now seek to determine the spurious velocity
zero-points. We make the ansatz:

$\displaystyle v_{{\rm encode},i}=v_{{\rm true},i}+v_{0}(\mathbf{s}_{i})+{\rm
noise.}$ (10)

Assuming that we have enough spectra to observe different true RVs at
approximately fixed activity $\mathbf{s}$, we can isolate a constant baseline

$\begin{split}v_{\rm 0}(\mathbf{s}_{i})&\approx\langle
v\rangle_{i}=\frac{\sum_{j\neq i}w_{ij}v_{{\rm encode},j}}{\sum_{j\neq
i}w_{ij}},\\\ {\rm
where}~{}w_{ij}&=\exp{\left[-\frac{|\mathbf{s}_{i}-\mathbf{s}_{j}|^{2}}{2\sigma^{2}_{R}}\right]}.\end{split}$
(11)

Here, $w_{ij}$ denotes the weight of the encoded RV $v_{{\rm encode},j}$ for
neighbor $j$, and $\sigma_{R}$ denotes the characteristic radius of the latent
distribution. We calculate $\sigma_{R}$ as the average distance among the 10
nearest neighbors in the distribution. This calculation is mathematically
equivalent to Gaussian smoothing in latent space. In theory, each point in the
distribution should have a non-zero weight and be included in the calculation,
but samples farther than 3-$\sigma_{R}$ exert negligible effects. In practice,
the summation is carried out over the 5% nearest neighbors in latent space,
with a minimum of 10 nearest samples.

To obtain our final estimate of the true Doppler RVs, we subtract the
activity-dependent RV zero-points from the encoded RVs:

$\displaystyle v_{{\rm correct},i}=v_{{\rm encode},i}-\langle v\rangle_{i}$
(12)

The resulting corrected RVs, $v_{{\rm correct},i}$, should be closer to the
true Doppler RVs, $v_{{\rm true},i}$. We expect the efficacy of this method to
increase with larger sample sizes because that we can observe the star with at
different Doppler shifts in nearly the same activity state.

##  3 Applying AESTRA to Simulated Data

###  3.1 Generating Synthetic Spectra from SOAP CCFs

To assess the performance of AESTRA, we simulated a large number of sythetic
spectra spectra affected by magnetic activity.

First, we generated a “quiet” spectrum, $\mathbf{y}_{\rm quiet}$, without any
activity perturbations or Doppler shift. We used an evenly sampled wavelength
grid ranging from 5000Å to 5050Å with a wavelength spacing of 0.025Å111 Our
synthetic spectra offer better line sampling than typically seen in Sun-like
stars. Additional tests with narrower line widths (0.05Å) show that AESTRA’s
performance is not significantly impacted by poor line sampling. , resulting
in an observed wavelength vector $\mathbf{\lambda}_{\rm obs}$ of length
$L=2000$.

Figure 5:  CCF ratios (${\rm CCF_{\rm active}/CCF_{\rm quiet}}$) versus
velocity offset ($\Delta\mathbf{v}_{\rm CCF}$) for a solar-type star with a
single active region, calculated by the SOAP 2.0 code. The color of each curve
indicates the brightness variations due to the darkening effect of the active
region at different phases of stellar rotation. The left panels depict the
effect of spots with varying sizes and latitudes, and the right panels show
the impact of plages. The gray shaded area represents the approximate photon
noise level corresponding to a single-spectrum RV precision of 0.3
$\rm{m\,s}^{-1}$ .

To produce realistic absorption lines, we randomly selected 70 line locations
within the observed wavelength range. For each line, we determined a random
width from a Gaussian distribution $r\sim
N(\mu=0.2{\rm\AA},\sigma=0.02{\rm\AA})$ and a random amplitude from a uniform
distribution $a\sim U(0.01,0.70)$. We then calculated the quiet spectrum as
the product of all absorption lines:

$\displaystyle\mathbf{y}_{\rm
quiet}=\prod_{k}^{K}\left(1-a_{k}\exp{\left[-\frac{1}{2}\frac{(\lambda_{\rm
obs}-\lambda_{k})^{2}}{r_{k}^{2}}\right]}\right).$ (13)

The spectral deformations are based on the SOAP 2.0 code (Spot Oscillation And
Planet; Dumusque et al. (2014)). Since SOAP 2.0 only calculates the impact of
spots and plages on the CCF rather than the entire spectrum, we needed to
devise a method to perturb the synthetic spectrum using the CCFs.

We define ${\rm CCF_{quiet}}$ as the quiet CCF of a Sun-like star evaluated at
velocity offsets $\Delta\mathbf{v}_{\rm CCF}$, and ${\rm CCF_{active}}$ as the
CCF affected by activity. We introduced spectral perturbations to the
simulated spectrum by perturbing all the absorption lines using the same ${\rm
CCF_{active}}$. At each epoch, ${\rm CCF_{active}}$ was calculated by randomly
placing four active regions on the star’s surface. The activity type
(spot/plage), size, and phase of each active region were randomly drawn from
the priors provided in Table 1. For each realization of ${\rm CCF_{active}}$,
we defined an activity function $f_{\rm act}$ to measure the activity-induced
perturbation as a function of velocity offset:

$\displaystyle f_{\rm act}(\Delta\mathbf{v}_{\rm CCF}/c)={\rm
CubicSpline}\left(\frac{\Delta\mathbf{v}_{\rm CCF}}{c},\frac{\rm
CCF_{active}}{\rm CCF_{quiet}}\right).$ (14)

Figure 6:  Simulated spectrum of a star affected by multiple active regions
and photon noise ($\mathbf{y}_{\rm obs}$, gray), its reconstruction
($\mathbf{y}^{\prime}_{\rm obs}$, red), and the noise-free underlying spectrum
(black). In the top panel, $\mathbf{y}_{\rm obs}$ and
$\mathbf{y}^{\prime}_{\rm obs}$ are difficult to distinguish because they
overlap nearly perfectly. The bottom panel displays the corresponding residual
spectra, with the gray spectrum predominantly showing photon noise.

Figure 5 shows a collection of ${\rm CCF_{active}/CCF_{quiet}}$ due to a
single active region. To perturb a specific line, we aligned the activity
function $f_{\rm act}$ with the line center $\lambda_{k}$ and interpreted the
wavelength offset as a velocity offset, given by
$\Delta\lambda/\lambda_{k}=\Delta\mathbf{v}_{\rm CCF}/c$. By incorporating the
individual line perturbations, we constructed the entire intrinsic spectrum:

$\displaystyle\mathbf{y}_{\rm intrinsic}=\mathbf{y}_{\rm
quiet}\odot\prod_{k}^{K}f_{\rm act}\left(\frac{\lambda_{\rm
obs}-\lambda_{k}}{\lambda_{k}}\right).$ (15)

Table 1: Priors of the SOAP active region parameters Parameter | Distribution | Prior
---|---|---
Activity type | Binary | $P({\rm Spot})=P({\rm Plage})=0.5$
Size [$R_{\star}$] | Gaussian |  $\mu=0.10$, $\sigma=0.05$
Longitude [∘] | Uniform | $0^{\circ}\leq\phi<360^{\circ}$
Latitude [∘] | Uniform | $-30^{\circ}\leq\theta\leq 30^{\circ}$

For the recovery test, we Doppler-shifted the perturbed intrinsic spectrum
according to the true RV and interpolated it onto the grid of observed
wavelengths using a cubic spline function. To simulate photon noise, we added
independent Gaussian random numbers to the spectrum:

$\displaystyle\mathbf{y}_{\rm obs}={\rm DopplerShift}(\mathbf{y}_{\rm
intrinsic},v_{\rm true})+{\rm noise.}$ (16)

The true planetary RV signal at time $t$ was taken to be $v_{\rm
true}(t)=K\sin{\left({2\pi t}/{P}\right)}$, where $K$ is the velocity
semiamplitude and $P$ is the orbital period. For the following tests, we
adjusted the photon-noise level of the simulated spectra to correspond to a
Cramer-Rao theoretical limit of 0.3 $\rm{m\,s}^{-1}$ on the precision of a
single RV measurement. The value of 0.3 $\rm{m\,s}^{-1}$ was chosen to be
comparable to the anticipated capabilities of frontier EPRV instruments.

###  3.2 Defining RV Estimates for Model Assessment

To evaluate the performance of the AESTRA model, we track four radial velocity
estimates:

Apparent RV ($v_{\rm apparent}$): The apparent RV is computed by finding the
Doppler shift that brings the quiet baseline spectrum and the observed
spectrum into the best possible agreement (minimum $\chi^{2}$). Affected by
stellar activity, this metric provides insight into the RV noise level due to
line distortions.

De-trended RV using traditional activity indicators ($v_{\rm traditional}$):
This quantity represents the traditional technique for activity mitigation
whereby the apparent RVs are corrected to control for the observed variations
in activity indicators that are derived from the spectra. We use a quiet
spectrum ($\mathbf{y}_{\rm quiet}$) as a template to extract the CCF of the
observed, activity-perturbed spectrum ($\mathbf{y}_{\rm obs}$). Three activity
indicators are measured: the bisector span, the full-width at half-maximum of
the CCF (FWHM), and the depth of the CCF. Following Dumusque et al. (2014),
the bisector span is defined as the difference between the top (10%$-$40%
depth) and bottom (60%$-$90% depth) bisectors. Multilinear regression is
performed between the time series of $v_{\rm apparent}$ and the time series of
the three activity indicators, and the results are used to subtract the
estimated contributions to the RV variations due to changes in the activity
indicators.

Corrected RV ($v_{\rm correct}$): This quantity is derived from the AESTRA
model by de-trending the output of the RV estimator ($v_{\rm encode}$) against
the components of the latent vector, i.e., the generalized activity indicators
(see Section 2.3 for details). These values can be used to assess the model’s
performance in isolating true Doppler RVs from activity-induced RV offsets.

Reference RV ($v_{\rm ref}$): Serving as a benchmark, the reference RV is
calculated by fitting the noise-free underlying rest-frame spectrum
($\mathbf{y}_{\rm intrinsic}$) to the observed spectrum ($\mathbf{y}_{\rm
obs}$). The scatter in the reference RVs around the true RVs is slightly
higher than the Cramer-Rao bound of 0.30 $\rm{m\,s}^{-1}$ , probably due to
interpolation imperfections. This serves as a reference for the best possible
RV precision in the absence of stellar activity.

For the $i$-th input spectrum, let $v_{i}$ denote the RV measurement and
$v_{{\rm true},i}$ denote the true planetary RV. We define RV scatter as the
standard deviation of the residual RV ($\Delta v_{i}=v_{i}-v_{{\rm true},i}$):

$\displaystyle(\Delta v)_{\rm std}\equiv\sqrt{\frac{1}{N_{\rm
spec}}\sum_{i}(\Delta v_{i}-({\Delta v})_{\rm avg})^{2}}$ (17)

To assess the model’s performance, we compare $v_{\rm correct}$ to both
$v_{\rm traditional}$ and $v_{\rm ref}$. An effective model should yield
values for $v_{\rm correct}$ that approach the benchmark values of $v_{\rm
ref}$, and an advantageous model should reduce the RV scatter below the
scatter in $v_{\rm traditional}$.

###  3.3 Case I: Time-Independent Activity

Our overall goal was to see if AESTRA could recover 0.10 $\rm{m\,s}^{-1}$
Earth-like signals in the presence of RV jitter that is approximately 30 times
larger in amplitude. First, we considered spectral perturbations that are
randomly and independently assigned to each spectrum. We generated 1,000
synthetic spectra using SOAP CCFs, with each spectrum affected by four active
regions randomly drawn from the priors in Table 1, without any temporal
structure imposed. The apparent RVs, derived by fitting a quiet baseline to
the simulated spectra, are predominantly influenced by random stellar
activity.

We trained an AESTRA model using a two-phase strategy as described in Section
2.2.222For this case and for the others described in this paper, training
required several hours on an NVIDIA A100 GPU. Table 2 summarizes the training
results, and Figure 6 illustrates an example spectrum and its reconstruction.
The AESTRA model achieved high reconstruction quality and RV consistency, with
$L_{\rm fid}\approx 1$ and $L_{\rm RV}<1$. The fact that the self-reported RV
precision was better than the photon noise limit suggests that the RV
estimator performed well and was not the limiting factor in the RV
determinations. Instead, the achievable RV precision was mainly limited by the
decorrelation process.

Table 2: Model performance on simulated datasets with random stellar activity (Case I) and starspot evolution (Case II). $L_{\rm reg}$ is evaluated assuming $k_{\rm reg}=1$. Model | $\mathbf{N_{spec}}$ |  $\mathbf{K}$ [$\rm{m\,s}^{-1}$ ] | $\mathbf{L_{fid}}$ | $\mathbf{L_{RV}}$ | $\mathbf{L_{reg}}$
---|---|---|---|---|---
Case I | 1000 | 0.10 | 1.007 | 0.587 | 0.025
Case I | 100 | 0.30 | 0.973 | 0.559 | 0.034
Case II | 200 | 0.30 | 0.982 | 0.604 | 0.025
Figure 7: Top: Time series of apparent RV (black, left y-axis) and the
fraction of active area on the stellar surface (orange, right y-axis),
dominated by random stellar activity. Middle: The corrected model RVs (black)
are compared to the 0.1 $\rm{m\,s}^{-1}$ true planetary signal (red). Bottom:
The phase-folded and re-binned model RVs (black) are in good agreement with
the true planetary signals (red).  Figure 8:  Top: Lomb-Scargle periodogram of
the apparent RVs ($v_{\rm apparent}$, without any correction) affected by four
random active regions (spots and plages). Middle: De-trended RVs using
traditional indicators ($v_{\rm traditional}$) reveal a subtle peak at the
true period (red dashed line). Bottom: AESTRA corrected RVs ($v_{\rm
correct}$) show a prominent peak at the correct period. Figure 9: Corner plot
illustrating the posterior distribution of the orbital parameters for 1,000
synthetic spectra affected by random activity, with truth values indicated by
blue lines.

Table 3 summarizes the performance of different methods for extracting RVs
from the simulated dataset of 1,000 spectra. The method involving traditional
activity indicators managed to reduce the apparent RV scatter from 3.38
$\rm{m\,s}^{-1}$ to 0.98 $\rm{m\,s}^{-1}$ . The AESTRA model achieved a
further reduction of the RV scatter to 0.46 $\rm{m\,s}^{-1}$ , closely
approaching the RV benchmark for this dataset. Figure 8 presents Lomb-Scargle
periodograms of the RVs de-trended by traditional activity indicators ($v_{\rm
traditional}$) and those corrected by AESTRA ($v_{\rm correct}$). While the
traditional analysis yields a peak at the true period, there are also spurious
peaks with similar amplitudes. The periodogram of the AESTRA-corrected RVs
shows a more prominent peak, allowing for a clearer detection of 0.10
$\rm{m\,s}^{-1}$ planetary signal.

To further evaluate the model’s performance in recovering planetary signals,
we performed Markov Chain Monte Carlo (MCMC) sampling using the emcee package
(Foreman-Mackey et al., 2013) to retrieve the orbital parameters from the
corrected RV $v_{\rm correct}$. Currently, we set the uncertainty for model
RVs to be the same as the RV scatter $(\Delta v_{\rm correct})_{\rm std}$, as
obtaining principled uncertainties directly from AESTRA remains a task for the
future. We intend to explore alternative approaches, such as noise propagation
through the network or least-squares fit error, in the future development of
AESTRA.

Table 3: RV estimates and parameter constraints through MCMC sampling | Case I | Case II
---|---|---
RV [$\rm{m\,s}^{-1}$ ] | $\mathbf{N=1000}$ | $\mathbf{N=100}$ | $\mathbf{N=200}$
$(\Delta v_{\rm apparent})_{\rm std}$ | 3.38 | 3.27 | 3.51
$(\Delta v_{\rm traditional})_{\rm std}$ | 0.98 | 0.96 | 0.98
$(\Delta v_{\rm correct})_{\rm std}$ | 0.46 | 0.73 | 0.44
Reference RV |  |  |
$(\Delta v_{\rm ref})_{\rm std}$ | 0.39 | 0.38 | 0.41
Parameters |  |  |
Period [day] | ${100.0}^{+0.27}_{-0.26}$ | ${115.1}^{+7.6}_{-7.8}$ | ${101.9}^{+1.3}_{-1.3}$
$K$ [$\rm{m\,s}^{-1}$ ] | ${0.12}^{+0.01}_{-0.01}$ | ${0.36}^{+0.08}_{-0.07}$ | ${0.29}^{+0.03}_{-0.03}$
Phase [∘] | ${5.0}^{+12.0}_{-12.0}$ | ${70.0}^{+25.0}_{-27.0}$ | ${37.0}^{+13.0}_{-13.0}$

Assuming circular orbits for simplicity, we assigned uniform priors for the
orbital period between 10 and 200 days, semi-amplitude between 0.0 and 10.0
$\rm{m\,s}^{-1}$ , and phase between $-180$ and 180 degrees. We initialized 32
walkers near the maximum a posteriori (MAP) solution and ran the MCMC sampler
for 10,000 steps with a 2,000-step burn-in. As shown in Figure 9, the RVs
measured by AESTRA constrained the orbital period to be ${100.0}\pm{0.27}$
days and the signal amplitude to be ${0.12}\pm 0.01$ $\rm{m\,s}^{-1}$ . The
full parameter constraints are summarized in Table 3.

To assess AESTRA’s performance in recovering planetary signals with different
amplitudes, we created new datasets using the same activity perturbations but
varying injected signal amplitudes ($K$ values). The model, trained on the
original dataset with 0.1 $\rm{m\,s}^{-1}$ planetary signals, was applied to
recover the signal amplitudes in the new datasets. Figure 10 shows that the
model accurately recovered weaker signals and slightly underestimated
amplitudes $>$ 1 $\rm{m\,s}^{-1}$ . This probably occurs because planetary RVs
are treated as noise during de-trending, causing stronger signals to introduce
more uncertainty in activity-dependent RVs.

Deep learning methods generally require large datasets, making it important to
understand how AESTRA performance scales with sample size. We generated a
smaller dataset of 100 spectra affected by random stellar activity and
injected a planetary signal of 0.30 m/s, comparable to the RV scatter caused
by pure photon noise. The apparent RV scatter in the $N_{\rm spec}=100$
dataset was 3.27 m/s. We followed the same procedure to train the AESTRA
model. The training performance and resulting RV estimates are detailed in
Table 2 and Table 3.

For the 100-spectrum simulated dataset, the best-fit AESTRA model reduced the
RV scatter from 3.27 $\rm{m\,s}^{-1}$ to 0.73 $\rm{m\,s}^{-1}$ , lower than
the scatter of 0.96 $\rm{m\,s}^{-1}$ achieved by the method involving
traditional activity indicators. While the improvement was not as good as it
was in the 1,000-spectrum dataset, the AESTRA-corrected RVs still allowed for
a 5-$\sigma$ detection of the injected planetary signal, and gave a correct
estimate of the RV semi-amplitude, $K={0.36}^{+0.08}_{-0.07}$ $\rm{m\,s}^{-1}$
.

We believe that for small sample sizes, the primary source of RV uncertainty
in $v_{\rm correct}$ arises from the decorrelation process, which relies on
Gaussian smoothing in the latent space to estimate the activity-dependent RV
zero-points. If the latent space is underpopulated, i.e., if a given type of
stellar perturbation is not observed at different Doppler shifts, it becomes
challenging to calibrate the RV zero-points. This highlights the importance of
capturing potential stellar activity perturbations through repeated
observations.

Figure 10: Recovered versus true planetary signal amplitudes. The model
accurately recovers small amplitude planetary signals but tends to
underestimate amplitudes $>$1 $\rm{m\,s}^{-1}$ .

###  3.4 Case II: Time-Dependent Activity

Figure 11:  Top panel: Time-series of apparent RVs (black, left y-axis) and
the percentage of the stellar surface covered by active regions (orange, right
y-axis). Middle panel: Corrected model RVs (black) compared to injected
Doppler RVs (red). Bottom panel: Phase-folded, re-binned corrected model RVs
(black error bars), showing good agreement with the true RVs (red).  Figure
12:  Lomb-Scargle Periodograms illustrating the impact of complex spot
evolution on apparent RVs (top), RVs de-trended using traditional activity
indicators (middle), and the AESTRA-corrected RVs (bottom). The true planetary
period (red) is detected by both methods.

Next, we performed simulations in which stellar activity exhibited realistic
temporal correlations. We generated a dataset of 200 synthetic spectra
affected by evolving starspots using the SOAP 2.0 code, incorporating the
time-structured nature of starspot growth, decay, and rotation, as detailed
below. A planetary signal with amplitude 0.30 $\rm{m\,s}^{-1}$ was injected
into the dataset for recovery testing.

As before, we placed four spots on the stellar surface, with new spots
generated at random locations as old spots decayed. Each spot’s size evolved
over time, following a Gaussian function:

$\displaystyle S(t)=S_{\max}\exp\left[-\frac{1}{2}\left(\frac{t-t_{\rm
peak}}{\tau}\right)^{2}\right],$ (18)

where $S(t)$ represents the spot’s size at time $t$, $S_{\text{max}}$ is the
maximum spot size, chosen randomly between 0.5% and 20% stellar radius,
$t_{\text{peak}}$ is the time when the spot reaches its maximum size, and
$\tau$ is the spot’s lifetime, determined by $\tau\propto S_{\rm max}^{2}$,
with $\tau_{\rm max}\sim P_{\rm rot}$ for the largest spots. The spots’
longitudes and latitudes were randomly chosen based on the priors described in
Table 1. The growth and decay of starspots results in distinct time structures
in the apparent RVs, as shown in the top panel of Figure 11.

Training the AESTRA model on this dataset led to a reduction in the apparent
RV scatter from 3.51 $\rm{m\,s}^{-1}$ to 0.44 $\rm{m\,s}^{-1}$ , as presented
in Table 3. The AESTRA-based RV scatter of 0.44 $\rm{m\,s}^{-1}$ is about one-
half of the scatter obtained with the traditional method, and leads to a solid
detection of the 0.3 $\rm{m\,s}^{-1}$ planetary signals in the phase-folded
RVs shown in Figure 11.

Both the AESTRA model and the traditional analysis identified significant
peaks at the true planetary period in their Lomb-Scargle periodograms (Figure
12). However, the traditional analysis gave $K=0.52\pm 0.09$ $\rm{m\,s}^{-1}$
, an overestimate of the true RV semi-amplitude. The AESTRA model gave a
correct estimate, $K$ = ${0.29}\pm 0.03$ $\rm{m\,s}^{-1}$ .

As summarized in Table 3, when the AESTRA RVs were fitted with a circular-
orbit model using an MCMC method, the posteriors for the orbital parameters
overlapped with the true values, despite the presence of complex and time-
structured stellar activity. These results also underscore the importance of
good temporal coverage for robust RV estimation.

##  4 Conclusions and Outlook

AESTRA is a deep learning method designed to separate Doppler shifts from
spectral perturbations due to stellar activity. It operates on a collection of
spectra of the same star rather than a single spectrum or a cross-correlation
function. The autoencoder network summarizes spectral perturbations using a
small number of latent parameters, yielding a compact representation that can
be utilized as generalized activity indices. The RV estimator network is
jointly trained with the autoencoder to determine RV offsets. Training is
performed using a combined loss function that balances the various objectives,
resulting in a model that provides rest-frame spectral reconstruction and RV
estimation.

In our simulations, AESTRA successfully detected planetary signals with
amplitudes as low as 0.10 $\rm{m\,s}^{-1}$ despite activity-induced
perturbations with amplitudes of 3 $\rm{m\,s}^{-1}$ and photon noise of 0.30 m
s-1. Additionally, AESTRA showed encouraging performance even on datasets
consisting of as few as 100 spectra. The model recovered planetary signals in
a series of simulated spectra dominated by time-dependent stellar activity,
even without utilizing time-domain information.

While AESTRA has shown promising results on synthetic spectra, there is no
guarantee that it will be effective on real data. We plan to apply AESTRA to
analyze solar spectra obtained with instruments such as HARPS-N (Dumusque et
al., 2015), NEID (Lin et al., 2022), and EXPRES (Petersburg et al., 2020).
This validation step will assess the model’s real-world applicability and will
undoubtedly reveal limitations or biases that are not apparent in synthetic
data. This may provide valuable insights into solar variability across
different time scales and set the stage for applications to other stars, with
and without known planets.

AESTRA is at an early stage of development. We intend to explore some
potential extensions to improve the model’s performance. Currently, AESTRA
treats each spectrum independently, not considering the times of observation.
In our tests, AESTRA performed well in the presence of stellar activity
perturbations that were correlated in time (Case II), but we expect that it
would be even better if AESTRA took time-domain information into account. The
model could learn the temporal behavior of long-lived starspots, differential
rotation, and acoustic oscillations, which exhibit distinct timescales. Future
iterations of the AESTRA model could benefit from integrating a continuous
latent space time-series model and applying physically-motivated constraints
on the timescales associated with different types of activity.

More broadly, the current version of AESTRA is a general model that does not
incorporate any physical knowledge about stellar atmospheres or even any
built-in concept of a spectral absorption line. It seems likely that
integrating known physical processes and relationships that affect stellar
spectra into AESTRA will help it cope with more challenging and realistic
datasets. For example, future iterations could make use of models for telluric
absorption lines and for variations in the point-spread-function of the
spectrograph.

If AESTRA can be successfully validated using solar data, perhaps after
incorporating some of the improvements described above, the stage will be set
for applications to exoplanet data. We will explore AESTRA’s potential in
improving exoplanet characterization, analyzing low signal-to-noise
candidates, and searching for Earth-like planets. As we explore AESTRA’s
potential in our future work, the code is made publicly available on GitHub
(https://github.com/yanliang-astro/aestra), including a frozen version tagged
“paper-I” for reproducing the results of this work.

We thank Andrew Howard and David Hogg for constructive feedback and support.

## References

  * Ahumada et al. (2020) Ahumada, R., Prieto, C. A., Almeida, A., et al. 2020, The Astrophysical Journal Supplement Series, 249, 3
  * Beatty & Gaudi (2015) Beatty, T. G., & Gaudi, B. S. 2015, PASP, 127, 1240, doi: 10.1086/684264
  * Bedell et al. (2019) Bedell, M., Hogg, D. W., Foreman-Mackey, D., Montet, B. T., & Luger, R. 2019, The Astronomical Journal, 158, 164
  * Blackman et al. (2020) Blackman, R. T., Fischer, D. A., Jurgenson, C. A., et al. 2020, The Astronomical Journal, 159, 238
  * Bouchy et al. (2001) Bouchy, F., Pepe, F., & Queloz, D. 2001, Astronomy & Astrophysics, 374, 733
  * Brems et al. (2019) Brems, S. S., Kürster, M., Trifonov, T., Reffert, S., & Quirrenbach, A. 2019, Astronomy & Astrophysics, 632, A37
  * Collier Cameron et al. (2021) Collier Cameron, A., Ford, E., Shahaf, S., et al. 2021, Monthly Notices of the Royal Astronomical Society, 505, 1699
  * Crass et al. (2021) Crass, J., Gaudi, B. S., Leifer, S., et al. 2021, arXiv preprint arXiv:2107.14291
  * de Beurs et al. (2022) de Beurs, Z. L., Vanderburg, A., Shallue, C. J., et al. 2022, The Astronomical Journal, 164, 49
  * Dumusque et al. (2014) Dumusque, X., Boisse, I., & Santos, N. 2014, The Astrophysical Journal, 796, 132
  * Dumusque et al. (2015) Dumusque, X., Glenday, A., Phillips, D. F., et al. 2015, The Astrophysical Journal Letters, 814, L21
  * Foreman-Mackey et al. (2013) Foreman-Mackey, D., Hogg, D. W., Lang, D., & Goodman, J. 2013, Publications of the Astronomical Society of the Pacific, 125, 306
  * Giguere et al. (2016) Giguere, M. J., Fischer, D. A., Zhang, C. X., et al. 2016, The Astrophysical Journal, 824, 150
  * Hall et al. (2018) Hall, R. D., Thompson, S. J., Handley, W., & Queloz, D. 2018, Monthly Notices of the Royal Astronomical Society, 479, 2968
  * Haywood et al. (2014) Haywood, R., Collier Cameron, A., Queloz, D., et al. 2014, Monthly notices of the royal astronomical society, 443, 2517
  * He et al. (2015) He, K., Zhang, X., Ren, S., & Sun, J. 2015, in Proceedings of the IEEE international conference on computer vision, 1026–1034
  * Liang et al. (2023) Liang, Y., Melchior, P., Lu, S., Goulding, A., & Ward, C. 2023, AJ, 166, 75, doi: 10.3847/1538-3881/ace100
  * Lin et al. (2022) Lin, A. S., Monson, A., Mahadevan, S., et al. 2022, The Astronomical Journal, 163, 184
  * Mascareño et al. (2020) Mascareño, A. S., Faria, J., Figueira, P., et al. 2020, Astronomy & Astrophysics, 639, A77
  * Melchior et al. (2023) Melchior, P., Liang, Y., Hahn, C., & Goulding, A. 2023, The Astronomical Journal, 166, 74
  * Milbourne et al. (2019) Milbourne, T., Haywood, R., Phillips, D., et al. 2019, The Astrophysical Journal, 874, 107
  * Petersburg et al. (2020) Petersburg, R. R., Ong, J. M. J., Zhao, L. L., et al. 2020, AJ, 159, 187, doi: 10.3847/1538-3881/ab7e31
  * Robitaille et al. (2013) Robitaille, T. P., Tollerud, E. J., Greenfield, P., et al. 2013, Astronomy & Astrophysics, 558, A33
  * Smith & Geach (2023) Smith, M. J., & Geach, J. E. 2023, Royal Society Open Science, 10, 221454
  * Thompson et al. (2016) Thompson, S. J., Queloz, D., Baraffe, I., et al. 2016, in Ground-based and Airborne Instrumentation for Astronomy VI, Vol. 9908, SPIE, 1949–1961
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., et al. 2017, Advances in neural information processing systems, 30