Skip to content

MMSplice

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 FLOPs (M) MACs (M)
5 56,677 5.71 2.79

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

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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
>>> import torch
>>> from multimolecule import DnaTokenizer, MmSpliceForSequencePrediction

>>> tokenizer = DnaTokenizer.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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
>>> import torch
>>> from multimolecule import DnaTokenizer, MmSpliceForSequencePrediction

>>> tokenizer = DnaTokenizer.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] + "T" + 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/T 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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
@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 you use MultiMolecule in your research, you must cite the MultiMolecule project as follows:

BibTeX
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
@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

multimolecule.models.mmsplice

DnaTokenizer

Bases: Tokenizer

Tokenizer for DNA sequences.

Parameters:

Name Type Description Default

alphabet

Alphabet | str | List[str] | None

alphabet to use for tokenization.

  • If is None, the standard RNA alphabet will be used.
  • If is a string, it should correspond to the name of a predefined alphabet. The options include
    • standard
    • iupac
    • streamline
    • nucleobase
  • If is an alphabet or a list of characters, that specific alphabet will be used.
 None

nmers

int

Size of kmer to tokenize.

 1

codon

bool

Whether to tokenize into codons.

 False

replace_U_with_T

bool

Whether to replace U with T.

 True

do_upper_case

bool

Whether to convert input to uppercase.

 True

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
>>> from multimolecule import DnaTokenizer
>>> tokenizer = DnaTokenizer()
>>> tokenizer('<pad><cls><eos><unk><mask><null>ACGTNRYSWKMBDHVX|.*-?')["input_ids"]
[1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 2]
>>> tokenizer('acgt')["input_ids"]
[1, 6, 7, 8, 9, 2]
>>> tokenizer('acgu')["input_ids"]
[1, 6, 7, 8, 9, 2]
>>> tokenizer = DnaTokenizer(replace_U_with_T=False)
>>> tokenizer('acgu')["input_ids"]
[1, 6, 7, 8, 3, 2]
>>> tokenizer = DnaTokenizer(nmers=3)
>>> tokenizer('tataaagta')["input_ids"]
[1, 84, 21, 81, 6, 8, 19, 71, 2]
>>> tokenizer = DnaTokenizer(codon=True)
>>> tokenizer('tataaagta')["input_ids"]
[1, 84, 6, 71, 2]
>>> tokenizer('tataaagtaa')["input_ids"]
Traceback (most recent call last):
ValueError: length of input sequence must be a multiple of 3 for codon tokenization, but got 10
Source code in multimolecule/tokenisers/dna/tokenization_dna.py
Python
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
class DnaTokenizer(Tokenizer):
    """
    Tokenizer for DNA sequences.

    Args:
        alphabet: alphabet to use for tokenization.

            - If is `None`, the standard RNA alphabet will be used.
            - If is a `string`, it should correspond to the name of a predefined alphabet. The options include
                + `standard`
                + `iupac`
                + `streamline`
                + `nucleobase`
            - If is an alphabet or a list of characters, that specific alphabet will be used.
        nmers: Size of kmer to tokenize.
        codon: Whether to tokenize into codons.
        replace_U_with_T: Whether to replace U with T.
        do_upper_case: Whether to convert input to uppercase.

    Examples:
        >>> from multimolecule import DnaTokenizer
        >>> tokenizer = DnaTokenizer()
        >>> tokenizer('<pad><cls><eos><unk><mask><null>ACGTNRYSWKMBDHVX|.*-?')["input_ids"]
        [1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 2]
        >>> tokenizer('acgt')["input_ids"]
        [1, 6, 7, 8, 9, 2]
        >>> tokenizer('acgu')["input_ids"]
        [1, 6, 7, 8, 9, 2]
        >>> tokenizer = DnaTokenizer(replace_U_with_T=False)
        >>> tokenizer('acgu')["input_ids"]
        [1, 6, 7, 8, 3, 2]
        >>> tokenizer = DnaTokenizer(nmers=3)
        >>> tokenizer('tataaagta')["input_ids"]
        [1, 84, 21, 81, 6, 8, 19, 71, 2]
        >>> tokenizer = DnaTokenizer(codon=True)
        >>> tokenizer('tataaagta')["input_ids"]
        [1, 84, 6, 71, 2]
        >>> tokenizer('tataaagtaa')["input_ids"]
        Traceback (most recent call last):
        ValueError: length of input sequence must be a multiple of 3 for codon tokenization, but got 10
    """

    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        alphabet: Alphabet | str | List[str] | None = None,
        nmers: int = 1,
        codon: bool = False,
        replace_U_with_T: bool = True,
        do_upper_case: bool = True,
        additional_special_tokens: List | Tuple | None = None,
        **kwargs,
    ):
        if codon and (nmers > 1 and nmers != 3):
            raise ValueError("Codon and nmers cannot be used together.")
        if codon:
            nmers = 3  # set to 3 to get correct vocab
        if not isinstance(alphabet, Alphabet):
            alphabet = get_alphabet(alphabet, nmers=nmers)
        super().__init__(
            alphabet=alphabet,
            nmers=nmers,
            codon=codon,
            replace_U_with_T=replace_U_with_T,
            do_upper_case=do_upper_case,
            additional_special_tokens=additional_special_tokens,
            **kwargs,
        )
        self.replace_U_with_T = replace_U_with_T
        self.nmers = nmers
        self.codon = codon

    def _tokenize(self, text: str, **kwargs):
        if self.do_upper_case:
            text = text.upper()
        if self.replace_U_with_T:
            text = text.replace("U", "T")
        if self.codon:
            if len(text) % 3 != 0:
                raise ValueError(
                    f"length of input sequence must be a multiple of 3 for codon tokenization, but got {len(text)}"
                )
            return [text[i : i + 3] for i in range(0, len(text), 3)]
        if self.nmers > 1:
            return [text[i : i + self.nmers] for i in range(len(text) - self.nmers + 1)]  # noqa: E203
        return list(text)

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/T channels; N, padding, special, and unknown tokens are encoded as all-zero columns by the embedding layer. Defaults to 4 (the ACGT 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
1
2
3
4
5
6
7
>>> 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
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/T` channels; `N`, padding, special, and unknown tokens are
            encoded as all-zero columns by the embedding layer.
            Defaults to 4 (the `ACGT` 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 Type 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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
>>> 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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
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
1
2
3
4
5
6
7
8
9
>>> 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
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
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
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
@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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
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"]

    def _init_weights(self, module):
        if isinstance(module, nn.Conv1d):
            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Linear):
            nn.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
                nn.init.uniform_(module.bias, -bound, bound)
        elif isinstance(module, (nn.BatchNorm1d, nn.LayerNorm)):
            nn.init.constant_(module.weight, 1)
            nn.init.constant_(module.bias, 0)