MMSplice¶

Modular modeling of the effects of genetic variants on splicing.

Disclaimer¶

This is an UNOFFICIAL implementation of the MMSplice: modular modeling improves the predictions of genetic variant effects on splicing by Jun Cheng, et al.

The OFFICIAL repository of MMSplice is at gagneurlab/MMSplice_MTSplice.

Tip

The MultiMolecule team has confirmed that the provided model and checkpoints are producing the same intermediate representations as the original implementation.

The team releasing MMSplice did not write this model card for this model so this model card has been written by the MultiMolecule team.

Model Details¶

MMSplice is a modular neural network for predicting the effect of genetic variants on pre-mRNA splicing. It decomposes an exon together with its flanking introns into five regions and scores each region with an independent small convolutional sub-network. For variant-effect estimation, the model is run on both the reference and the alternative sequence, and the per-module score deltas are combined by a fixed linear model into a delta-logit-PSI splicing-effect score. Please refer to the Training Details section for more information on the training process.

Model Specification¶

Num Modules	Num Parameters (M)	FLOPs (M)	MACs (M)
5	0.057	5.71	2.79

(FLOPs and MACs measured on a 220 bp exon-with-flanks input.)

Links¶

Code: multimolecule.mmsplice
Data: Human splice-site and exon data with MPRA exon-skipping variant-effect measurements
Paper: MMSplice: modular modeling improves the predictions of genetic variant effects on splicing
Developed by: Jun Cheng, Thi Yen Duong Nguyen, Kamil J. Cygan, Muhammed Hasan Çelik, William G. Fairbrother, Žiga Avsec, Julien Gagneur
Model type: Modular 1D CNN with five region-specific sub-networks for splice variant-effect prediction
Original Repository: gagneurlab/MMSplice_MTSplice

Usage¶

The model file depends on the multimolecule library. You can install it using pip:

Bash
1	`pip install multimolecule`

Direct Use¶

Module Scores¶

Python
>>> import torch
>>> from multimolecule import RnaTokenizer, MmSpliceForSequencePrediction

>>> tokenizer = RnaTokenizer.from_pretrained("multimolecule/mmsplice")
>>> model = MmSpliceForSequencePrediction.from_pretrained("multimolecule/mmsplice")
>>> _ = model.eval()
>>> left_intron = "A" * 100
>>> exon = "C" * 20
>>> right_intron = "G" * 100
>>> reference = tokenizer(left_intron + exon + right_intron, add_special_tokens=False, return_tensors="pt")
>>> output = model.model(**reference)
>>> output["logits"].shape
torch.Size([1, 5])

Variant Effect¶

Python
>>> import torch
>>> from multimolecule import RnaTokenizer, MmSpliceForSequencePrediction

>>> tokenizer = RnaTokenizer.from_pretrained("multimolecule/mmsplice")
>>> model = MmSpliceForSequencePrediction.from_pretrained("multimolecule/mmsplice")
>>> _ = model.eval()
>>> left_intron = "A" * 100
>>> exon = "C" * 20
>>> right_intron = "G" * 100
>>> reference = tokenizer(left_intron + exon + right_intron, add_special_tokens=False, return_tensors="pt")
>>> alternative_exon = exon[:10] + "U" + exon[11:]
>>> alternative = tokenizer(left_intron + alternative_exon + right_intron, add_special_tokens=False, return_tensors="pt")
>>> output = model(
...     reference["input_ids"],
...     alternative_input_ids=alternative["input_ids"],
... )
>>> output["logits"].shape
torch.Size([1, 1])

Interface¶

Input length: exon sequence with 100 nt upstream intronic context + 100 nt downstream intronic context
Tokenization: disable special tokens; the embedding layer maps A/C/G/U ids to the four upstream channels and maps N, padding, special, and unknown tokens to all-zero columns
Output (reference-only call, input_ids / inputs_embeds): per-module score vector logits of shape (batch_size, 5)

Variant Effect¶

Reference + alternative call (also pass alternative_input_ids / alternative_inputs_embeds): additionally returns alternative_logits and per-module delta_logits = alternative_logits - logits
MmSpliceForSequencePrediction: requires both reference and alternative; returns the combined scalar delta-logit-PSI score of shape (batch_size, 1)

