Step Functions

Step Functions Default arguments

The default arguments passed to all step functions are collected in the StepFunctionBaseArgs class.

class inseq.attr.step_functions.StepFunctionBaseArgs(attribution_model: AttributionModel, forward_output: ModelOutput, target_ids: Int64[Tensor, 'batch_size'], decoder_input_ids: Int64[Tensor, 'batch_size seq_len'], decoder_input_embeds: Float[Tensor, 'batch_size seq_len embed_size'], decoder_attention_mask: Int64[Tensor, 'batch_size seq_len'], is_attributed_fn: bool)[source]

Base class for step function base arguments. These arguments are passed to all step functions and are complemented by the ones defined in the step function signature.

attribution_model

The attribution model used in the current step.

Type:: AttributionModel

forward_output

The output of the model’s forward pass.

Type:: ModelOutput

target_ids

Tensor of target token ids of size (batch_size,) corresponding to the target predicted tokens for the next generation step.

Type:: torch.Tensor

is_attributed_fn

Whether the step function is being used as attribution target. Defaults to False. Enables custom behavior that is different whether the fn is used as target or not.

Type:: bool, optional, defaults to False

encoder_input_ids

Tensor of ids of encoder input tokens of size (batch_size, source_seq_len), representing encoder inputs at the present step. Available only for encoder-decoder models.

Type:: torch.Tensor

decoder_input_ids

Tensor of ids of decoder input tokens of size (batch_size, target_seq_len), representing decoder inputs at the present step.

Type:: torch.Tensor

encoder_input_embeds

Tensor of embeddings of encoder input tokens of size (batch_size, source_seq_len, hidden_size), representing encoder inputs at the present step. Available only for encoder-decoder models.

Type:: torch.Tensor

decoder_input_embeds

Tensor of embeddings of decoder input tokens of size (batch_size, target_seq_len, hidden_size), representing decoder inputs at the present step.

Type:: torch.Tensor

encoder_attention_mask

Tensor of attention mask of encoder input tokens of size (batch_size, source_seq_len), used for masking padding tokens in the encoder input. Available only for encoder-decoder models.

Type:: torch.Tensor

decoder_attention_mask

Tensor of attention mask of decoder input tokens of size (batch_size, target_seq_len), used for masking padding tokens in the decoder input.

Type:: torch.Tensor

Pre-registered Step Functions

The following functions can be used out-of-the-box as attribution targets or step functions in the inseq.models.AttributionModel.attribute() function call simply by passing their string identifier (function name minus the _fn suffix).

inseq.attr.step_functions.logit_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs) → Float32[Tensor, 'batch_size'][source]: Compute the logit of the target_ids from the model’s output logits.

inseq.attr.step_functions.probability_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs, logprob: bool = False) → Float32[Tensor, 'batch_size'][source]: Compute the probabilty of target_ids from the model’s output logits.

inseq.attr.step_functions.entropy_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs) → Float32[Tensor, 'batch_size'][source]: Compute the entropy of the model’s output distribution.

inseq.attr.step_functions.crossentropy_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs) → Float32[Tensor, 'batch_size'][source]: Compute the cross entropy between the target_ids and the logits. See: https://github.com/ZurichNLP/nmtscore/blob/master/src/nmtscore/models/m2m100.py#L99.

inseq.attr.step_functions.perplexity_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs) → Float32[Tensor, 'batch_size'][source]: Compute perplexity of the target_ids from the logits. Perplexity is the weighted branching factor. If we have a perplexity of 100, it means that whenever the model is trying to guess the next word it is as confused as if it had to pick between 100 words. Reference: https://chiaracampagnola.io/2020/05/17/perplexity-in-language-models/.

Returns the logit of a generation target given contrastive context or target prediction alternative. If only contrast_targets are specified, the logit of the contrastive prediction is computed given same context. The logit for the same token given contrastive source/target preceding context can also be computed using contrast_sources without specifying contrast_targets.

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

Returns the probability of a generation target given contrastive context or target prediction alternative. If only contrast_targets are specified, the probability of the contrastive prediction is computed given same context. The probability for the same token given contrastive source/target preceding context can also be computed using contrast_sources without specifying contrast_targets.

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

