# Copyright 2020 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Gradient-based attribution."""

from typing import cast, Optional

from absl import logging

from lit_nlp.api import components as lit_components

from lit_nlp.api import dataset as lit_dataset

from lit_nlp.api import dtypes

from lit_nlp.api import model as lit_model

from lit_nlp.api import types

from lit_nlp.components.citrus import utils as citrus_utils

from lit_nlp.lib import utils

import numpy as np

JsonDict = types.JsonDict

Spec = types.Spec

CLASS_KEY = 'Class to explain'

NORMALIZATION_KEY = 'Normalize'

INTERPOLATION_KEY = 'Interpolation steps'

class GradientNorm(lit_components.Interpreter):

"""Salience map from gradient L2 norm."""

def find_fields(self, output_spec: Spec) -> list[str]:

# Find TokenGradients fields

supported_fields: list[str] = []

# Check that these are aligned to Tokens fields

for f in utils.find_spec_keys(output_spec, types.TokenGradients):

tokens_field = cast(types.TokenGradients, output_spec[f]).align

is_valid_tokens = (

tokens_field is not None and tokens_field in output_spec and

isinstance(output_spec[tokens_field], types.Tokens))

if not is_valid_tokens:

logging.info('Skipping %s. Invalid tokens field, %s', str(f),

str(tokens_field))

continue

supported_fields.append(f)

return supported_fields

def _interpret(self, grads: np.ndarray, tokens: np.ndarray):

assert grads.shape[0] == len(tokens)

# Norm of dy/d(embs)

grad_norm = np.linalg.norm(grads, axis=1)

grad_norm /= np.sum(grad_norm)

# <float32>[num_tokens]

return grad_norm

def run(self,

inputs: list[JsonDict],

model: lit_model.Model,

dataset: lit_dataset.Dataset,

model_outputs: Optional[list[JsonDict]] = None,

config: Optional[JsonDict] = None) -> Optional[list[JsonDict]]:

"""Run this component, given a model and input(s)."""

del dataset, config

# Find gradient fields to interpret

output_spec = model.output_spec()

grad_fields = self.find_fields(output_spec)

logging.info('Found fields for gradient attribution: %s', str(grad_fields))

if len(grad_fields) == 0: # pylint: disable=g-explicit-length-test

return None

# Run model, if needed.

if model_outputs is None:

model_outputs = list(model.predict(inputs))

assert len(model_outputs) == len(inputs)

all_results: list[JsonDict] = []

for o in model_outputs:

result: dict[str, dtypes.TokenSalience] = {}

for grad_field in grad_fields:

token_field = cast(types.TokenGradients, output_spec[grad_field]).align

tokens = o[token_field]

scores = self._interpret(o[grad_field], tokens)

result[grad_field] = dtypes.TokenSalience(tokens, scores)

all_results.append(result)

return all_results

def is_compatible(self, model: lit_model.Model,

dataset: lit_dataset.Dataset) -> bool:

del dataset # Unused by Grad L2 Norm

return bool(self.find_fields(model.output_spec()))

def meta_spec(self) -> types.Spec:

return {'saliency': types.TokenSalience(autorun=True, signed=False)}

class GradientDotInput(lit_components.Interpreter):

"""Salience map using the values of gradient * input as attribution."""

def find_fields(self, output_spec: Spec) -> list[str]:

# Find and check that TokenGradients fields are aligned to Tokens fields

aligned_fields = []

for f in utils.find_spec_keys(output_spec, types.TokenGradients):

field_spec = cast(types.TokenGradients, output_spec[f])

tokens_field = field_spec.align

is_valid_tokens = (

tokens_field is not None and tokens_field in output_spec and

isinstance(output_spec[tokens_field], types.Tokens))

if not is_valid_tokens:

logging.info('Skipping %s. Invalid tokens field, %s', str(f),

str(tokens_field))

continue

embeddings_field = field_spec.grad_for

is_valid_embeddings = (

embeddings_field is not None and embeddings_field in output_spec and

isinstance(output_spec[embeddings_field], types.TokenEmbeddings))

if not is_valid_embeddings:

logging.info('Skipping %s. Invalid emebeddings field, %s.', str(f),

str(tokens_field))