Training Details¶

MMSplice was trained as five independent modules on splicing data and the modules were combined with a linear model to predict variant effects on percent-spliced-in (PSI).

Training Data¶

The acceptor, donor, exon, and intron modules were trained on splice-site and exon data derived from human reference transcripts. The combining linear model was fit against a massively parallel reporter assay (MPRA) of exon-skipping variants.

Training Procedure¶

Pre-training¶

Each module was trained with a sequence-to-scalar objective scoring its region. The module scores (and their reference/alternative deltas) were then combined by a fixed linear model into a delta-logit-PSI splicing-effect score.

Citation¶

BibTeX
@article{cheng2019mmsplice,
  title     = {MMSplice: modular modeling improves the predictions of genetic variant effects on splicing},
  author    = {Cheng, Jun and Nguyen, Thi Yen Duong and Cygan, Kamil J and {\c{C}}elik, Muhammed Hasan and Fairbrother, William G and Avsec, {\v{Z}}iga and Gagneur, Julien},
  journal   = {Genome Biology},
  volume    = 20,
  number    = 1,
  pages     = {48},
  year      = 2019,
  publisher = {Springer},
  doi       = {10.1186/s13059-019-1653-z}
}

Note

The artifacts distributed in this repository are part of the MultiMolecule project. If MultiMolecule supports your research, please cite the MultiMolecule project as follows:

