Multinomial diffusion

`multinomial_diffusion`

`MODEL_TYPE = 'diffusion'` `module-attribute`

`logger = ColorLog(console, name).logger` `module-attribute`

`InstaNovoPlus(config: DictConfig, transition_model: nn.Module, diffusion_schedule: Float[torch.Tensor, ' time'], residue_set: ResidueSet)`

Bases: Module

This class implements Multinomial Diffusion as described in Hoogeboom et al. 2021.

PARAMETER	DESCRIPTION
`config`	The model configuration. This should have keys: - 'name': the model name identifier. - 'time_steps': the number of time steps in the diffusion process - 'max_length': the maximum sequence for the model - 'device': the device where the `Pytorch` model should be loaded e.g. `cpu`, `cuda:0` etc. - 'vocab_size': the number of residues in the vocabulary - 'transition_model': the `DictConfig` for the transition model This information is necessary for saving and loading the model. TYPE: `DictConfig`
`transition_model`	The model that predictions the initial sequence given the sequence sampled the current time step and the sequence sampled the previous time step. This is just a sequence tagging model. TYPE: `Module`
`diffusion_schedule`	The sequence of diffusion probabilities. Note that `diffusion_schedule[t]` is \alpha_t in the paper's terminology, not \beta_t. TYPE: `FloatTensor[time_steps]`
`residue_set`	The residue vocabulary. This holds a mapping between residues and indices and residue masses. TYPE: `ResidueSet`

`config_path: str` `instance-attribute`

`schedule_path: str` `instance-attribute`

`checkpoint_path: str` `instance-attribute`

`config = config` `instance-attribute`

`time_steps = config.time_steps` `instance-attribute`

`residue_set = residue_set` `instance-attribute`

`transition_model = transition_model` `instance-attribute`

`save(path: str, ckpt_details: str, overwrite: bool = False, temp_dir: str | None = None, use_legacy_format: bool = False) -> None`

Save the model to a directory.

PARAMETER	DESCRIPTION
`path`	Path to the base directory where the model is saved. The model is saved in a subdirectory with the model's name identifier. TYPE: `str`
`ckpt_details`	Additional checkpoint details to include in model save directory. TYPE: `str`
`overwrite`	Whether to overwrite the directory if one already exists for the model. Defaults to False. TYPE: `bool` DEFAULT: `False`
`temp_dir`	Temporary directory to save intermediate files to. Defaults to None. TYPE: `str \| None` DEFAULT: `None`
`use_legacy_format`	Whether to save the model in the legacy folder format. If False, saves as a single file. Defaults to False. TYPE: `bool` DEFAULT: `False`

RAISES	DESCRIPTION
`FileExistsError`	If `overwrite` is `False` and a directory already exists for the model identifier.

`load(path: str, override_config: DictConfig | dict | None = None) -> Tuple[InstaNovoPlus, DictConfig]` `classmethod`

Load a saved model.

