pandora_llm.features.ModelStealing

Module Contents

class pandora_llm.features.ModelStealing.ModelStealing(*args, **kwargs)[source]

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

Model stealing attack

prepare_projection(svd_dataloader, proj_type='rademacher', proj_dim=512, proj_seed=229, device=None, fp16=False)[source]

Compute the embedding projection layer for the gray-box model-stealing attack

Parameters:
  • svd_dataloader (torch.utils.data.DataLoader) – input data to estimate projection layer

  • proj_type (str) – projection type (defualt “rademacher”)

  • proj_dim (int) – project to how many dimensions (default 512)

  • proj_seed (int) – random seed for random projection

  • device (str) – device to compute on: “cuda” or “cpu”

  • fp16 (bool) – whether to use fp16

Returns:

grad norm of input IDs

Return type:

torch.Tensor or list

compute_features(dataloader, svd_embedding_projection_layer, projector, device=None, fp16=None)[source]

Computes dataloader gradients with jl dimensionality reduction. :param dataloader: DataLoader of samples. :param svd_embedding_projection_layer: dictionary of dimensionality reduction functions :param device: CPU or GPU :param fp16: use half precision floats for model

Returns:

data for input IDs

Return type:

torch.Tensor or list

Raises:

Exception – must load the model first

Parameters:
  • dataloader (torch.utils.data.DataLoader)

  • svd_embedding_projection_layer (jaxtyping.Float[torch.Tensor, ...])

  • projector (trak.projectors.AbstractProjector)

  • device (str)

  • fp16 (bool)

pandora_llm.features.ModelStealing.compute_dataloader_basis_changes(model, dataloader, svd_embedding_projection_layer, projector, device=None, fp16=False)[source]

Computes dataloader gradients with jl dimensionality reduction. :param model: HuggingFace model. :param dataloader: DataLoader of samples. :param projector: dictionary of dimensionality reduction functions :param device: CPU or GPU :param fp16: use half precision floats for model

Returns:

JL-reduced Dataloader gradients for input IDs

Parameters:
  • model (transformers.AutoModelForCausalLM)

  • dataloader (torch.utils.data.DataLoader)

  • svd_embedding_projection_layer (jaxtyping.Float[torch.Tensor, ...])

  • projector (trak.projectors.AbstractProjector)

  • device (str)

  • fp16 (bool)

Return type:

jaxtyping.Float[torch.Tensor, n proj_dim]

pandora_llm.features.ModelStealing.compute_basis_change(model, svd_embedding_projection_layer, projector, input_ids, attention_mask=None, device=None)[source]

This computes the basis change for the last layer (Carlini et al. gray-box attack), and returns it with the norms of that layer.

Parameters:
  • model (transformers.AutoModelForCausalLM) – HuggingFace model.

  • input_ids (jaxtyping.Integer[torch.Tensor, batch seq]) – tensor of input IDs.

  • svd_embedding_projection_layer (jaxtyping.Float[torch.Tensor, ...]) – embedding layer

  • projector (trak.projectors.AbstractProjector) – dimensionality reduction function

  • device – CPU or GPU

  • attention_mask (jaxtyping.Bool[torch.Tensor, batch seq])

Returns:

Basis change for the last layer

Return type:

jaxtyping.Float[torch.Tensor, …]

pandora_llm.features.ModelStealing.compute_dataloader_logits_embedding(model, dataloader, device=None, fp16=False)[source]

Computes logits of text in dataloader

Parameters:
  • model (transformers.AutoModelForCausalLM) – HuggingFace model.

  • dataloader (torch.utils.data.DataLoader) – DataLoader with tokens.

  • device (str) – CPU or GPU

  • fp16 (bool) – use half precision floats for model

Returns:

Tensor of logits for last token

Return type:

jaxtyping.Float[torch.Tensor, n batch vocab]

pandora_llm.features.ModelStealing.compute_input_ids_logits(model, input_ids, attention_mask=None, device=None)[source]

Compute logits of last token in input ids

Parameters:
  • model (transformers.AutoModelForCausalLM) – HuggingFace model.

  • input_ids (jaxtyping.Integer[torch.Tensor, batch seq]) – tensor of input IDs

  • attention_mask (jaxtyping.Bool[torch.Tensor, batch seq]) – attention_mask

  • device (str) – CPU or GPU

Returns:

Logits of last token in input ids

Return type:

jaxtyping.Float[torch.Tensor, batch vocab]