{"id":18293448,"url":"https://github.com/archinetai/vat-pytorch","last_synced_at":"2025-04-05T11:31:02.378Z","repository":{"id":40646611,"uuid":"506182367","full_name":"archinetai/vat-pytorch","owner":"archinetai","description":"Virtual Adversarial Training (VAT) techniques in PyTorch","archived":false,"fork":false,"pushed_at":"2022-07-19T20:58:15.000Z","size":45,"stargazers_count":17,"open_issues_count":0,"forks_count":2,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-04-04T09:02:18.230Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/archinetai.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2022-06-22T09:32:59.000Z","updated_at":"2024-12-25T05:51:36.000Z","dependencies_parsed_at":"2022-07-14T04:40:36.109Z","dependency_job_id":null,"html_url":"https://github.com/archinetai/vat-pytorch","commit_stats":null,"previous_names":[],"tags_count":9,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/archinetai%2Fvat-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/archinetai%2Fvat-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/archinetai%2Fvat-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/archinetai%2Fvat-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/archinetai","download_url":"https://codeload.github.com/archinetai/vat-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247330550,"owners_count":20921651,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-05T14:24:34.744Z","updated_at":"2025-04-05T11:31:02.078Z","avatar_url":"https://github.com/archinetai.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"\u003cimg src=\"./LOGO.png\"\u003e\u003c/img\u003e\n\nA collection of VAT (Virtual Adversarial Training) methods, in PyTorch.\n\n## Install\n\n```bash\n$ pip install vat-pytorch\n```\n\n[![PyPI - Python Version](https://img.shields.io/pypi/v/vat-pytorch?style=flat\u0026colorA=0f0f0f\u0026colorB=0f0f0f)](https://pypi.org/project/vat-pytorch/) \n\n\n## API \n\n### SMART \nThe \u003ca href=\"https://aclanthology.org/2020.acl-main.197/\"\u003eSMART paper\u003c/a\u003e proposes to find the noise that maximally perturbs the logits when added to the embedding layer, and to use a loss function to make sure that the perturbed logits are as close as possible to the predicted logits. \n\n```py \nfrom vat_pytorch import SMARTLoss, inf_norm \n\nloss = SMARTLoss(\n    model: nn.Module,\n    loss_fn: Callable,\n    loss_last_fn: Callable = None, \n    norm_fn: Callable = inf_norm, \n    num_steps: int = 1,\n    step_size: float = 1e-3, \n    epsilon: float = 1e-6,\n    noise_var: float = 1e-5\n)\n```\n\n### ALICE \n\nThe \u003ca href=\"https://arxiv.org/abs/2005.08156\"\u003eALICE paper\u003c/a\u003e is analogous to the SMART paper, but adds an additional term to make sure that the perturbed logits are as close as possible to both the predicted logits *and* the ground truth labels. \n\n```py \nfrom vat_pytorch import ALICELoss, inf_norm \n\nloss = ALICEPPLoss(\n    model: nn.Module,\n    loss_fn: Callable,\n    num_classes: int,\n    loss_last_fn: Callable = None,\n    gold_loss_fn: Callable = None,\n    gold_loss_last_fn: Callable = None,\n    norm_fn: Callable = inf_norm,\n    alpha: float = 1,\n    num_steps: int = 1,\n    step_size: float = 1e-3,\n    epsilon: float = 1e-6,\n    noise_var: float = 1e-5,\n)\n```\n\n### ALICE++ \n\nThe \u003ca href=\"https://aclanthology.org/2021.paclic-1.40/\"\u003eALICE++ paper\u003c/a\u003e is analogous to the ALICE paper, but instead of adding noise to the embedding layer, it picks a random layer from the network at each iteration on which to add the noise. \n\n```py \nfrom vat_pytorch import ALICEPPLoss, ALICEPPModule, inf_norm \n\nloss = ALICEPPLoss(\n    model: ALICEPPModule,\n    num_classes: int,\n    loss_fn: Callable,\n    num_layers: int,\n    max_layer: int = None,\n    loss_last_fn: Callable = None,\n    gold_loss_fn: Callable = None, \n    gold_loss_last_fn: Callable = None, \n    norm_fn: Callable = inf_norm, \n    alpha: float = 1,\n    num_steps: int = 1,\n    step_size: float = 1e-3, \n    epsilon: float = 1e-6,\n    noise_var: float = 1e-5,\n)\n```\n\n## Usage (Classification)\n\n### Extract Model\nThe first thing we have to do is extract the chunk of the model that we want to perturb adversarially. A generic example with Huggingface's RoBERTa for sequence classification is given. \n\n```py \nimport torch.nn as nn \nfrom transformers import AutoModelForSequenceClassification\n\nclass ExtractedRoBERTa(nn.Module):\n\n    def __init__(self):\n        super().__init__()\n        model = AutoModelForSequenceClassification.from_pretrained('roberta-base')\n        self.roberta = model.roberta\n        self.layers = model.roberta.encoder.layer  \n        self.classifier = model.classifier \n        self.attention_mask = None \n        self.num_layers = len(self.layers) - 1 \n\n    def forward(self, hidden, with_hidden_states = False, start_layer = 0):\n        \"\"\" Forwards the hidden value from self.start_layer layer to the logits. \"\"\"\n        hidden_states = [hidden] \n        \n        for layer in self.layers[start_layer:]:\n            hidden = layer(hidden, attention_mask = self.attention_mask)[0]\n            hidden_states += [hidden]\n\n        logits = self.classifier(hidden)\n\n        return (logits, hidden_states) if with_hidden_states else logits \n\n    def get_embeddings(self, input_ids):\n        \"\"\" Computes first embedding layer given inputs_ids \"\"\" \n        return self.roberta.embeddings(input_ids)\n\n    def set_attention_mask(self, attention_mask):\n        \"\"\" Sets the correct mask on all subsequent forward passes \"\"\" \n        self.attention_mask = self.roberta.get_extended_attention_mask(\n            attention_mask, \n            input_shape = attention_mask.shape, \n            device = attention_mask.device\n        ) # (b, 1, 1, s) \n```\nThe function `set_attention_mask` is used to fix the attention mask for all subsequent forward calls, this is necessary if we want to use a mask using any VAT loss. The parameter `start_layer` in the forward function is necessary only if we are using `ALICEPPLoss` since the loss function needs a way to change the start layer internally. \n\n\n### SMART\n\n```py\nimport torch.nn as nn  \nimport torch.nn.functional as F \nfrom vat_pytorch import SMARTLoss, kl_loss, sym_kl_loss\n\nclass SMARTClassificationModel(nn.Module):\n    # b: batch_size, s: sequence_length, d: hidden_size , n: num_labels\n\n    def __init__(self, extracted_model, weight = 1.0):\n        super().__init__()\n        self.model = extracted_model \n        self.weight = weight\n        self.vat_loss = SMARTLoss(model = extracted_model, loss_fn = kl_loss, loss_last_fn = sym_kl_loss)\n\n    def forward(self, input_ids, attention_mask, labels):\n        \"\"\" input_ids: (b, s), attention_mask: (b, s), labels: (b,) \"\"\"\n        # Get input embeddings \n        embeddings = self.model.get_embeddings(input_ids)\n        # Set mask and compute logits \n        self.model.set_attention_mask(attention_mask)\n        logits = self.model(embeddings)\n        # Compute CE loss  \n        ce_loss = F.cross_entropy(logits.view(-1, 2), labels.view(-1))\n        # Compute VAT loss\n        vat_loss = self.vat_loss(embeddings, logits) \n        # Merge losses \n        loss = ce_loss + self.weight * vat_loss\n        return logits, loss\n```\n\n### ALICE\n\n```py\nimport torch.nn as nn  \nimport torch.nn.functional as F \nfrom vat_pytorch import ALICELoss, kl_loss\n\nclass ALICEClassificationModel(nn.Module):\n    # b: batch_size, s: sequence_length, d: hidden_size , n: num_labels\n\n    def __init__(self, extracted_model):\n        super().__init__()\n        self.model = extracted_model \n        self.vat_loss = ALICELoss(model = extracted_model, loss_fn = kl_loss, num_classes = 2)\n\n    def forward(self, input_ids, attention_mask, labels):\n        \"\"\" input_ids: (b, s), attention_mask: (b, s), labels: (b,) \"\"\"\n        # Get input embeddings \n        embeddings = self.model.get_embeddings(input_ids)\n        # Set iteration specific data (e.g. attention mask) \n        self.model.set_attention_mask(attention_mask)\n        # Compute logits \n        logits = self.model(embeddings)\n        # Compute VAT loss\n        loss = self.vat_loss(embeddings, logits, labels) \n        return logits, loss\n```\n\n### ALICE++\n\n```py\nimport torch.nn as nn  \nimport torch.nn.functional as F \nfrom vat_pytorch import ALICEPPLoss, kl_loss\n\nclass ALICEPPClassificationModel(nn.Module):\n    # b: batch_size, s: sequence_length, d: hidden_size , n: num_labels\n\n    def __init__(self, extracted_model):\n        super().__init__()\n        self.model = extracted_model \n        self.vat_loss = ALICEPPLoss(\n            model = extracted_model, \n            loss_fn = kl_loss,\n            num_layers = self.model.num_layers,\n            num_classes = 2 \n        )\n        \n    def forward(self, input_ids, attention_mask, labels):\n        \"\"\" input_ids: (b, s), attention_mask: (b, s), labels: (b,) \"\"\"\n        # Get input embeddings \n        embeddings = self.model.get_embeddings(input_ids)\n        # Set iteration specific data (e.g. attention mask) \n        self.model.set_attention_mask(attention_mask)\n        # Compute logits \n        logits, hidden_states = self.model(embeddings, with_hidden_states = True) \n        # Compute VAT loss \n        loss = self.vat_loss(hidden_states, logits, labels) \n        return logits, loss\n```\n\nNote that `extracted_model` requires a function with the following signature `forward(self, hidden: Tensor, *, start_layer: int) -\u003e Tensor`, the interface `ALICEPPModule` (`from vat_pytorch import ALICEPPModule`) can be used instead of the `nn.Module` class on the extracted model to make sure that the method is present. \n\n\n### Wrapped Model Usage \nAny of the above losses can be used as follows with the extracted model. \n```py \nimport torch \nfrom transformers import AutoTokenizer \n\nextracted_model = ExtractedRoBERTa()\ntokenizer = AutoTokenizer.from_pretrained('roberta-base')\n# Pick one: \nmodel = SMARTClassificationModel(extracted_model)\nmodel = ALICEClassificationModel(extracted_model)\nmodel = ALICEPPClassificationModel(extracted_model)\n# Compute inputs \ntext = [\"This text belongs to class 1...\", \"This text belongs to class 0...\"]\ninputs = tokenizer(text, return_tensors='pt')\nlabels = torch.tensor([1, 0]) \n# Compute logits and loss \nlogits, loss = model(input_ids = inputs['input_ids'], attention_mask = inputs['attention_mask'], labels = labels)\n# To finetune do this for many steps  \nloss.backward() \n```\n\n## Citations\n\n```bibtex\n@inproceedings{Jiang2020SMARTRA,\n  title={SMART: Robust and Efficient Fine-Tuning for Pre-trained Natural Language Models through Principled Regularized Optimization},\n  author={Haoming Jiang and Pengcheng He and Weizhu Chen and Xiaodong Liu and Jianfeng Gao and Tuo Zhao},\n  booktitle={ACL},\n  year={2020}\n}\n```\n\n```bibtex\n@article{Pereira2020AdversarialTF,\n  title={Adversarial Training for Commonsense Inference},\n  author={Lis Kanashiro Pereira and Xiaodong Liu and Fei Cheng and Masayuki Asahara and Ichiro Kobayashi},\n  journal={ArXiv},\n  year={2020},\n  volume={abs/2005.08156}\n}\n```\n\n```bibtex\n@inproceedings{Pereira2021ALICEAT,\n  title={ALICE++: Adversarial Training for Robust and Effective Temporal Reasoning},\n  author={Lis Kanashiro Pereira and Fei Cheng and Masayuki Asahara and Ichiro Kobayashi},\n  booktitle={PACLIC},\n  year={2021}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farchinetai%2Fvat-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Farchinetai%2Fvat-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Farchinetai%2Fvat-pytorch/lists"}