continue

aligned_fields.append(f)

return aligned_fields

def _interpret(self, grads: np.ndarray, embs: np.ndarray):

assert grads.shape == embs.shape

# dot product of gradients and embeddings

# <float32>[num_tokens]

grad_dot_input = np.sum(grads * embs, axis=-1)

scores = citrus_utils.normalize_scores(grad_dot_input)

return scores

def run(self,

inputs: list[JsonDict],

model: lit_model.Model,

dataset: lit_dataset.Dataset,

model_outputs: Optional[list[JsonDict]] = None,

config: Optional[JsonDict] = None) -> Optional[list[JsonDict]]:

"""Run this component, given a model and input(s)."""

# Find gradient fields to interpret

output_spec = model.output_spec()

grad_fields = self.find_fields(output_spec)

logging.info('Found fields for gradient attribution: %s', str(grad_fields))

if len(grad_fields) == 0: # pylint: disable=g-explicit-length-test

return None

# Run model, if needed.

if model_outputs is None:

model_outputs = list(model.predict(inputs))

assert len(model_outputs) == len(inputs)

all_results: list[JsonDict] = []

for o in model_outputs:

result: dict[str, dtypes.TokenSalience] = {}

for grad_field in grad_fields:

embeddings_field = cast(types.TokenGradients,

output_spec[grad_field]).grad_for

scores = self._interpret(o[grad_field], o[embeddings_field])

token_field = cast(types.TokenGradients, output_spec[grad_field]).align

tokens = o[token_field]

result[grad_field] = dtypes.TokenSalience(tokens, scores)

all_results.append(result)

return all_results

def is_compatible(self, model: lit_model.Model,

dataset: lit_dataset.Dataset) -> bool:

del dataset # Unused by Grad*Input

return bool(self.find_fields(model.output_spec()))

def meta_spec(self) -> types.Spec:

return {'saliency': types.TokenSalience(autorun=True, signed=True)}

class IntegratedGradients(lit_components.Interpreter):

"""Salience map from Integrated Gradients.

Integrated Gradients is an attribution method originally proposed in

Sundararajan et al. (https://arxiv.org/abs/1703.01365), which attributes an

importance value for each input feature based on the gradients of the model

output with respect to the input. The feature attribution values are

calculated by taking the integral of gradients along a straight path from a

baseline to the input being analyzed. The original implementation can be

found at: https://github.com/ankurtaly/Integrated-Gradients/blob/master/

BertModel/bert_model_utils.py

This component requires that the following fields in the model spec. Field

names like `embs` are placeholders; you can call them whatever you like,

and as with other LIT components multiple segments are supported.

Output:

- TokenEmbeddings (`embs`) to return the input embeddings

- TokenGradients (`grads`) to return gradients w.r.t. `embs`

- A label field (`target`) to return the label that `grads`

was computed for. This is usually a CategoryLabel, but can be anything

since it will just be fed back into the model.

Input

- TokenEmbeddings (`embs`) to accept the modified input embeddings

- A label field to (`target`) to pin the gradient target to the same

label for all integral steps, since the argmax prediction may change.

"""

def __init__(self,

autorun: bool = False,

class_key: str = '',

interpolation_steps: int = 30,

normalize: bool = True):

"""Cretaes an IntegratedGradients interpreter.

Args:

autorun: Determines if this intepreter should run automatically.

class_key: The class to explain.

interpolation_steps: The number of steps to interpolate.

normalize: Flag to enable/disable normalization.

"""

self._autorun: bool = autorun

self._class_key: str = class_key

self._interpolation_steps: int = interpolation_steps

self._normalize: bool = normalize

def find_fields(self, input_spec: Spec, output_spec: Spec) -> list[str]:

# Find and check that TokenGradients fields are aligned to Tokens fields

aligned_fields = []

for f in utils.find_spec_keys(output_spec, types.TokenGradients):

field_spec = cast(types.TokenGradients, output_spec[f])

tokens_field = field_spec.align

embeddings_field = field_spec.grad_for

grad_key = field_spec.grad_target_field_key

if not isinstance(output_spec.get(tokens_field), types.Tokens):

logging.info('Skipping %s. Invalid tokens field, %s.', str(f),

str(tokens_field))