Equivalent to contrast_prob_diff_fn but for logits. The original target function used in Yin and Neubig (2022)

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

Returns the difference between next step probability for a candidate generation target vs. a contrastive alternative. Can be used as attribution target to answer the question: “Which features were salient in the choice of picking the selected token rather than its contrastive alternative?”. Follows the implementation of Yin and Neubig (2022). Can also be used to compute the difference in probability for the same token given contrastive source/target preceding context using contrast_sources without specifying contrast_targets.

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

Compute the pointwise conditional cross-mutual information (P-CXMI) of target ids given original and contrastive input options. The P-CXMI is defined as the negative log-ratio between the conditional probability of the target given the original input and the conditional probability of the target given the contrastive input, as defined by Yin et al. (2021).

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

inseq.attr.step_functions.kl_divergence_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs, contrast_sources: str | Sequence[str] | BatchEncoding | Batch | None = None, contrast_targets: str | Sequence[str] | BatchEncoding | Batch | None = None, contrast_targets_alignments: list[list[tuple[int, int]]] | None = None, top_k: int = 0, top_p: float = 1.0, min_tokens_to_keep: int = 1, contrast_force_inputs: bool = False) → Float32[Tensor, 'batch_size'][source]

Compute the pointwise Kullback-Leibler divergence of target ids given original and contrastive input options. The KL divergence is the expectation of the log difference between the probabilities of regular (P) and contrastive (Q) inputs.

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.
top_k (int) – If set to a value > 0, only the top top_k tokens will be considered for computing the KL divergence. Defaults to 0 (no top-k selection).
top_p (float) – If set to a value > 0 and < 1, only the tokens with cumulative probability above top_p will be considered for computing the KL divergence. Defaults to 1.0 (no filtering), applied before top_k filtering.
min_tokens_to_keep (int) – Minimum number of tokens to keep with top_p filtering. Defaults to 1.

Returns the in-context pointwise V-usable information as defined by Lu et al. (2023). In-context PVI is a variant of P-CXMI that captures the amount of usable information in a given contextual example, i.e. how much context information contributes to model’s prediction. In-context PVI was used by Lu et al. (2023) to estimate example difficulty for a given model, and by Prasad et al. (2023) to measure the informativeness of intermediate reasoning steps in chain-of-thought prompting.

Reference implementation: https://github.com/boblus/in-context-pvi/blob/main/in_context_pvi.ipynb

Parameters:

contrast_sources (str or list(str)) – Source text(s) used as contrastive inputs to compute the contrastive step function for encoder-decoder models. If not specified, the source text is assumed to match the original source text. Defaults to None.
contrast_targets (str or list(str)) – Contrastive target text(s) to be compared to the original target text. If not specified, the original target text is used as contrastive target (will result in same output unless contrast_sources are specified). Defaults to None.
contrast_targets_alignments (list(tuple(int, int)), optional) – A list of tuples of indices, where the first element is the index of the original target token and the second element is the index of the contrastive target token, used only if contrast_targets is specified. If an explicit alignment is not specified, the alignment of the original and contrastive target texts is assumed to be 1:1 for all available tokens. Defaults to None.
contrast_force_inputs (bool, optional, defaults to False) – Whether to force the contrastive inputs to be used for attribution. Defaults to False.

inseq.attr.step_functions.mc_dropout_prob_avg_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs, n_mcd_steps: int = 5, logprob: bool = False)[source]

Returns the average of probability scores using a pool of noisy prediction computed with MC Dropout. Can be used as an attribution target to compute more robust attribution scores.

Note

In order to obtain meaningful results, the attribution_model must contain dropout layers or other sources of noise in the forward pass.

Parameters:: n_mcd_steps (int) – The number of prediction steps that should be used to normalize the original output.

inseq.attr.step_functions.top_p_size_fn(args: StepFunctionEncoderDecoderArgs | StepFunctionDecoderOnlyArgs, top_p: float)[source]

Returns the number of tokens that have cumulative probability above top_p in the model’s output logits.

Parameters:: top_p (float) – The cumulative probability threshold to use for filtering the logits.