OmniSpectra: A Unified Foundation Model for Universal Spectra Representation Learning

AI summary

Overview

Research area: Astronomical machine learning / foundation models for spectroscopy, with methods drawn from sequence modeling, implicit neural representations, and physics-informed representation learning (astro-ph.IM).

Technical level: Advanced. The paper assumes familiarity with Transformer attention, rotary position embeddings, masked autoencoding, and Fourier feature decoding, though the central idea can be grasped without that background.

Scope (one sentence): The paper introduces OmniSpectra, a 12.2M-parameter wavelength-native foundation model that replaces integer pixel indexing with continuous log-wavelength coordinates so that a single set of weights can ingest spectra from eight heterogeneous surveys without resampling.

What This Paper Is About

Spectroscopic surveys sample the same physical wavelength range at very different pixel densities, so a feature like the Hα line lands at a different pixel index on each instrument. Standard sequence models encode position by token rank, which means the same physical feature receives a different positional code per survey, and practitioners are forced to destructively resample spectra onto a common grid — blurring narrow absorption and emission lines and injecting interpolation artifacts.

OmniSpectra's goal is to remove that coordinate mismatch entirely: it anchors token position to physical log-wavelength rather than sequence index, so one model can ingest native spectra from heterogeneous instruments and transfer across them without interpolation or per-instrument adaptation.

Key Contributions

1. A unified wavelength-anchored positional encoding, realized with continuous 1D Rotary Position Encoding (RoPE) over log-wavelength. Because cosmological redshift is multiplicative in wavelength, it becomes an additive shift in log λ that cancels from RoPE's relative token distances, leaving attention geometry grounded in the physical ratio of wavelengths (log(λ_A/λ_B) = log(E_B/E_A)) rather than in detector sampling.

2. A resolution-gated implicit Fourier decoder that reconstructs flux at each instrument's native coordinates using a Fourier basis attenuated by a gate derived from the instrument's reported line-spread function (LSF) width, providing anti-aliasing against the spectral bias of implicit neural decoders.

3. An attention-pooled patch tokenizer that separates content (local line shape, extracted from fixed spans of P = 32 pixels) from position (a continuous code in log λ), so that token content carries no patch rank or instrument identifier.

4. A unified multi-instrument pretraining framework trained on 6.4M spectra from eight surveys (SDSS Legacy, BOSS, eBOSS, SEGUE-1, SEGUE-2, and DESI in the optical; APOGEE and GALAH in the near-infrared), spanning resolving power R ≈ 1,500–34,000 and 3,865–16,384 native pixels per spectrum.

Main Findings

• State of the art at small scale: OmniSpectra reaches state-of-the-art performance with only 12.2M parameters, outperforming substantially larger models — AstroCLIP at 370M and AION-1 at 314M parameters are named explicitly. The authors frame this as evidence that encoding physical structure can substitute for parameter scale.

• Stellar parameters: Best on all three targets (T_eff, log g, [Fe/H]) on both SDSS and APOGEE. SDSS: T_eff R² 0.93 (σ_MAD 158), log g 0.31/0.22, [Fe/H] 0.58/0.21. APOGEE: T_eff 0.96 (σ_MAD 85), log g 0.96 (0.18), [Fe/H] 0.91 (0.07). The gap to baselines persists as labels are reduced.

• Galaxy properties (DESI PROVABGS): Best on sSFR (R² 0.72, σ_MAD 1.00), t_age (0.55/0.57), Z_MW (0.69/0.41), and log M* (0.89/0.20), with substantial margins on three of the four targets. The authors note a tradeoff: GalSpecNet has a lower σ_MAD on log M* (0.18 versus 0.20), and several baselines have lower σ_MAD on t_age.

• Source classification (SDSS): Best F1 for every class — Galaxy 0.99, Quasar 0.93, Star 0.97. AION-1's public checkpoint collapses on this task (Quasar 0.05, Star 0.40). OmniSpectra retains significant performance with fewer than 500 labeled samples, though other models slowly close the gap as labels increase.

• Redshift estimation (DESI PROVABGS): R² 0.99, σ_MAD 0.0001, and a 0.06% catastrophic-outlier rate (|r| > 0.15). GaSNet-III ties on σ_MAD (0.0001) but at a 0.28% outlier rate — nearly 5× higher.

• Cross-survey generalization: Pretrained on SDSS alone (3800–9200 Å) and evaluated on DESI and APOGEE with no further pretraining, OmniSpectra transfers with little degradation, even though APOGEE (15100–17000 Å) lies entirely outside the pretraining wavelength range. Perceiver's index-based encoding degrades sharply on redshift under the same protocol.

• Inference efficiency: On identical DESI batches (batch size 64), OmniSpectra and Shen et al. process shorter sequences and reach the highest throughput, with OmniSpectra using far less memory than per-pixel and fixed-length baselines.

• Ablation (leave-one-component-out): Full model beats the all-alternatives base on every task — stellar R² 0.599 vs 0.544, APOGEE 0.940 vs 0.797, source F1 0.985 vs 0.968, galaxy 0.713 vs 0.668, redshift 0.991 vs 0.965. The positional code matters most on APOGEE (0.759 without it), consistent with transfer from optical to near-infrared; the LSF gate matters most on source classification. Some partial configurations underperform the base, indicating non-additive component interactions.

• Anti-aliasing evidence: At 8× native sampling, an ungated reconstruction of a SEGUE patch develops between-pixel oscillations absent from an independent higher-resolution GALAH observation of the same star; the LSF gate substantially reduces them.

Methodology in Plain English

