Reference / API reference / dorothy / dorothy.models / dorothy.models.multi_head_mlp

dorothy.models.multi_head_mlp¶

Multi-head MLP architecture for multi-survey stellar parameter inference.

This module implements a multi-head architecture where each survey has its own encoder network, feeding into a shared trunk that produces stellar parameter predictions. This allows handling surveys with different wavelength grids while learning shared representations.

Architecture:

Survey Encoders (one per survey):: Input: (batch, 3, n_wavelengths_i) -> Flatten -> Hidden layers -> (batch, latent_dim)
Shared Trunk:: Input: (batch, latent_dim) or (batch, n_surveys * latent_dim) for combined Hidden layers -> (batch, trunk_output_dim)
Output Head (a LabelAwareOutput per label source):: Label decoder: the trunk output -> Hidden layers -> mu Error decoder: the trunk output and the log label errors -> Hidden layers of the same sizes -> ln s Output: (batch, 2, n_params) = [mu, ln s]

With residual_stream (the default), the trunk and the decoders’ hidden layers are residual blocks of width latent_dim (see SharedTrunk and LabelAwareOutput).

For multi-survey training:

Each sample may have data from one or more surveys
Encoders only process surveys with valid data (indicated by has_data mask)
Encoded representations are combined (averaged or concatenated) before trunk

Classes

`MultiHeadMLP`(survey_configs[, n_parameters, ...])	Multi-head MLP for multi-survey stellar parameter prediction.
`OutputHead`(input_dim[, n_parameters, ...])	Output head that produces stellar parameter predictions.
`SharedTrunk`(input_dim[, output_dim, ...])	Shared trunk network that processes combined survey embeddings.
`SurveyEncoder`(survey_name, n_wavelengths[, ...])	Encoder network for a single survey's spectral data.

class dorothy.models.multi_head_mlp.SurveyEncoder(survey_name, n_wavelengths, latent_dim=256, hidden_layers=None, normalization='layernorm', activation='gelu', dropout=0.0)[source]¶

Bases: Module

Encoder network for a single survey’s spectral data.

Takes 3-channel spectral input (flux, rho, mask) and produces a fixed-size latent representation.

survey_name¶: Name of the survey this encoder handles.

n_wavelengths¶: Number of wavelength bins for this survey.

latent_dim¶: Output dimension of the encoder.

forward(x)[source]¶

Encode spectral data to latent representation.

Parameters:: x (Tensor) – Input tensor of shape (batch, 3, n_wavelengths).
Returns:: Latent tensor of shape (batch, latent_dim).
Return type:: Tensor

class dorothy.models.multi_head_mlp.SharedTrunk(input_dim, output_dim=128, hidden_layers=None, normalization='layernorm', activation='gelu', dropout=0.0, residual=False)[source]¶

Bases: Module

Shared trunk network that processes combined survey embeddings.

Takes latent representations from survey encoders and produces features for the output head. The plain trunk is one hidden block (Linear, normalization, activation) per hidden layer, then a block to output_dim. The residual trunk is a stream of width output_dim: a linear map from input_dim when the sizes differ, one pre-norm ResidualBlock per hidden layer, whose hidden size it is, then a normalization.

input_dim¶: Input dimension (latent_dim for single survey or n_surveys * latent_dim for concatenation).

output_dim¶: Output dimension for the output head.

forward(x)[source]¶

Process combined embeddings.

Parameters:: x (Tensor) – Combined latent tensor of shape (batch, input_dim).
Returns:: Features tensor of shape (batch, output_dim).
Return type:: Tensor

class dorothy.models.multi_head_mlp.OutputHead(input_dim, n_parameters=11, hidden_layers=None, normalization='layernorm', activation='gelu', dropout=0.0, label_aware_uncertainty=True, residual=False)[source]¶

Bases: Module

Output head that produces stellar parameter predictions.

Takes features from the shared trunk and produces predictions in the standard (batch, 2, n_params) format. By default the head is a LabelAwareOutput: its label decoder maps the features through hidden_layers to mu, and its error decoder maps the features and the label errors through layers of the same sizes to ln s.