BibTeX
@software{chen_2024_12638419,
  author    = {Chen, Zhiyuan and Zhu, Sophia Y.},
  title     = {MultiMolecule},
  doi       = {10.5281/zenodo.12638419},
  publisher = {Zenodo},
  url       = {https://doi.org/10.5281/zenodo.12638419},
  year      = 2024,
  month     = may,
  day       = 4
}

Contact¶

Please use GitHub issues of MultiMolecule for any questions or comments on the model card.

Please contact the authors of the MMSplice paper for questions or comments on the paper/model.

License¶

This model implementation is licensed under the GNU Affero General Public License.

For additional terms and clarifications, please refer to our License FAQ.

Text Only
1	`SPDX-License-Identifier: AGPL-3.0-or-later`

API Reference¶

MmSpliceConfig ¶

Bases: PreTrainedConfig

This is the configuration class to store the configuration of a MmSpliceModel. It is used to instantiate a MMSplice model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a configuration to that of the MMSplice gagneurlab/MMSplice_MTSplice architecture.

Configuration objects inherit from PreTrainedConfig and can be used to control the model outputs. Read the documentation from PreTrainedConfig for more information.

MMSplice (Cheng et al. 2019, Genome Biology) is a modular model. Five independent sub-networks (acceptor_intron, acceptor, exon, donor, donor_intron) each score one region of the exon-with-flanking- introns sequence. The five scalar scores form the module score vector. For variant-effect estimation the model is run on the reference and the alternative sequence and the per-module score deltas are combined by the fixed upstream linear model into a delta-logit-PSI splicing-effect score.

The default module configurations replicate the upstream pretrained weights exactly (see [MmSpliceModuleConfig]).

Parameters:

Name	Type	Description	Default
`vocab_size` ¶	`int`	Vocabulary size of the MMSplice model. Defines the number of feature channels derived from the one-hot encoded `input_ids`. MMSplice uses four `A/C/G/U` channels; `N`, padding, special, and unknown tokens are encoded as all-zero columns by the embedding layer. Defaults to 4 (the `ACGU` nucleobase alphabet).	`4`
`modules` ¶	`dict \| None`	Per sub-module architecture configuration. A mapping from module name to a [`MmSpliceModuleConfig`]. The default defines the five canonical MMSplice modules with their upstream architectures.	`None`
`modules_config` ¶	`dict \| None`	Alias used when loading serialized configs. Prefer `modules` when constructing configs directly.	`None`
`acceptor_intron_cut` ¶	`int`	Number of bp removed from the 3’ end of the acceptor intron (the part considered the acceptor site).	`6`
`donor_intron_cut` ¶	`int`	Number of bp removed from the 5’ end of the donor intron (the part considered the donor site).	`6`
`acceptor_intron_length` ¶	`int`	Intron length consumed by the acceptor splice-site module.	`50`
`acceptor_exon_length` ¶	`int`	Exon flank length consumed by the acceptor splice-site module.	`3`
`donor_exon_length` ¶	`int`	Exon flank length consumed by the donor splice-site module.	`5`
`donor_intron_length` ¶	`int`	Intron length consumed by the donor splice-site module.	`13`
`num_labels` ¶	`int`	Number of sequence-prediction labels. MMSplice emits one scalar delta-logit-PSI score, so this must be 1.	`1`
`head` ¶	`HeadConfig \| None`	Loss configuration for [`MmSpliceForSequencePrediction`]. The upstream variant-effect combiner emits one scalar delta-logit-PSI score, so the default head config has `num_labels=1`.	`None`

Examples:

Python Console Session
>>> from multimolecule import MmSpliceConfig, MmSpliceModel
>>> # Initializing a MMSplice multimolecule/mmsplice style configuration
>>> configuration = MmSpliceConfig()
>>> # Initializing a model (with random weights) from the multimolecule/mmsplice style configuration
>>> model = MmSpliceModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config

Source code in multimolecule/models/mmsplice/configuration_mmsplice.py

Python
class MmSpliceConfig(PreTrainedConfig):
    r"""
    This is the configuration class to store the configuration of a
    [`MmSpliceModel`][multimolecule.models.MmSpliceModel]. It is used to instantiate a MMSplice model according to the
    specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a
    configuration to that of the MMSplice
    [gagneurlab/MMSplice_MTSplice](https://github.com/gagneurlab/MMSplice_MTSplice) architecture.

    Configuration objects inherit from [`PreTrainedConfig`][multimolecule.models.PreTrainedConfig] and can be used to
    control the model outputs. Read the documentation from [`PreTrainedConfig`][multimolecule.models.PreTrainedConfig]
    for more information.

    MMSplice (Cheng et al. 2019, *Genome Biology*) is a *modular* model. Five
    independent sub-networks (``acceptor_intron``, ``acceptor``, ``exon``,
    ``donor``, ``donor_intron``) each score one region of the exon-with-flanking-
    introns sequence. The five scalar scores form the module score vector. For
    variant-effect estimation the model is run on the reference and the
    alternative sequence and the per-module score deltas are combined by the
    fixed upstream linear model into a delta-logit-PSI splicing-effect score.

    The default module configurations replicate the upstream pretrained weights
    exactly (see [`MmSpliceModuleConfig`]).

    Args:
        vocab_size:
            Vocabulary size of the MMSplice model. Defines the number of feature
            channels derived from the one-hot encoded `input_ids`. MMSplice uses
            four `A/C/G/U` channels; `N`, padding, special, and unknown tokens are
            encoded as all-zero columns by the embedding layer.
            Defaults to 4 (the `ACGU` nucleobase alphabet).
        modules:
            Per sub-module architecture configuration. A mapping from module name
            to a [`MmSpliceModuleConfig`]. The default defines the five canonical
            MMSplice modules with their upstream architectures.
        modules_config:
            Alias used when loading serialized configs. Prefer `modules` when
            constructing configs directly.
        acceptor_intron_cut:
            Number of bp removed from the 3' end of the acceptor intron (the part
            considered the acceptor site).
        donor_intron_cut:
            Number of bp removed from the 5' end of the donor intron (the part
            considered the donor site).
        acceptor_intron_length:
            Intron length consumed by the acceptor splice-site module.
        acceptor_exon_length:
            Exon flank length consumed by the acceptor splice-site module.
        donor_exon_length:
            Exon flank length consumed by the donor splice-site module.
        donor_intron_length:
            Intron length consumed by the donor splice-site module.
        num_labels:
            Number of sequence-prediction labels. MMSplice emits one scalar
            delta-logit-PSI score, so this must be 1.
        head:
            Loss configuration for [`MmSpliceForSequencePrediction`]. The
            upstream variant-effect combiner emits one scalar delta-logit-PSI
            score, so the default head config has `num_labels=1`.

    Examples:
        >>> from multimolecule import MmSpliceConfig, MmSpliceModel
        >>> # Initializing a MMSplice multimolecule/mmsplice style configuration
        >>> configuration = MmSpliceConfig()
        >>> # Initializing a model (with random weights) from the multimolecule/mmsplice style configuration
        >>> model = MmSpliceModel(configuration)
        >>> # Accessing the model configuration
        >>> configuration = model.config
    """

    model_type = "mmsplice"

    def __init__(
        self,
        vocab_size: int = 4,
        modules: dict | None = None,
        modules_config: dict | None = None,
        acceptor_intron_cut: int = 6,
        donor_intron_cut: int = 6,
        acceptor_intron_length: int = 50,
        acceptor_exon_length: int = 3,
        donor_exon_length: int = 5,
        donor_intron_length: int = 13,
        num_labels: int = 1,
        head: HeadConfig | None = None,
        problem_type: str | None = "regression",
        bos_token_id: int | None = None,
        eos_token_id: int | None = None,
        pad_token_id: int = 4,
        **kwargs,
    ):
        if num_labels != 1:
            raise ValueError(f"MMSplice emits one delta-logit-PSI score; `num_labels` must be 1, got {num_labels}")
        super().__init__(num_labels=num_labels, pad_token_id=pad_token_id, **kwargs)
        self.vocab_size = vocab_size
        self.bos_token_id = bos_token_id  # type: ignore[assignment]
        self.eos_token_id = eos_token_id  # type: ignore[assignment]
        if modules is None:
            modules = modules_config
        if modules is None:
            acceptor_region = acceptor_intron_length + acceptor_exon_length
            donor_region = donor_exon_length + donor_intron_length
            modules = {
                "acceptor_intron": MmSpliceModuleConfig(
                    architecture="conv",
                    conv_channels=256,
                    conv_kernel_size=13,
                    conv_activation="relu",
                ),
                "acceptor": MmSpliceModuleConfig(
                    architecture="dense",
                    region_length=acceptor_region,
                    conv_channels=32,
                    conv_kernel_size=15,
                    conv_activation="relu",
                    conv_batch_norm=True,
                    pointwise_channels=32,
                    hidden_sizes=[],
                    flatten_dropout=True,
                ),
                "exon": MmSpliceModuleConfig(
                    architecture="conv",
                    conv_channels=128,
                    conv_kernel_size=11,
                    conv_activation="relu",
                    conv_batch_norm=True,
                    pool_mask_zeros=True,
                ),
                "donor": MmSpliceModuleConfig(
                    architecture="dense",
                    region_length=donor_region,
                    hidden_sizes=[128, 64],
                ),
                "donor_intron": MmSpliceModuleConfig(
                    architecture="conv",
                    conv_channels=256,
                    conv_kernel_size=13,
                    conv_activation="relu",
                ),
            }
        self.modules_config = {
            name: cfg if isinstance(cfg, MmSpliceModuleConfig) else MmSpliceModuleConfig(**cfg)
            for name, cfg in modules.items()
        }
        self.acceptor_intron_cut = acceptor_intron_cut
        self.donor_intron_cut = donor_intron_cut
        self.acceptor_intron_length = acceptor_intron_length
        self.acceptor_exon_length = acceptor_exon_length
        self.donor_exon_length = donor_exon_length
        self.donor_intron_length = donor_intron_length
        # The backbone hidden representation is the per-module score vector.
        self.hidden_size = len(MODULE_ORDER)
        self.problem_type = problem_type
        if head is None:
            head = HeadConfig(num_labels=num_labels, hidden_size=1, problem_type=problem_type)
        elif not isinstance(head, HeadConfig):
            head = HeadConfig(**head)
        self.head = head

        missing = sorted(set(MODULE_ORDER) - set(self.modules_config))
        if missing:
            raise ValueError(f"Missing required MMSplice modules in modules config: {missing}.")
        unexpected = sorted(set(self.modules_config) - set(MODULE_ORDER))
        if unexpected:
            raise ValueError(f"Unexpected MMSplice modules in modules config: {unexpected}.")
        if min(acceptor_intron_length, donor_intron_length) <= 0:
            raise ValueError("Intron region lengths must be positive.")
        if min(acceptor_exon_length, donor_exon_length) < 0:
            raise ValueError("Exon flank lengths must be non-negative.")

MmSpliceModuleConfig ¶

Bases: FlatDict

Configuration for a single MMSplice sub-module.

MMSplice is a modular model: each genomic region (the acceptor and donor splice sites, the exon body, and the two flanking intron stubs) is scored by an independent small network. The five upstream sub-networks (gagneurlab/MMSplice_MTSplice) do not share an architecture, so each is described by its own MmSpliceModuleConfig.

Two architecture families exist:

conv (acceptor_intron, exon, donor_intron): a single length-preserving 1D convolution, optional batch-norm, ReLU, global average pooling over positions, then a linear projection to a scalar. Length-independent.
dense (acceptor, donor): a fixed-length splice-site network. The region is one-hot encoded, optionally passed through one or more convolutions, flattened, and projected to a scalar with a stack of dense + batch-norm + ReLU blocks. A final sigmoid produces a probability that is converted to a logit score.

Parameters:

Name	Description	Default
`architecture` ¶	Either `conv` or `dense` (see above).	required
`region_length` ¶	Fixed input length the module consumes. `0` for the length- independent `conv` modules (`acceptor_intron` / `donor_intron`).	required
`conv_channels` ¶	Output channels of the (first) convolution.	required
`conv_kernel_size` ¶	Kernel size of the (first) convolution.	required
`conv_activation` ¶	Activation applied to the (first) convolution.	required
`conv_batch_norm` ¶	Whether a batch-norm follows the (first) convolution.	required
`pool_mask_zeros` ¶	Whether global average pooling ignores all-zero input positions. Upstream `exon` uses masked pooling to ignore `N` padding.	required
`pointwise_channels` ¶	Output channels of the `dense`-family `1x1` convolution. `0` disables it. Followed by a batch-norm.	required
`hidden_sizes` ¶	Output sizes of the dense blocks of a `dense`-family head.	required
`flatten_dropout` ¶	Whether a dropout is applied right after the flatten (before the dense blocks). Upstream `acceptor` uses this; `donor` instead applies dropout inside each dense block.	required
`dropout` ¶	Dropout probability used by `dense`-family heads.	required
`batch_norm_eps` ¶	Epsilon of every batch-norm (upstream Keras default `1e-3`).	required

Source code in multimolecule/models/mmsplice/configuration_mmsplice.py

Python
class MmSpliceModuleConfig(FlatDict):
    r"""
    Configuration for a single MMSplice sub-module.

    MMSplice is a *modular* model: each genomic region (the acceptor and donor
    splice sites, the exon body, and the two flanking intron stubs) is scored by
    an independent small network. The five upstream sub-networks
    ([gagneurlab/MMSplice_MTSplice](https://github.com/gagneurlab/MMSplice_MTSplice))
    do **not** share an architecture, so each is described by its own
    `MmSpliceModuleConfig`.

    Two architecture families exist:

    - `conv` (``acceptor_intron``, ``exon``, ``donor_intron``): a single
      length-preserving 1D convolution, optional batch-norm, ReLU, global
      average pooling over positions, then a linear projection to a scalar.
      Length-independent.
    - `dense` (``acceptor``, ``donor``): a fixed-length splice-site network.
      The region is one-hot encoded, optionally passed through one or more
      convolutions, flattened, and projected to a scalar with a stack of
      dense + batch-norm + ReLU blocks. A final sigmoid produces a probability
      that is converted to a logit score.

    Args:
        architecture:
            Either `conv` or `dense` (see above).
        region_length:
            Fixed input length the module consumes. `0` for the length-
            independent `conv` modules (``acceptor_intron`` / ``donor_intron``).
        conv_channels:
            Output channels of the (first) convolution.
        conv_kernel_size:
            Kernel size of the (first) convolution.
        conv_activation:
            Activation applied to the (first) convolution.
        conv_batch_norm:
            Whether a batch-norm follows the (first) convolution.
        pool_mask_zeros:
            Whether global average pooling ignores all-zero input positions.
            Upstream `exon` uses masked pooling to ignore `N` padding.
        pointwise_channels:
            Output channels of the `dense`-family `1x1` convolution. `0`
            disables it. Followed by a batch-norm.
        hidden_sizes:
            Output sizes of the dense blocks of a `dense`-family head.
        flatten_dropout:
            Whether a dropout is applied right after the flatten (before the
            dense blocks). Upstream `acceptor` uses this; `donor` instead applies
            dropout inside each dense block.
        dropout:
            Dropout probability used by `dense`-family heads.
        batch_norm_eps:
            Epsilon of every batch-norm (upstream Keras default `1e-3`).
    """

    architecture: str = "conv"
    region_length: int = 0
    conv_channels: int = 0
    conv_kernel_size: int = 0
    conv_activation: str = "linear"
    conv_batch_norm: bool = False
    pool_mask_zeros: bool = False
    pointwise_channels: int = 0
    hidden_sizes: list = []  # noqa: RUF012
    flatten_dropout: bool = False
    dropout: float = 0.2
    batch_norm_eps: float = 1e-3

MmSpliceForSequencePrediction ¶

Bases: MmSplicePreTrainedModel

MMSplice with a sequence-level prediction head.

The head consumes the per-module score deltas for a reference/alternative sequence pair and applies the fixed upstream linear combiner to produce the delta-logit-PSI splicing-effect score.

Examples:

Python Console Session
>>> import torch
>>> from multimolecule import MmSpliceConfig, MmSpliceForSequencePrediction
>>> config = MmSpliceConfig()
>>> model = MmSpliceForSequencePrediction(config)
>>> _ = model.eval()
>>> input_ids = torch.randint(4, (1, 400))
>>> alternative_input_ids = torch.randint(4, (1, 400))
>>> output = model(input_ids, alternative_input_ids=alternative_input_ids)
>>> output["logits"].shape
torch.Size([1, 1])

Source code in multimolecule/models/mmsplice/modeling_mmsplice.py

Python
class MmSpliceForSequencePrediction(MmSplicePreTrainedModel):
    """
    MMSplice with a sequence-level prediction head.

    The head consumes the per-module score deltas for a reference/alternative
    sequence pair and applies the fixed upstream linear combiner to produce the
    delta-logit-PSI splicing-effect score.

    Examples:
        >>> import torch
        >>> from multimolecule import MmSpliceConfig, MmSpliceForSequencePrediction
        >>> config = MmSpliceConfig()
        >>> model = MmSpliceForSequencePrediction(config)
        >>> _ = model.eval()
        >>> input_ids = torch.randint(4, (1, 400))
        >>> alternative_input_ids = torch.randint(4, (1, 400))
        >>> output = model(input_ids, alternative_input_ids=alternative_input_ids)
        >>> output["logits"].shape
        torch.Size([1, 1])
    """

    def __init__(self, config: MmSpliceConfig):
        super().__init__(config)
        self.model = MmSpliceModel(config)
        self.prediction = MmSpliceDeltaLogitPsiHead()
        head = config.head
        if head is None:
            raise ValueError("MmSpliceForSequencePrediction requires `config.head` to be set")
        self.criterion = Criterion(head)
        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        alternative_input_ids: Tensor | NestedTensor | None = None,
        alternative_attention_mask: Tensor | None = None,
        alternative_inputs_embeds: Tensor | NestedTensor | None = None,
        labels: Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | SequencePredictorOutput:
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            alternative_input_ids=alternative_input_ids,
            alternative_attention_mask=alternative_attention_mask,
            alternative_inputs_embeds=alternative_inputs_embeds,
            return_dict=True,
            **kwargs,
        )

        if outputs.delta_logits is None:
            raise ValueError(
                "MmSpliceForSequencePrediction requires an alternative sequence to compute delta-logit-PSI. "
                "Use MmSpliceModel for reference-only module scores."
            )
        logits = self.prediction(outputs.delta_logits)
        loss = self.criterion(logits, labels) if labels is not None else None

        return SequencePredictorOutput(loss=loss, logits=logits)