continue

is_embs_valid = (

isinstance(input_spec.get(embeddings_field),

types.TokenEmbeddings) and

isinstance(output_spec.get(embeddings_field), types.TokenEmbeddings))

if not is_embs_valid:

logging.info('Skipping %s. Invalid embeddings field, %s.', str(f),

str(tokens_field))

continue

is_grad_cls_valid = grad_key in input_spec

if not is_grad_cls_valid:

logging.info('Skipping %s. Invalid gradient class field, %s.', str(f),

str(tokens_field))

continue

aligned_fields.append(f)

return aligned_fields

def get_interpolated_inputs(self, baseline: np.ndarray, target: np.ndarray,

num_steps: int) -> np.ndarray:

"""Gets num_step linearly interpolated inputs from baseline to target."""

if num_steps <= 0: return np.array([])

if num_steps == 1: return np.array([baseline, target])

delta = target - baseline # <float32>[num_tokens, emb_size]

# Creates scale values array of shape [num_steps, num_tokens, emb_dim],

# where the values in scales[i] are the ith step from np.linspace.

# <float32>[num_steps, 1, 1]

scales = np.linspace(0, 1, num_steps + 1,

dtype=np.float32)[:, np.newaxis, np.newaxis]

shape = (num_steps + 1,) + delta.shape

# <float32>[num_steps, num_tokens, emb_size]

deltas = scales * np.broadcast_to(delta, shape)

interpolated_inputs = baseline + deltas

return interpolated_inputs # <float32>[num_steps, num_tokens, emb_size]

def estimate_integral(self, path_gradients: np.ndarray) -> np.ndarray:

"""Estimates the integral of the path_gradients using trapezoid rule."""

path_gradients = (path_gradients[:-1] + path_gradients[1:]) / 2

# There are num_steps elements in the path_gradients. Summing num_steps - 1

# terms and dividing by num_steps - 1 is equivalent to taking

# the average.

return np.average(path_gradients, axis=0)

def get_baseline(self, embeddings: np.ndarray) -> np.ndarray:

"""Returns baseline embeddings to use in Integrated Gradients."""

# Replaces embeddings in the original input with the zero embedding, or

# with the specified token embedding.

baseline = np.zeros_like(embeddings)

# TODO(ellenj): Add option to use a token's embedding as the baseline.

return baseline

def get_salience_result(self, model_input: JsonDict, model: lit_model.Model,

interpolation_steps: int, normalize: bool,

class_to_explain: str, model_output: JsonDict,

grad_fields: list[str]):

result = {}

output_spec = model.output_spec()

# We ensure that the embedding and gradient class fields are present in the

# model's input spec in find_fields().

embeddings_fields = [

cast(types.TokenGradients,

output_spec[grad_field]).grad_for for grad_field in grad_fields]

# The gradient class input is used to specify the target class of the

# gradient calculation (if unspecified, this option defaults to the argmax,

# which could flip between interpolated inputs).

# If class_to_explain is emptystring, then explain the argmax class.

grad_class_key = cast(types.TokenGradients,

output_spec[grad_fields[0]]).grad_target_field_key

if class_to_explain == '': # pylint: disable=g-explicit-bool-comparison

if grad_class_key in model_output:

grad_class = model_output[grad_class_key]

else:

# If this is not present, should error because we can't infer

# what class to use.

grad_class = model_input[grad_class_key]

else:

grad_class = class_to_explain

interpolated_inputs = {}

all_embeddings = []

all_baselines = []

for embed_field in embeddings_fields:

# <float32>[num_tokens, emb_size]

embeddings = np.array(model_output[embed_field])

all_embeddings.append(embeddings)

# Starts with baseline of zeros. <float32>[num_tokens, emb_size]

baseline = self.get_baseline(embeddings)

all_baselines.append(baseline)

# Get interpolated inputs from baseline to original embedding.

# <float32>[interpolation_steps, num_tokens, emb_size]

interpolated_inputs[embed_field] = self.get_interpolated_inputs(

baseline, embeddings, interpolation_steps)

# Create model inputs and populate embedding field(s).

inputs_with_embeds = []

for i in range(interpolation_steps):