n_parameters¶: Number of stellar parameters to predict.

forward(x, label_errors=None)[source]¶

Produce parameter predictions.

Parameters:

x (Tensor) – Features tensor of shape (batch, input_dim).
label_errors (Tensor | None) – Reported label errors of shape (batch, n_parameters), in normalized units, for the log-scatter (see LabelAwareOutput). None uses the learned missing vector for every entry.

Returns:

Predictions of shape (batch, 2, n_parameters) where – - output[:, 0, :] contains predicted means - output[:, 1, :] contains predicted log-scatter

Raises:

ValueError – If label_errors is given to a head built with label_aware_uncertainty=False, or has the wrong shape.

Return type:

Tensor

class dorothy.models.multi_head_mlp.MultiHeadMLP(survey_configs, n_parameters=11, latent_dim=256, encoder_hidden=None, trunk_hidden=None, output_hidden=None, combination_mode='concat', normalization='layernorm', activation='gelu', dropout=0.0, label_sources=None, head_dropout=None, label_aware_uncertainty=True, residual_stream=True)[source]¶

Bases: Module

Multi-head MLP for multi-survey stellar parameter prediction.

This architecture has separate encoder networks for each survey, a shared trunk that combines the encodings, and one output head per label source.

The model can handle: - Single survey input (standard case) - Multiple surveys with per-sample availability masks

survey_names¶: List of survey names in order.

n_parameters¶: Number of stellar parameters being predicted.

combination_mode¶: How to combine multi-survey embeddings (“mean” or “concat”).

encoders¶: ModuleDict of survey encoders.

trunk¶: Shared trunk network.

latent_dim¶: Output dimension of each survey encoder.

label_sources¶: Label source of each output head; [“default”] when none were given.

output_heads¶: ModuleDict mapping each label source to its output head.

output_head¶: The output head of the first label source.

label_aware_uncertainty¶: Whether the heads are LabelAwareOutputs (True) or end in an AdditiveOutput.

residual_stream¶: Whether the trunk and the heads’ decoders form a residual stream of width latent_dim (True) or are plain blocks.

encode_present_rows_only¶: With multi-survey input, run each encoder only on the rows that have that survey’s data (default True). The rows without data are discarded either way, so this changes results only by float rounding (a different GEMM size); False encodes every row.

Examples

>>> model = MultiHeadMLP(
...     survey_configs={"boss": 4506, "lamost": 3700},
...     n_parameters=11,
... )
>>> # Single survey forward
>>> x_boss = torch.randn(32, 3, 4506)
>>> output = model.forward_single(x_boss, "boss")
>>>
>>> # Multi-survey forward
>>> inputs = {"boss": x_boss, "lamost": torch.randn(32, 3, 3700)}
>>> has_data = {"boss": torch.ones(32, dtype=torch.bool),
...             "lamost": torch.ones(32, dtype=torch.bool)}
>>> output = model.forward_multi(inputs, has_data)

classmethod from_config(config)[source]¶

Create a MultiHeadMLP from a MultiHeadModelConfig.

Parameters:: config (MultiHeadModelConfig) – Multi-head model configuration.
Returns:: Configured MultiHeadMLP instance.
Return type:: MultiHeadMLP

forward_single(x, survey)[source]¶

Forward pass for single-survey input.

This is the standard case where all samples come from one survey.

Parameters:

x (Tensor) – Input tensor of shape (batch, 3, n_wavelengths).
survey (str) – Name of the survey.

Returns:

Predictions of shape (batch, 2, n_parameters).

Raises:

KeyError – If survey is not in the model’s encoders.

Return type:

Tensor

forward_multi(inputs, has_data)[source]¶

Forward pass for multi-survey input.

Handles samples that may have data from different subsets of surveys.

Parameters:

inputs (dict[str, Tensor]) – Dict mapping survey names to input tensors. Each tensor has shape (batch, 3, n_wavelengths_survey).
has_data (dict[str, Tensor]) – Dict mapping survey names to boolean masks. Each mask has shape (batch,) indicating which samples have valid data from that survey.