MmSpliceModel ¶

Bases: MmSplicePreTrainedModel

The bare MMSplice modular backbone.

MMSplice scores the exon-with-flanking-introns sequence with five independent sub-networks. The backbone returns the per-module score vector. For variant effect prediction, pass both a reference and an alternative sequence; the backbone then also returns the per-module score deltas.

The five sub-networks do not share an architecture; each faithfully replicates the corresponding upstream gagneurlab/MMSplice_MTSplice Keras module (Cheng et al. 2019, Genome Biology).

Examples:

Python Console Session
>>> import torch
>>> from multimolecule import MmSpliceConfig, MmSpliceModel
>>> config = MmSpliceConfig()
>>> model = MmSpliceModel(config)
>>> _ = model.eval()
>>> input_ids = torch.randint(4, (1, 400))
>>> output = model(input_ids)
>>> output["logits"].shape
torch.Size([1, 5])

Source code in multimolecule/models/mmsplice/modeling_mmsplice.py

Python
class MmSpliceModel(MmSplicePreTrainedModel):
    """
    The bare MMSplice modular backbone.

    MMSplice scores the exon-with-flanking-introns sequence with five independent
    sub-networks. The backbone returns the per-module score vector. For variant
    effect prediction, pass both a reference and an alternative sequence; the
    backbone then also returns the per-module score deltas.

    The five sub-networks do not share an architecture; each faithfully
    replicates the corresponding upstream
    [gagneurlab/MMSplice_MTSplice](https://github.com/gagneurlab/MMSplice_MTSplice)
    Keras module (Cheng et al. 2019, *Genome Biology*).

    Examples:
        >>> import torch
        >>> from multimolecule import MmSpliceConfig, MmSpliceModel
        >>> config = MmSpliceConfig()
        >>> model = MmSpliceModel(config)
        >>> _ = model.eval()
        >>> input_ids = torch.randint(4, (1, 400))
        >>> output = model(input_ids)
        >>> output["logits"].shape
        torch.Size([1, 5])
    """

    def __init__(self, config: MmSpliceConfig):
        super().__init__(config)
        self.embeddings = MmSpliceEmbedding(config)
        self.region_models = nn.ModuleDict({name: MmSpliceModule(config, name) for name in MODULE_ORDER})
        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        alternative_input_ids: Tensor | NestedTensor | None = None,
        alternative_attention_mask: Tensor | None = None,
        alternative_inputs_embeds: Tensor | NestedTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> MmSpliceModelOutput | tuple[Tensor, ...]:
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        if input_ids is None and inputs_embeds is None:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        reference = self._score(input_ids, attention_mask, inputs_embeds)

        delta = None
        alternative = None
        has_alternative = alternative_input_ids is not None or alternative_inputs_embeds is not None
        if has_alternative:
            if alternative_input_ids is not None and alternative_inputs_embeds is not None:
                raise ValueError("You cannot specify both alternative_input_ids and alternative_inputs_embeds")
            alternative = self._score(
                alternative_input_ids,
                alternative_attention_mask,
                alternative_inputs_embeds,
            )
            delta = alternative - reference

        return MmSpliceModelOutput(
            logits=reference,
            alternative_logits=alternative,
            delta_logits=delta,
        )

    def _score(
        self,
        input_ids: Tensor | NestedTensor | None,
        attention_mask: Tensor | None,
        inputs_embeds: Tensor | NestedTensor | None,
    ) -> Tensor:
        if isinstance(input_ids, NestedTensor):
            if attention_mask is None:
                attention_mask = input_ids.mask
            input_ids = input_ids.tensor
        if isinstance(inputs_embeds, NestedTensor):
            if attention_mask is None:
                attention_mask = inputs_embeds.mask
            inputs_embeds = inputs_embeds.tensor
        embedding_output = self.embeddings(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
        )
        scores = []
        for name in MODULE_ORDER:
            module = self.region_models[name]
            if self.gradient_checkpointing and self.training:
                score = self._gradient_checkpointing_func(module.__call__, embedding_output)
            else:
                score = module(embedding_output)
            scores.append(score)
        return torch.cat(scores, dim=-1)