# Interpolates embeddings for all inputs simultaneously.

# Each entry is <float32>[num_tokens, emb_size]

updates = {k: interpolated_inputs[k][i] for k in embeddings_fields}

updates[grad_class_key] = grad_class

input_copy = utils.make_modified_input(model_input, updates, 'IG')

inputs_with_embeds.append(input_copy)

embed_outputs = model.predict(inputs_with_embeds)

# Create list with concatenated gradients for each interpolate input.

gradients = []

for o in embed_outputs:

# <float32>[total_num_tokens, emb_size]

interp_gradients = np.concatenate([o[field] for field in grad_fields])

gradients.append(interp_gradients)

# <float32>[interpolation_steps, total_num_tokens, emb_size]

path_gradients = np.stack(gradients, axis=0)

# Calculate integral

# <float32>[total_num_tokens, emb_size]

integral = self.estimate_integral(path_gradients)

# <float32>[total_num_tokens, emb_size]

concat_embeddings = np.concatenate(all_embeddings)

# <float32>[total_num_tokens, emb_size]

concat_baseline = np.concatenate(all_baselines)

# <float32>[total_num_tokens, emb_size]

integrated_gradients = integral * (np.array(concat_embeddings) -

np.array(concat_baseline))

# Dot product of integral values and (embeddings - baseline).

# <float32>[total_num_tokens]

attributions = np.sum(integrated_gradients, axis=-1)

# <float32>[total_num_tokens]

scores = citrus_utils.normalize_scores(

attributions) if normalize else attributions

for grad_field in grad_fields:

# Format as salience map result.

token_field = cast(types.TokenGradients, output_spec[grad_field]).align

tokens = model_output[token_field]

# Only use the scores that correspond to the tokens in this grad_field.

# The gradients for all input embeddings were concatenated in the order

# of the grad fields, so they can be sliced out in the same order.

sliced_scores = scores[:len(tokens)] # <float32>[num_tokens in field]

scores = scores[len(tokens):] # <float32>[num_remaining_tokens]

assert len(tokens) == len(sliced_scores)

result[grad_field] = dtypes.TokenSalience(tokens, sliced_scores)

return result

def run(self,

inputs: list[JsonDict],

model: lit_model.Model,

dataset: lit_dataset.Dataset,

model_outputs: Optional[list[JsonDict]] = None,

config: Optional[JsonDict] = None) -> Optional[list[JsonDict]]:

"""Run this component, given a model and input(s)."""

config = config or {}

class_to_explain = config.get(CLASS_KEY, self._class_key)

try:

interpolation_steps = int(config.get(INTERPOLATION_KEY,

self._interpolation_steps))

except ValueError as parse_error:

raise RuntimeError(

'Failed to parse interpolation steps'

f'from "{config[INTERPOLATION_KEY]}".') from parse_error

normalization = config.get(NORMALIZATION_KEY, self._normalize)

# Find gradient fields to interpret

input_spec = model.input_spec()

output_spec = model.output_spec()

grad_fields = self.find_fields(input_spec, output_spec)

logging.info('Found fields for integrated gradients: %s', str(grad_fields))

if len(grad_fields) == 0: # pylint: disable=g-explicit-length-test

return None

# Run model, if needed.

if model_outputs is None:

model_outputs = list(model.predict(inputs))

all_results = []

for model_output, model_input in zip(model_outputs, inputs):

result = self.get_salience_result(model_input, model, interpolation_steps,

normalization, class_to_explain,

model_output, grad_fields)

all_results.append(result)

return all_results

def is_compatible(self, model: lit_model.Model,

dataset: lit_dataset.Dataset) -> bool:

del dataset # Unused by IG

return bool(self.find_fields(model.input_spec(), model.output_spec()))

def config_spec(self) -> types.Spec:

return {

CLASS_KEY: types.TextSegment(default=self._class_key),

NORMALIZATION_KEY: types.Boolean(default=self._normalize),

INTERPOLATION_KEY:

types.Scalar(

min_val=5,

max_val=100,

default=self._interpolation_steps,

step=1)

}

def meta_spec(self) -> types.Spec:

return {'saliency': types.TokenSalience(autorun=self._autorun, signed=True)}