Add 'classic' rollout

app.py

@@ -8,27 +8,30 @@ sys.path.append("lib")
 import torch
 
 from roberta2 import RobertaForSequenceClassification
+from transformers import AutoTokenizer
+
 from gradient_rollout import GradientRolloutExplainer
+from rollout import RolloutExplainer
 from integrated_gradients import IntegratedGradientsExplainer
-from transformers import AutoModelForSequenceClassification
-from transformers import AutoTokenizer
-from captum.attr import LayerIntegratedGradients
-from captum.attr import visualization
-import util
-import torch
 
-ig_explainer = IntegratedGradientsExplainer()
-gr_explainer = GradientRolloutExplainer()
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+model = RobertaForSequenceClassification.from_pretrained("textattack/roberta-base-SST-2").to(device)
+tokenizer = AutoTokenizer.from_pretrained("textattack/roberta-base-SST-2")
+
+ig_explainer = IntegratedGradientsExplainer(model, tokenizer)
+gr_explainer = GradientRolloutExplainer(model, tokenizer)
+ro_explainer = RolloutExplainer(model, tokenizer)
 
-def run(sent, rollout, ig, ig_baseline):
-    a = gr_explainer(sent, rollout)
-    b = ig_explainer(sent, ig, ig_baseline)
-    return a, b
+def run(sent, gradient, rollout, ig, ig_baseline):
+    a = gr_explainer(sent, gradient)
+    b = ro_explainer(sent, rollout)
+    c = ig_explainer(sent, ig, ig_baseline)
+    return a, b, c
 
 examples = pandas.read_csv("examples.csv").to_numpy().tolist()
 
 with gradio.Blocks(title="Explanations with attention rollout") as iface:
-    util.Markdown(pathlib.Path("description.md"))
+    gradio.Markdown(pathlib.Path("description.md").read_text)
     with gradio.Row(equal_height=True):
         with gradio.Column(scale=4):
             sent = gradio.Textbox(label="Input sentence")
@@ -36,19 +39,54 @@ with gradio.Blocks(title="Explanations with attention rollout") as iface:
             but = gradio.Button("Submit")
     with gradio.Row(equal_height=True):
         with gradio.Column():
-            rollout_layer = gradio.Slider(minimum=0, maximum=12, value=8, step=1, label="Select rollout start layer")
+            rollout_layer = gradio.Slider(
+                    minimum=1,
+                    maximum=12,
+                    value=1,
+                    step=1,
+                    label="Select rollout start layer"
+                )
+        with gradio.Column():
+            gradient_layer = gradio.Slider(
+                    minimum=1,
+                    maximum=12,
+                    value=8,
+                    step=1,
+                    label="Select gradient rollout start layer"
+                )
+        with gradio.Column():
+            ig_layer = gradio.Slider(
+                    minimum=0,
+                    maximum=12,
+                    value=0,
+                    step=1,
+                    label="Select IG layer"
+                )
+            ig_baseline = gradio.Dropdown(
+                    label="Baseline token",
+                    choices=['Unknown', 'Padding'], value="Unknown"
+                )
+    with gradio.Row(equal_height=True):
+        with gradio.Column():
+            gradio.Markdown("### Attention Rollout")
             rollout_result = gradio.HTML()
         with gradio.Column():
-            ig_layer = gradio.Slider(minimum=0, maximum=12, value=0, step=1, label="Select IG layer")
-            ig_baseline = gradio.Dropdown(label="Baseline token", choices=['Unknown', 'Padding'], value="Unknown")
+            gradio.Markdown("### Gradient-weighted Attention Rollout")
+            gradient_result = gradio.HTML()
+        with gradio.Column():
+            gradio.Markdown("### Layer-Integrated Gradients")
             ig_result = gradio.HTML()
     gradio.Examples(examples, [sent])
     with gradio.Accordion("Some more details"):
-        util.Markdown(pathlib.Path("notice.md"))
+        gradio.Markdown(pathlib.Path("notice.md").read_text)
 
-    rollout_layer.change(gr_explainer, [sent, rollout_layer], rollout_result)
+    gradient_layer.change(gr_explainer, [sent, gradient_layer], gradient_result)
+    rollout_layer.change(ro_explainer, [sent, rollout_layer], rollout_result)
     ig_layer.change(ig_explainer, [sent, ig_layer, ig_baseline], ig_result)
