pandora_llm.features.Quantile

Module Contents

class pandora_llm.features.Quantile.Quantile(model_name, ref_model_name, model_revision=None, model_cache_dir=None, ref_model_revision=None, ref_model_cache_dir=None)

Bases: pandora_llm.features.base.FeatureComputer, pandora_llm.features.base.LLMHandler

Computes the quantile score from Zhang et al. 2024 (https://arxiv.org/pdf/2409.14513). It trains a weak ensemble to estimate the mean and stddev conditional on nonmember data using a “gaussian+pinball” loss. The reference is then mean+normal_cdf^{-1}(1-alpha)*stddev. Currently supports GPTNeoX models as reference.

Parameters:

• model_name (str)

• ref_model_name (str)

• model_revision (str)

• model_cache_dir (str)

• ref_model_revision (str)

• ref_model_cache_dir (str)

ref_model_name

ref_model_revision = None

ref_model_cache_dir = None

ref_model_paths = []

load_model(model_index)

Loads the specified model into memory.

Parameters:

model_index (int) – Index of the model to load. Base model corresponds to index 0; quantile regression models are 1-indexed.

Raises:

• IndexError – If the model_index is out of bounds.

• Exception – If a model is already loaded.

Return type:

None

finetune_ref(nonmember_train_tokens, neg_log_probs, accelerator)

Fine-tunes reference model ensemble to do quantile regression

Parameters:

• nonmember_train_tokens (jaxtyping.Integer[torch.Tensor, n seq]) – nonmember train tokens

• neg_log_probs (jaxtyping.Float[torch.Tensor, n]) – negative log probs of nonmember train tokens

• accelerator (accelerate.Accelerator) – accelerator object

Return type:

None

compute_features(dataloader, accelerator, mode='primary')

Computes the negative log-likelihood (NLL) feature for the given dataloader.

Parameters:

• dataloader (torch.utils.data.DataLoader) – The dataloader providing input sequences.

• accelerator (accelerate.Accelerator) – The Accelerator object for distributed or mixed-precision training.

• mode (str) – whether to compute the “primary” score or the “ref” score

Returns:

The NLL feature for each sequence in the dataloader.

Raises:

• Exception – If the model is not loaded before calling this method.

• ValueError – If the model stage is unexpected

Return type:

jaxtyping.Float[torch.Tensor, n]

static reduce(primary_log_probs, ref_mu_sigma)

Computes quantile attack by computing primary_log_probs minus reference. The reference is mean+normal_cdf^{-1}(1-alpha)*stddev.

Parameters:

• primary_log_probs (jaxtyping.Float[torch.Tensor, n]) – Log probs from primary model

• ref_mu_sigma (jaxtyping.Float[torch.Tensor, n num_models 2]) – Nonmember mu and sigma from reference model

Returns:

primary minus ref

Return type:

jaxtyping.Float[torch.Tensor, n]

pandora_llm.features.Quantile.gaussian_loss_fn(score, target, eps=0.0001, quantile=None, kl_weight=0.0, return_kl=False, ignore_index=-100)

pandora_llm.features.Quantile.pinball_loss_fn(score, target, quantiles, ignore_index=-100)

pandora_llm.features.Quantile.gaussian_pinball_loss_fn(score, target, eps=0.0001, quantile=None, ignore_index=-100)

pandora_llm.features.Quantile.compute_metrics(p)

Parameters:

p (transformers.EvalPrediction)

class pandora_llm.features.Quantile.GPTNeoXForQuantileRegression(config)

Bases: transformers.models.gpt_neox.GPTNeoXPreTrainedModel

num_labels = 2

var_nonlin

config

gpt_neox

score

forward(input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, head_mask=None, past_key_values=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None)

labels (torch.LongTensor of shape (batch_size,), optional):

Labels for computing the sequence classification/regression loss. Indices should be in [0, …, config.num_labels - 1]. If config.num_labels == 1 a regression loss is computed (Mean-Square loss), If config.num_labels > 1 a classification loss is computed (Cross-Entropy).

Parameters:

• input_ids (Optional[torch.LongTensor])

• attention_mask (Optional[torch.FloatTensor])

• position_ids (Optional[torch.LongTensor])

• inputs_embeds (Optional[torch.FloatTensor])

• head_mask (Optional[torch.FloatTensor])

• past_key_values (Optional[Tuple[Tuple[torch.FloatTensor]]])

• labels (Optional[torch.LongTensor])

• use_cache (Optional[bool])

• output_attentions (Optional[bool])

• output_hidden_states (Optional[bool])

• return_dict (Optional[bool])

Return type:

Union[Tuple[torch.Tensor], transformers.modeling_outputs.SequenceClassifierOutputWithPast]

pandora_llm.features.Quantile.compute_quantile_regression_dl(model, dataloader, accelerator)

Parameters:

• model (GPTNeoXForQuantileRegression)

• dataloader (torch.utils.data.DataLoader)

• accelerator (accelerate.Accelerator)

Return type:

jaxtyping.Float[torch.Tensor, n 2]