PARAMETER	DESCRIPTION
`path`	Path to model checkpoint file or directory where model is saved. TYPE: `str`
`override_config`	Optional override config values with a DictConfig or dict, defaults to None. TYPE: `DictConfig \| dict \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`(InstaNovoPlus, DictConfig)`	The loaded model and config.

`get_pretrained() -> list[str]` `staticmethod`

Get a list of pretrained model ids.

`from_pretrained(model_id: str, override_config: DictConfig | dict | None = None) -> Tuple['InstaNovoPlus', 'DictConfig']` `classmethod`

Download and load by model id or model path.

`prepare_fine_tuning(residue_set: ResidueSet) -> None`

Prepare a model for fine-tuning on a dataset with a new residue vocabulary.

PARAMETER	DESCRIPTION
`residue_set`	The residue vocabulary for the new dataset. TYPE: `ResidueSet`

`mixture_categorical(log_x: Float[ResidueLogProbabilities, 'batch token'], log_alpha: float, log_alpha_complement: float) -> Float[ResidueLogProbabilities, 'batch token']`

A categorical mixture between a base distribution and a uniform distribution.

PARAMETER	DESCRIPTION
`log_x`	The base distribution. TYPE: `FloatTensor[..., num_classes]`
`log_alpha`	The log of the mixture weight. TYPE: `float`
`log_alpha_complement`	The log of 1 minus the mixture weight. TYPE: `float`

RETURNS	DESCRIPTION
`Float[ResidueLogProbabilities, 'batch token']`	torch.FloatTensor[..., num_classes]: The log-probabilities of the mixture.

`forward(log_x_t: Float[ResidueLogProbabilities, 'batch token'], log_x_0: Float[ResidueLogProbabilities, 'batch token'], t: Integer[TimeStep, ' batch']) -> Float[ResidueLogProbabilities, 'batch token']`

Calculate the log-posterior of t-1-th process values given the 0-th and t-th values.

PARAMETER	DESCRIPTION
`log_x_t`	The log one-hot representation of the process values at the `t`-th time step. TYPE: `FloatTensor[batch_size, sequence_length, num_classes]`
`log_x_0`	The log one-hot representation of the process values at the `t`-th time step. TYPE: `FloatTensor[batch_size, sequence_length, num_classes]`
`t`	The time step. TYPE: `int`

RETURNS	DESCRIPTION
`Float[ResidueLogProbabilities, 'batch token']`	torch.FloatTensor[batch_size, sequence_length, num_classes]: The log-posterior probabilities of the process values at the `t-1`-th time step given the values at the 0-th and `t`-th time step i.e. q( x_{t-1} \| x_{t}, x_0 ).

`reverse_distribution(x_t: Integer[Peptide, 'batch token'], time: Integer[TimeStep, ' batch'], **kwargs: dict) -> Float[ResidueLogProbabilities, 'batch token']`

Calculate the reverse transition distribution of the diffusion process.

PARAMETER	DESCRIPTION
`x_t`	The values at the `t`-th time step of the reverse process. TYPE: `LongTensor[batch_size, sequence_length]`
`time`	The time step. TYPE: `int`

RETURNS	DESCRIPTION
`Float[ResidueLogProbabilities, 'batch token']`	torch.FloatTensor[batch_size, sequence_length, num_classes]: The log-probabilities of values for the `t-1`-th time step given values at the `t`-th time step i.e. `log p( x_{t-1} \| x_{t} )`.

`DiffusionLoss(model: InstaNovoPlus)`

Bases: Module

Holds logic for calculating the diffusion loss.

PARAMETER	DESCRIPTION
`model`	The multinomial diffusion class. TYPE: `InstaNovoPlus`

`model = model` `instance-attribute`

`base_model = model.module if hasattr(model, 'module') else model` `instance-attribute`

`time_steps = self.base_model.time_steps` `instance-attribute`

`kl_divergence(log_probs_first: Float[ResidueLogProbabilities, '...'], log_probs_second: Float[ResidueLogProbabilities, '...']) -> Float[torch.Tensor, '...']` `staticmethod`

Calculate the Kullback-Liebler divergence between two multinomial distributions.

PARAMETER	DESCRIPTION
`log_probs_first`	The log-probabilities of the base distribution. TYPE: `FloatTensor[..., num_classes]`
`log_probs_second`	The log-probabilities of the comparison distribution. TYPE: `FloatTensor[..., num_classes]`

RETURNS	DESCRIPTION
`Float[Tensor, '...']`	torch.FloatTensor[1]: The KL-divergence averaged over all but the final dimension.

`forward(x_0: Integer[Peptide, 'batch token'], **kwargs: dict) -> Float[torch.Tensor, '1']`

Calculate a single Monte Carlo estimate of the multinomial diffusion loss (-ELBO).

PARAMETER	DESCRIPTION
`x_0`	A batch of padded sequences. TYPE: `LongTensor[batch_size, sequence_length]`

RETURNS	DESCRIPTION
`Float[Tensor, '1']`	torch.FloatTensor[1]: The loss estimate.

`cosine_beta_schedule(timesteps: int, s: float = 0.008) -> Float[torch.Tensor, ' time']`

Cosine schedule as proposed in https://arxiv.org/abs/2102.09672 .

Returns alpha parameters, NOT Beta

Multinomial diffusion