-    but.click(run, [sent, rollout_layer, ig_layer, ig_baseline], [rollout_result, ig_result])
+    but.click(run,
+            inputs=[sent, gradient_layer, rollout_layer, ig_layer, ig_baseline],
+            outputs=[gradient_result, rollout_result, ig_result]
+        )
 
 
 iface.launch()

lib/ExplanationGenerator.py

@@ -25,8 +25,8 @@ class Generator:
         self.key = key
         self.model.eval()
 
-    def forward(self, input_ids, attention_mask):
-        return self.model(input_ids, attention_mask)
+    def tokens_from_ids(self, ids):
+        return list(map(lambda s: s[1:] if s[0] == "Ġ" else s, self.tokenizer.convert_ids_to_tokens(ids)))
 
     def _calculate_gradients(self, output, index, do_relprop=True):
         if index == None:
@@ -72,7 +72,6 @@ class Generator:
         rollout[:, 0, 0] = rollout[:, 0].min()
         return rollout[:, 0]
 
-
     def generate_LRP_last_layer(self, input_ids, attention_mask,
                      index=None):
         output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]
@@ -117,7 +116,7 @@ class Generator:
             all_layer_attentions.append(avg_heads)
         rollout = compute_rollout_attention(all_layer_attentions, start_layer=start_layer)
         rollout[:, 0, 0] = 0
-        return rollout[:, 0]
+        return output, rollout[:, 0]
 
     def generate_attn_gradcam(self, input_ids, attention_mask, index=None):
         output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]
@@ -148,12 +147,14 @@ class Generator:
             return torch.matmul(cam_ss, R_ss)
 
         output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]
-        blocks = _get_module_from_name(self.model, self.key)
+
+        self._calculate_gradients(output, index, do_relprop=False)
 
         num_tokens = input_ids.size(-1)
         R = torch.eye(num_tokens).expand(output.size(0), -1, -1).clone().to(output.device)
 
-        for i, blk in enumerate(model.roberta.encoder.layer):
+        blocks = _get_module_from_name(self.model, self.key)
+        for i, blk in enumerate(blocks):
             if i < start_layer:
                 continue
             grad = blk.attention.self.get_attn_gradients().detach()
@@ -161,5 +162,7 @@ class Generator:
             cam = avg_heads(cam, grad)
             joint = apply_self_attention_rules(R, cam)
             R += joint
-        return R[:, 0, 1:-1]
+        # 0 because we look at the influence *on* the CLS token
+        # 1:-1 because we don't want the influence *from* the CLS/SEP tokens
+        return output, R[:, 0, 1:-1]
 

lib/gradient_rollout.py

@@ -4,68 +4,22 @@ from captum.attr import visualization
 
 from roberta2 import RobertaForSequenceClassification
 from util import visualize_text, PyTMinMaxScalerVectorized
+from ExplanationGenerator import Generator
 
 classifications = ["NEGATIVE", "POSITIVE"]
 
-class GradientRolloutExplainer:
-    def __init__(self):
-        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-        self.model = RobertaForSequenceClassification.from_pretrained("textattack/roberta-base-SST-2").to(self.device)
-        self.model.eval()
-        self.tokenizer = AutoTokenizer.from_pretrained("textattack/roberta-base-SST-2")
-
-    def tokens_from_ids(self, ids):
-        return list(map(lambda s: s[1:] if s[0] == "Ġ" else s, self.tokenizer.convert_ids_to_tokens(ids)))
-
-    def run_attribution_model(self, input_ids, attention_mask, index=None, start_layer=0):
-        def avg_heads(cam, grad):
-            cam = (grad * cam).clamp(min=0).mean(dim=-3)
-            # set negative values to 0, then average
-            #    cam = cam.clamp(min=0).mean(dim=0)
-            return cam
-
-        def apply_self_attention_rules(R_ss, cam_ss):
-            R_ss_addition = torch.matmul(cam_ss, R_ss)
-            return R_ss_addition
-
-        output = self.model(input_ids=input_ids, attention_mask=attention_mask)[0]
-        if index == None:
-            # index = np.expand_dims(np.arange(input_ids.shape[1])
-            # by default explain the class with the highest score
-            index = output.argmax(axis=-1).detach().cpu().numpy()
-
-        # create a one-hot vector selecting class we want explanations for
-        one_hot = (
-            torch.nn.functional.one_hot(
-                torch.tensor(index, dtype=torch.int64), num_classes=output.size(-1)
-            )
-            .to(torch.float)
-            .requires_grad_(True)
-        ).to(self.device)
-        one_hot = torch.sum(one_hot * output)
-        self.model.zero_grad()
-        # create the gradients for the class we're interested in
-        one_hot.backward(retain_graph=True)
-
-        num_tokens = self.model.roberta.encoder.layer[0].attention.self.get_attn().shape[-1]
-        R = torch.eye(num_tokens).expand(output.size(0), -1, -1).clone().to(self.device)
-
-        for i, blk in enumerate(self.model.roberta.encoder.layer):
-            if i < start_layer:
-                continue
-            grad = blk.attention.self.get_attn_gradients()
-            cam = blk.attention.self.get_attn()
-            cam = avg_heads(cam, grad)
-            joint = apply_self_attention_rules(R, cam)
-            R += joint
-        return output, R[:, 0, 1:-1]
+class GradientRolloutExplainer(Generator):
+    def __init__(self, model, tokenizer):
+        super().__init__(model, key="roberta.encoder.layer")
+        self.device = model.device
+        self.tokenizer = tokenizer
 
     def build_visualization(self, input_ids, attention_mask, index=None, start_layer=8):
         # generate an explanation for the input
         vis_data_records = []
 
         for index in range(2):
