# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from typing import Tuple
import einops
import torch
from jaxtyping import Float
from typeguard import typechecked
@torch.no_grad()
@typechecked
def get_contributions(
parts: torch.Tensor,
whole: torch.Tensor,
distance_norm: int = 1,
) -> torch.Tensor:
"""
Compute contributions of the `parts` vectors into the `whole` vector.
Shapes of the tensors are as follows:
parts: p_1 ... p_k, v_1 ... v_n, d
whole: v_1 ... v_n, d
result: p_1 ... p_k, v_1 ... v_n
Here
* `p_1 ... p_k`: dimensions for enumerating the parts
* `v_1 ... v_n`: dimensions listing the independent cases (batching),
* `d` is the dimension to compute the distances on.
The resulting contributions will be normalized so that
for each v_: sum(over p_ of result(p_, v_)) = 1.
"""
EPS = 1e-5
k = len(parts.shape) - len(whole.shape)
assert k >= 0
assert parts.shape[k:] == whole.shape
bc_whole = whole.expand(parts.shape) # new dims p_1 ... p_k are added to the front
distance = torch.nn.functional.pairwise_distance(parts, bc_whole, p=distance_norm)
whole_norm = torch.norm(whole, p=distance_norm, dim=-1)
distance = (whole_norm - distance).clip(min=EPS)
sum = distance.sum(dim=tuple(range(k)), keepdim=True)
return distance / sum
@torch.no_grad()
@typechecked
def get_contributions_with_one_off_part(
parts: torch.Tensor,
one_off: torch.Tensor,
whole: torch.Tensor,
distance_norm: int = 1,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Same as computing the contributions, but there is one additional part. That's useful
because we always have the residual stream as one of the parts.
See `get_contributions` documentation about `parts` and `whole` dimensions. The
`one_off` should have the same dimensions as `whole`.
Returns a pair consisting of
1. contributions tensor for the `parts`
2. contributions tensor for the `one_off` vector
"""
assert one_off.shape == whole.shape
k = len(parts.shape) - len(whole.shape)
assert k >= 0
# Flatten the p_ dimensions, get contributions for the list, unflatten.
flat = parts.flatten(start_dim=0, end_dim=k - 1)
flat = torch.cat([flat, one_off.unsqueeze(0)])
contributions = get_contributions(flat, whole, distance_norm)
parts_contributions, one_off_contributions = torch.split(
contributions, flat.shape[0] - 1
)
return (
parts_contributions.unflatten(0, parts.shape[0:k]),
one_off_contributions[0],
)
@torch.no_grad()
@typechecked
def get_attention_contributions(
resid_pre: Float[torch.Tensor, "batch pos d_model"],
resid_mid: Float[torch.Tensor, "batch pos d_model"],
decomposed_attn: Float[torch.Tensor, "batch pos key_pos head d_model"],
distance_norm: int = 1,
) -> Tuple[
Float[torch.Tensor, "batch pos key_pos head"],
Float[torch.Tensor, "batch pos"],
]:
"""
Returns a pair of
- a tensor of contributions of each token via each head
- the contribution of the residual stream.
"""
# part dimensions | batch dimensions | vector dimension
# ----------------+------------------+-----------------
# key_pos, head | batch, pos | d_model
parts = einops.rearrange(
decomposed_attn,
"batch pos key_pos head d_model -> key_pos head batch pos d_model",
)
attn_contribution, residual_contribution = get_contributions_with_one_off_part(
parts, resid_pre, resid_mid, distance_norm
)
return (
einops.rearrange(
attn_contribution, "key_pos head batch pos -> batch pos key_pos head"
),
residual_contribution,
)
@torch.no_grad()
@typechecked
def get_mlp_contributions(
resid_mid: Float[torch.Tensor, "batch pos d_model"],
resid_post: Float[torch.Tensor, "batch pos d_model"],
mlp_out: Float[torch.Tensor, "batch pos d_model"],
distance_norm: int = 1,
) -> Tuple[Float[torch.Tensor, "batch pos"], Float[torch.Tensor, "batch pos"]]:
"""
Returns a pair of (mlp, residual) contributions for each sentence and token.
"""
contributions = get_contributions(
torch.stack((mlp_out, resid_mid)), resid_post, distance_norm
)
return contributions[0], contributions[1]
@torch.no_grad()
@typechecked
def get_decomposed_mlp_contributions(
resid_mid: Float[torch.Tensor, "d_model"],
resid_post: Float[torch.Tensor, "d_model"],
decomposed_mlp_out: Float[torch.Tensor, "hidden d_model"],
distance_norm: int = 1,
) -> Tuple[Float[torch.Tensor, "hidden"], float]:
"""
Similar to `get_mlp_contributions`, but it takes the MLP output for each neuron of
the hidden layer and thus computes a contribution per neuron.
Doesn't contain batch and token dimensions for sake of saving memory. But we may
consider adding them.
"""
neuron_contributions, residual_contribution = get_contributions_with_one_off_part(
decomposed_mlp_out, resid_mid, resid_post, distance_norm
)
return neuron_contributions, residual_contribution.item()
@torch.no_grad()
def apply_threshold_and_renormalize(
threshold: float,
c_blocks: torch.Tensor,
c_residual: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Thresholding mechanism used in the original graphs paper. After the threshold is
applied, the remaining contributions are renormalized on order to sum up to 1 for
each representation.
threshold: The threshold.
c_residual: Contribution of the residual stream for each representation. This tensor
should contain 1 element per representation, i.e., its dimensions are all batch
dimensions.
c_blocks: Contributions of the blocks. Could be 1 block per representation, like
ffn, or heads*tokens blocks in case of attention. The shape of `c_residual`
must be a prefix if the shape of this tensor. The remaining dimensions are for
listing the blocks.
"""
block_dims = len(c_blocks.shape)
resid_dims = len(c_residual.shape)
bound_dims = block_dims - resid_dims
assert bound_dims >= 0
assert c_blocks.shape[0:resid_dims] == c_residual.shape
c_blocks = c_blocks * (c_blocks > threshold)
c_residual = c_residual * (c_residual > threshold)
denom = c_residual + c_blocks.sum(dim=tuple(range(resid_dims, block_dims)))
return (
c_blocks / denom.reshape(denom.shape + (1,) * bound_dims),
c_residual / denom,
)