MmSpliceModelOutput `dataclass` ¶

Bases: ModelOutput

Base class for outputs of the MMSplice modular model.

Parameters:

Name	Type	Description	Default
`logits` ¶	`torch.FloatTensor` of shape `(batch_size, hidden_size)`	The per-module score vector for the (reference) input sequence. The module order is `acceptor_intron`, `acceptor`, `exon`, `donor`, `donor_intron`.	`None`
`alternative_logits` ¶	`torch.FloatTensor` of shape `(batch_size, hidden_size)`, optional	The per-module score vector for the alternative sequence, returned when an alternative sequence is provided.	`None`
`delta_logits` ¶	`torch.FloatTensor` of shape `(batch_size, hidden_size)`, optional	`alternative_logits - logits`, the per-module variant-effect deltas, returned when an alternative sequence is provided.	`None`

Source code in multimolecule/models/mmsplice/modeling_mmsplice.py

Python
@dataclass
class MmSpliceModelOutput(ModelOutput):
    """
    Base class for outputs of the MMSplice modular model.

    Args:
        logits (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
            The per-module score vector for the (reference) input sequence. The
            module order is `acceptor_intron`, `acceptor`, `exon`, `donor`,
            `donor_intron`.
        alternative_logits (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*):
            The per-module score vector for the alternative sequence, returned when
            an alternative sequence is provided.
        delta_logits (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*):
            `alternative_logits - logits`, the per-module variant-effect deltas,
            returned when an alternative sequence is provided.
    """

    logits: torch.FloatTensor | None = None
    alternative_logits: torch.FloatTensor | None = None
    delta_logits: torch.FloatTensor | None = None