The core idea. The authors argue that the positional encoding, not the tokenizer, is what breaks cross-instrument transfer. So they keep the tokenizer simple and make position physical.

Tokenizing. Each spectrum is chopped into patches of 32 consecutive detector pixels. Inside a patch, a small network looks at the local shape of the signal — line widths, asymmetries — and compresses it into a single content token via a learned attention-weighted average. Unlike mean pooling, learned weights can concentrate on a narrow absorption line instead of diluting it with the surrounding continuum. The token carries no information about which patch number it is or which instrument it came from.

Positioning. Each token is then given a position code computed as a continuous function of the patch's average log-wavelength, using sin/cos functions at log-spaced frequencies. Because the code is a formula rather than a lookup table, it can be evaluated at any wavelength, and the same physical region gets a comparable code on any grid. A rotary position embedding (RoPE) over log-wavelength supplies relative position, so attention compares tokens by the ratio of their wavelengths.

Why log-wavelength matters. Cosmological redshift stretches every wavelength by the same factor (1+z). On a pixel grid this changes how many tokens separate two lines, in a way that depends on the instrument's sampling. In log-wavelength, redshift simply adds a constant to every coordinate, so the constant cancels when attention computes relative distance. Identifying a feature therefore does not require already knowing the redshift — important, since redshift is itself a prediction target.

Reconstructing. For masked patches, the decoder queries the encoder's tokens by cross-attention and then evaluates a Fourier basis at each individual native pixel coordinate, gated by the instrument's reported spectral resolution so that frequencies finer than the instrument can actually resolve are suppressed. A final MLP combines the encoder context with this coordinate-dependent basis. The gate is an inductive bias, not a strict band-limit guarantee, because the MLP is nonlinear.

Training and evaluation. Pretraining is a masked autoencoder: each valid 32-pixel patch is masked with probability 0.5. Training runs 200 epochs on 4 NVIDIA A100 GPUs with AdamW (learning rate 10⁻⁴, weight decay 0.05, cosine annealing, batch size 64). Downstream, the encoder is frozen and only a lightweight cross-attention pooling head is trained for 10 epochs; supervised baselines are trained end-to-end per task. Results average over 5 random seeds. Shen et al. and Perceiver are re-pretrained by the authors on the identical corpus under a matched step budget, making those architecture-only comparisons.

Why This Matters

Research impact. The paper reframes cross-instrument misalignment as a structural problem of coordinate representation rather than a data or scale problem. If a 12.2M-parameter model with the right inductive biases can beat models 25–30× larger, that changes how the field thinks about scaling astronomical foundation models. The log-wavelength plus RoPE argument — that redshift becomes a translation that cancels from relative attention — is a reusable idea, and the LSF gate shows how an instrument's stated resolution can be turned directly into an architectural constraint rather than a hyperparameter to tune.

Real-world applications:

• Large spectroscopic surveys: Redshift estimation and source classification at scale for DESI, SDSS-V, and future petabyte-scale surveys, without per-instrument retraining.
• Cross-survey data fusion: Combining optical and near-infrared spectra from different facilities into a shared representation, useful for stellar population and galactic archaeology studies.
• Stellar parameter pipelines: Predicting T_eff, log g, and [Fe/H] from either optical (SDSS) or near-infrared (APOGEE) data with one model.
• Low-data regimes: With fewer than 500 labeled examples, the frozen-encoder approach still separates stars, galaxies, and quasars, which matters for rare-object searches and newly commissioned instruments that lack labeled training data.

Industry relevance. The underlying problem — continuous physical signals sampled on heterogeneous grids, where the same phenomenon receives a different code depending on the sensor — is not unique to astronomy. The paper itself cites the same mismatch in Earth observation (per-band physical wavelength conditioning), MRI (raw measurement coordinates), and mass spectrometry (continuous m/z per peak). Any domain with multiple sensors and inconsistent sampling could adopt the same recipe: anchor position to a physical coordinate, choose a transform that turns the dominant nuisance variation into a translation, and gate the decoder by the sensor's known resolution.

Future Directions

1. Wavelength coverage beyond the tuned range. The positional encoding's frequency periods are tuned to roughly 3,600–18,000 Å. Although the encoding is mathematically evaluable at any wavelength, performance at substantially different wavelengths — ultraviolet, far-infrared, radio — has not been evaluated and is the paper's own stated limitation.

2. Multimodal extension. The representations are currently unimodal. The authors name CLIP-style alignment to other modalities (photometry, images, maybe time series) as future work, which would test whether the wavelength-anchored token space can serve as a common substrate across data types.

3. Component interactions in the ablation. Several partial configurations underperform the base model, meaning the components interact non-additively and contributions are task-dependent. Understanding why — particularly why the positional code dominates APOGEE transfer and the LSF gate dominates classification — is an open question the ablation poses but does not resolve.

4. Pushing on parameter efficiency. The result that 12.2M parameters can beat 314M and 370M models invites a systematic study of how far physics-grounded inductive biases can substitute for scale, and where that substitution breaks down.

Target Audience

Astronomers and astrophysicists working with multi-survey spectroscopy who are tired of per-instrument pipelines and resampling. Machine learning researchers interested in geometric deep learning, coordinate-based representations, and physics-informed inductive biases for continuous signals. Practitioners in adjacent scientific domains — remote sensing, medical imaging, mass spectrometry — facing the same heterogeneous-sampling problem. Readers evaluating whether foundation-model scale is the right lever for scientific data, since this paper is a direct argument that structure can beat size. Prerequisites: comfort with Transformer attention and positional encoding will make the methodology sections substantially easier, though the intuition is accessible without them.

Authors’ abstract

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.