OCCAMS: A Multi-Task Learning-Based Investigation of the CMB Power Spectra
Authors: Ankit Mandal, Rahul Shah, Md. Wali Hossain, Somasri Sen, Anjan A Sen, Supratik Pal
Organizations: Physics and Applied Mathematics Unit, Indian Statistical Institute, 203 B.T. Road, Kolkata 700 108, India · Department of Physics, Jamia Millia Islamia, New Delhi 110025, India · Centre for Theoretical Physics, Jamia Millia Islamia, New Delhi 110025, India
How many latent variables are needed to fit the cosmic microwave background (CMB) temperature and polarisation power spectra when a neural network is given no cosmological model, and to what extent do they map onto the standard cosmological parameters? We use multi-task learning with variational autoencoders trained on CLASS-generated spectra (TT, TE, EE) over 2≤ℓ≤2500. No cosmological model is built into the network, and the number of latent variables, L, is varied over L=5, 6 and 7. Their fits are compared to the Planck 2018 data using χ^2, AIC, BIC and Bayesian evidence. Their correspondence with cosmological parameters is measured with normalised mutual information post-training. The loss weights are set within a single training run by uncertainty weighting, with the weight of the disentanglement term either fixed (semi-automated, SA) or learned (completely automated, CA). The reconstruction loss is split over multipole ranges in three configurations that separate the roles of temperature and polarisation and of the low- and high-ℓ ranges. In four of the six combinations of configuration and strategy, the fit improves from L=5 to L=6 and not beyond, and L=6 is the best, or close to the best, in five of the six cases. So, six latent variables are needed, as many as the free parameters of the standard, model-dependent approach. SA gives clearer disentanglement, while CA needs no a priori choice of the disentanglement weight and has the higher evidence at L=5 and 6: the two strategies are points on a trade-off between interpretability and fit to the data. Notably, the latent variables for τ and As share almost no information in any model, consistent with large-scale polarisation separating the two, as in standard analyses. A deep-learning approach can thus give a compact, data-driven parametrisation of the CMB spectra.
Figures & tables
Figure 1: The VAE network architecture. The input spectra are compressed via successive 1D convolutions into a latent distribution defined by μ and logσ2 . A stochastic sample z is drawn and subsequently upsampled via transposed convolutions to reconstruct the original spectra.
Parameter
Sampling Range
ωb
[0.020,0.024]
ωcdm
[0.10,0.13]
h
[0.62,0.80]
τ
[0.01,0.13]
ns
[0.92,1.01]
ln(1010As)
[2.90,3.18]
Table 1: Uniform sampling ranges used to generate the training data set via a Latin hypercube. Following Piras et al. (2025) , the ranges span approximately 10σ around the Planck 2018 baseline best-fit values. There are two exceptions to these bounds: the lower limit for the optical depth ( τ ) is restricted to 0.01 to prevent non-physical negative values, and the upper limit for the reduced Hubble constant ( h ) is broadened to encompass the local distance-ladder measurement from the SH0ES collaboration ( Riess et al., 2022 ) .
Configuration Name
Reconstruction Terms T
lowTTTEEE+full
T={Llowℓ,Lfullℓ}
lowTTTEEE+high
T={Llowℓ,Lhighℓ}
lowTT-TEEE+high
T={LlowTTℓ,LlowTEEEℓ,Lhighℓ}
Table 2: Summary of the three loss function configurations considered here.
Figure 2: The end-to-end Multi-Task Learning VAE pipeline for CMB parametrisation, shown for the CA strategy. Top: The theoretical Λ CDM model is simulated via CLASS to generate target temperature and polarisation spectra. Middle (violet dashed box): These spectra are compressed by the encoder into an L -dimensional probabilistic latent space, before being reconstructed by the decoder. Bottom (red dashed box): The extraction of data to compute the individual loss terms during training. A trainable parameter σi is attached to each term. Backpropagation (solid red curves) updates the network parameters ( ϕ,θ ) and the σi together, so that the loss weights are set in a single training run. In SA strategy, wd is fixed and the disentanglement loss term is not assigned any σi .
Figure 3: Absolute reconstruction residuals ∣ΔDℓ∣=∣DℓVAE−Dℓtrue∣ of the Planck 2018 best-fit Λ CDM TT (top), TE (middle) and EE (bottom) spectra. Each pair of columns shows one configuration, with SA (left) and CA (right). The latent dimensions are L=5 (blue), L=6 (orange) and L=7 (green).
Figure 4: Normalised MI between each Λ CDM parameter (rows) and each latent variable (columns) for all 18 trained models. The three blocks are the configurations, each shown for the SA and CA strategies, and the columns of panels correspond to L=5,6,7 . The number above each panel is the disentanglement score S . Solid outline: parameter-latent pair assigned by the matching, with normalised MI of at least 0.10; dashed outline: assigned pair below the threshold of 0.10, counted as not captured. All values are on a common colour scale that follows a power law with exponent 0.6 so that small values can be told apart.
Configuration
Strategy
L
χ^2
Δ AIC
Δ BIC
lnB
Score
lowTTTEEE+full
SA
5
0.9630
1.1
0.0
0.3
0.619
6
0.9607
0.7
4.0
0.0
0.737
7
0.9578
0.0
7.8
0.4
0.781
CA
5
0.9575
0.0
0.0
0.0
0.388
6
0.9603
2.8
7.2
3.5
0.144
7
0.9621
4.9
13.7
5.5
0.399
Table 3: Goodness of fit and model comparison for all 18 models, using the MCMC chains and the Plik_lite likelihood. Δ AIC, Δ BIC and lnB are measured from the best L within each configuration and strategy, and the lowest values are highlighted in bold. Score is the disentanglement score Stotal of equation ( 7 ), computed from the normalised MI matrices of Fig. 4 .
Configuration
L=5
L=6
L=7
lowTTTEEE+full
+4.3
+0.6
−1.0
lowTTTEEE+high
+5.6
+2.5
−0.5
lowTT-TEEE+high
+10.2
+3.4
−0.8
Table 4: Evidence comparison of the two disentanglement strategies at fixed latent dimension: lnBCA,SA=lnZCA−lnZSA for the same configuration. Positive values favour CA and negative values favour SA.
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Visualisation of the disentanglement penalty calculation during a training epoch. The network multiplies the transposed standardised parameter batch ( Θ~T ) by the standardised latent batch ( Z~ ). Each element Ri,j is the squared correlation between parameter i (e.g. τ ) and latent j . Normalising each column of R to sum to one gives M , from which the three terms of Ldis are computed: purity and strict from the columns of M , and collision from the off-diagonal elements of MTM . Minimising these terms makes each latent follow only a few parameters, and different latents follow different ones.
Configuration
Strategy
L
wl(104)
wh(102)
β
wd
lowTTTEEE+full
SA
5
0.168
0.283
0.253
5
6
0.994
1.16
0.209
5
7
0.973
1.21
0.238
5
CA
5
0.346
0.225
0.238
0.810
6
1.28
1.21
0.191
0.587
7
1.69
1.27
0.192
0.446
Appendix
Table 5: Loss weights at the end of training for all 18 models. wl and wh are the weights of the low- ℓ and high- ℓ reconstruction terms, β is that of the KL term, and wd is that of the disentanglement term. For lowTT-TEEE+high the wl column lists (wlowTT,wlowTEEE) , and for lowTTTEEE+full the wh column holds wfull .
Figure 6: Sensitivity of the disentanglement score to the MI threshold for the three configurations (solid: SA, dashed: CA; colours: L ). The dotted line marks the nominal threshold of 0.10 used in the main analysis.
Figure 7: Posterior distributions of the latent variables for the lowTTTEEE+high models with L=5 : SA (top) and CA (bottom). Contours enclose 68 and 95 per cent of the posterior.
Figure 8: Posterior distributions of the latent variables for the lowTTTEEE+high models with L=6 : SA (top) and CA (bottom). Contours enclose 68 and 95 per cent of the posterior.
Figure 9: Posterior distributions of the latent variables for the lowTTTEEE+high models with L=7 : SA (top) and CA (bottom). Contours enclose 68 and 95 per cent of the posterior.
Field-level inference of cosmological initial conditions from galaxy surveys requires a forward model that is simultaneously accurate in the non-linear regime, computationally efficient, and fully differentiable. Traditional N-body simulations are accurate but computationally prohibitive for iterative inference, while approximate solvers like Lagrangian Perturbation Theory (LPT) fail to capture the knotty halo-forming dynamics of the cosmic web at late times. We introduce the Lagrangian Neural Cellular Automaton (LNCA), a hybrid deep learning framework that can be applied to emulate structure formation as a local, iterative dynamical process on a comoving lattice. Unlike standard Eulerian Convolutional Neural Networks (CNNs) which map fixed density fields, the LNCA operates in the Lagrangian frame, advecting the computational graph itself to follow the flow of mass. By training the network to learn only the residual displacement corrections to the Zeldovich approximation, we achieve high-fidelity emulation of the non-linear physics while guaranteeing accuracy at large scales. We further constrain our model to produce complete trajectories, not just final states, by adopting an equivariant cellular automaton architecture, which recurrently iterates on its internal states to yield a dynamic history. The resulting model is strictly local, translationally and rotationally equivariant, and naturally supports continuous time integration, making it a reliable differentiable forward model for reconstructing the initial conditions of the universe from lightcone data. Our trained model supports percent-level precision in the power and cross spectra well into the non-linear regime (k≲0.5hMpc−1), while requiring ∼104 times fewer learned parameters than comparable models which take the form of an interpretable internal dynamic rule set.
Cooper Jacobus, Beatriz Tucci, Oliver Philcox
Department of Physics, Stanford University, Stanford, CA 94305, USA · Kavli Institute for Particle Astrophysics and Cosmology, Stanford · Leinweber Institute for Theoretical Physics, Stanford
Standard sequence models encode position by token rank rather than physical coordinate, typically forcing practitioners to destructively resample continuous physical signals such as spectra observed across heterogeneous instruments. We identify this coordinate mismatch as a structural limitation in cross-instrument representation learning, and propose OmniSpectra, a wavelength-native foundation model for spectral representation learning: a continuous-coordinate masked autoencoder that discards the integer pixel indexing of sequence modeling. Mapping physical coordinates to a logarithmic basis turns the multiplicative scaling induced by cosmological redshift into a uniform additive shift. In continuous rotary position embeddings, this shared shift cancels from relative token distances rather than producing grid-dependent positional distortion. We further counter the spectral bias inherent to implicit neural decoders with a physics-grounded, resolution-gated Fourier decoder that provides anti-aliasing. Pretrained on eight surveys and benchmarked against seven baselines, our model reaches state-of-the-art performance with only 12.2M parameters while outperforming substantially larger models, indicating that encoding physical structure can substitute for parameter scale.
Md Khairul Islam, Zeyu Xia, Adam S. Bolton +1
Computer Science Dept, University of Virginia · School of Data Science, University of Virginia
Recovering cosmological information beyond the power spectrum is a central goal for upcoming cosmological surveys, since late-time non-Gaussian signal in the matter density cannot be accessed through two-point statistics alone. Marked statistics fold part of this information back into the two-point level by reweighting the field with non-linear functions. We propose a neural marking scheme to generalize this process through a set of interpretable, physically motivated transformations that directly allow to interpret the gain in cosmological information at the morphological level. We employ a contrastive learning objective to align learnable marked summaries with the underlying cosmological parameters. At kmax=0.2hMpc−1, our neural mark tightens the marginalized constraint on σ8 by 2.9× and on Ωm by 1.8× compared to classical marks, breaking the Ωm−σ8 degeneracy at the Fisher information level. It further reduces the parameter MSE across our cosmological parameter prior by 1.45× over the best classical mark. The learned latent geometry aligns with the Ωm and σ8 directions in parameter space, indicating that the contrastive objective recovers the dominant axes of cosmological information. Our approach opens the door to more powerful, interpretable summary statistics for cosmological inference.
Federico Semenzato, Benjamin D. Wandelt, Michele Liguori +1
Dipartimento di Fisica Galileo Galilei, Universit`a di Padova, I-35131 Padova, Italy · INFN Sezione di Padova, I-35131 Padova, Italy · Department of Physics and Astronomy, Johns Hopkins University, 3400 North Charles Street, Baltimore, MD, 21218, USA +2