MmSplicePreTrainedModel ¶

Bases: PreTrainedModel

An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.

Source code in multimolecule/models/mmsplice/modeling_mmsplice.py

Python
class MmSplicePreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = MmSpliceConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _can_record_outputs: dict[str, Any] | None = None
    _no_split_modules = ["MmSpliceModule"]

    @torch.no_grad()
    def _init_weights(self, module):
        super()._init_weights(module)
        if isinstance(module, nn.Conv1d):
            init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
            if module.bias is not None:
                init.zeros_(module.bias)
        elif isinstance(module, nn.Linear):
            init.kaiming_uniform_(module.weight, a=math.sqrt(5))
            if module.bias is not None:
                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
                bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
                init.uniform_(module.bias, -bound, bound)
        elif isinstance(module, (nn.BatchNorm1d, nn.LayerNorm)):
            init.ones_(module.weight)
            init.zeros_(module.bias)

MMSplice¶

Disclaimer¶

Model Details¶

Model Specification¶

Links¶

Usage¶

Direct Use¶

Module Scores¶

Variant Effect¶

Interface¶

Variant Effect¶

Training Details¶

Training Data¶

Training Procedure¶

Pre-training¶

Citation¶

Contact¶

License¶

API Reference¶

MmSpliceConfig ¶