Returns:

Predictions of shape (batch, 2, n_parameters) from the first head.

Raises:

ValueError – If none of the keys of inputs is a survey of the model.

Return type:

Tensor

Note

A sample without data from any survey gets an output that carries no information.

forward(x, survey=None, has_data=None, label_errors=None)[source]¶

Unified forward pass.

Dispatches to forward_single or forward_multi based on input type.

Parameters:

x (Tensor | dict[str, Tensor]) – Either a single tensor (with survey name) or dict of tensors.
survey (str | None) – Survey name if x is a single tensor.
has_data (dict[str, Tensor] | None) – Data availability masks if x is a dict.
label_errors (Tensor | None) – Reported label errors (batch, n_parameters) for the first head’s log-scatter (see LabelAwareOutput). None uses the learned missing vector.

Returns:

Predictions of shape (batch, 2, n_parameters).

Raises:

KeyError – If survey is not in the model’s encoders.
ValueError – If x is a single tensor without survey on a multi-survey model, or a dict without any of the model’s surveys; or if label_errors is given to a model built with label_aware_uncertainty=False, or has the wrong shape.

Return type:

Tensor

forward_for_label_source(x, label_source, survey=None, has_data=None)[source]¶

Forward pass using a specific label source’s output head.

Use this for multi-label training where different samples may have labels from different sources.

Parameters:

x (Tensor | dict[str, Tensor]) – Input tensor or dict of tensors (same as forward()).
label_source (str) – Which label source’s output head to use.
survey (str | None) – Survey name if x is single tensor.
has_data (dict[str, Tensor] | None) – Data availability masks if x is dict.

Returns:

Predictions of shape (batch, 2, n_parameters).

Raises:

KeyError – If label_source is not in the model’s output heads.

Return type:

Tensor

forward_all_label_sources(x, survey=None, has_data=None, label_errors=None)[source]¶

Forward pass through all label source output heads.

Useful for inference when you want predictions from all label sources.

Parameters:

x (Tensor | dict[str, Tensor]) – Input tensor or dict of tensors.
survey (str | None) – Survey name if x is single tensor.
has_data (dict[str, Tensor] | None) – Data availability masks if x is dict.
label_errors (dict[str, Tensor] | None) – Label source name -> reported label errors (batch, n_parameters) for that head’s log-scatter (see LabelAwareOutput). Heads without an entry, and all heads when None, use their learned missing vector.

Returns:

Dict mapping label source names to predictions. Each prediction has shape (batch, 2, n_parameters).

Raises:

KeyError – If label_errors names a label source the model has no head for, or survey is not in the model’s encoders.
ValueError – If x is a single tensor without survey on a multi-survey model, or a dict without any of the model’s surveys; or if label_errors is given to a model built with label_aware_uncertainty=False, or has the wrong shape.

Return type:

dict[str, Tensor]

get_embeddings(x, survey, level='encoder')[source]¶

Get latent embeddings for a survey’s input.

Useful for anomaly detection and analysis.

Parameters:

x (Tensor) – Input tensor of shape (batch, 3, n_wavelengths).
survey (str) – Name of the survey.
level (str) – “encoder” for the survey encoder’s latent vector, or “trunk” for the shared trunk features (comparable across surveys, and the input to every output head).

Returns:

Embeddings of shape (batch, latent_dim) for “encoder” or (batch, trunk output dim) for “trunk”: 128, or latent_dim with residual_stream.

Raises:

KeyError – If survey is not in the model’s encoders.
ValueError – If level is not “encoder” or “trunk”.

Return type:

Tensor

count_parameters()[source]¶

Count the trainable parameters.

Returns:: Total number of trainable scalar parameters across all encoders, shared trunk, and output heads.
Return type:: int

count_parameters_by_component()[source]¶

Count the trainable parameters by component.

Returns:: Dict with the keys “encoder_<survey>”, “trunk”, “output_head_<source>” (one per head with several label sources, else “output_head”) and “total”, each the number of trainable parameters.
Return type:: dict[str, int]

extra_repr()[source]¶

Return string representation of model configuration.