-            output, expl = self.run_attribution_model(
+            output, expl = self.generate_rollout_attn_gradcam(
                 input_ids, attention_mask, index=index, start_layer=start_layer
             )
             # normalize scores

lib/integrated_gradients.py

@@ -6,15 +6,17 @@ from transformers import AutoTokenizer
 from captum.attr import LayerIntegratedGradients
 from captum.attr import visualization
 
+from roberta2 import RobertaForSequenceClassification
+from ExplanationGenerator import Generator
 from util import visualize_text
 
 classifications = ["NEGATIVE", "POSITIVE"]
 
 class IntegratedGradientsExplainer:
-    def __init__(self):
-        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-        self.model = AutoModelForSequenceClassification.from_pretrained("textattack/roberta-base-SST-2").to(self.device)
-        self.tokenizer = AutoTokenizer.from_pretrained("textattack/roberta-base-SST-2")
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.device = model.device
+        self.tokenizer = tokenizer
         self.baseline_map = {
                 'Unknown': self.tokenizer.unk_token_id,
                 'Padding': self.tokenizer.pad_token_id,

lib/rollout.py

@@ -0,0 +1,67 @@
+import torch
+from transformers import AutoTokenizer
+from captum.attr import visualization
+
+from roberta2 import RobertaForSequenceClassification
+from ExplanationGenerator import Generator
+from util import visualize_text, PyTMinMaxScalerVectorized
+
+classifications = ["NEGATIVE", "POSITIVE"]
+
+class RolloutExplainer(Generator):
+    def __init__(self, model, tokenizer):
+        super().__init__(model, key="roberta.encoder.layer")
+        self.device = model.device
+        self.tokenizer = tokenizer
+
+    def build_visualization(self, input_ids, attention_mask, start_layer=8):
+        # generate an explanation for the input
+        vis_data_records = []
+
+        output, expl = self.generate_rollout(
+            input_ids, attention_mask, start_layer=start_layer
+        )
+        # normalize scores
+        scaler = PyTMinMaxScalerVectorized()
+
+        norm = scaler(expl)
+        # get the model classification
+        output = torch.nn.functional.softmax(output, dim=-1)
+
+        for record in range(input_ids.size(0)):
+            classification = output[record].argmax(dim=-1).item()
+            class_name = classifications[classification]
+            nrm = norm[record]
+
+            # if the classification is negative, higher explanation scores are more negative
+            # flip for visualization
+            if class_name == "NEGATIVE":
+                nrm *= -1
+            tokens = self.tokens_from_ids(input_ids[record].flatten())[
+                1 : 0 - ((attention_mask[record] == 0).sum().item() + 1)
+            ]
+            vis_data_records.append(
+                visualization.VisualizationDataRecord(
+                    nrm,
+                    output[record][classification],
+                    classification,
+                    classification,
+                    classification,
+                    1,
+                    tokens,
+                    1,
+                )
+            )
+        return visualize_text(vis_data_records)
+
+    def __call__(self, input_text, start_layer=8):
+        if start_layer > 0:
+            start_layer -= 1
+
+        text_batch = [input_text]
+        encoding = self.tokenizer(text_batch, return_tensors="pt")
+        input_ids = encoding["input_ids"].to(self.device)
+        attention_mask = encoding["attention_mask"].to(self.device)
+
+        return self.build_visualization(input_ids, attention_mask, start_layer=int(start_layer))
+