vocab_size ¶

modules ¶

modules_config ¶

acceptor_intron_cut ¶

donor_intron_cut ¶

acceptor_intron_length ¶

acceptor_exon_length ¶

donor_exon_length ¶

donor_intron_length ¶

num_labels ¶

head ¶

MmSpliceModuleConfig ¶

architecture ¶

region_length ¶

conv_channels ¶

conv_kernel_size ¶

conv_activation ¶

conv_batch_norm ¶

pool_mask_zeros ¶

pointwise_channels ¶

hidden_sizes ¶

flatten_dropout ¶

dropout ¶

batch_norm_eps ¶

MmSpliceForSequencePrediction ¶

MmSpliceModel ¶

MmSpliceModelOutput dataclass ¶

logits ¶

alternative_logits ¶

delta_logits ¶

MmSplicePreTrainedModel ¶

`vocab_size` ¶

`modules` ¶

`modules_config` ¶

`acceptor_intron_cut` ¶

`donor_intron_cut` ¶

`acceptor_intron_length` ¶

`acceptor_exon_length` ¶

`donor_exon_length` ¶

`donor_intron_length` ¶

`num_labels` ¶

`head` ¶

`architecture` ¶

`region_length` ¶

`conv_channels` ¶

`conv_kernel_size` ¶

`conv_activation` ¶

`conv_batch_norm` ¶

`pool_mask_zeros` ¶

`pointwise_channels` ¶

`hidden_sizes` ¶

`flatten_dropout` ¶

`dropout` ¶

`batch_norm_eps` ¶

MmSpliceModelOutput `dataclass` ¶

`logits` ¶

`alternative_logits` ¶

`delta_logits` ¶