neuron_explainer/activations/derived_scalars/autoencoder.py

"""
This file contains code to compute derived scalars related to autoencoder latents post-hoc
(that is, from pre-existing MLP activations). Typically, the derived scalars consist of the
autoencoder latent activation multiplied by some other quantity.
"""

from typing import Callable

import torch

from neuron_explainer.activations.derived_scalars.derived_scalar_types import DerivedScalarType
from neuron_explainer.activations.derived_scalars.direct_effects import (
    convert_scalar_deriver_to_write_to_final_residual_grad,
)
from neuron_explainer.activations.derived_scalars.raw_activations import (
    convert_scalar_deriver_to_write_norm,
    convert_scalar_deriver_to_write_vector,
    no_op_tensor_calculate_derived_scalar_fn,
)
from neuron_explainer.activations.derived_scalars.reconstituted import make_apply_autoencoder
from neuron_explainer.activations.derived_scalars.reconstituter_class import (
    WriteLatentReconstituter,
)
from neuron_explainer.activations.derived_scalars.scalar_deriver import (
    DerivedScalarSource,
    DstConfig,
    RawScalarSource,
    ScalarDeriver,
)
from neuron_explainer.activations.derived_scalars.utils import detach_and_clone
from neuron_explainer.activations.derived_scalars.write_tensors import (
    get_autoencoder_write_tensor_by_layer_index,
    get_mlp_write_tensor_by_layer_index_with_autoencoder_context,
)
from neuron_explainer.models.autoencoder_context import AutoencoderContext
from neuron_explainer.models.model_component_registry import (
    ActivationLocationType,
    LayerIndex,
    NodeType,
    PassType,
)


def make_autoencoder_latent_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_autoencoder_latent_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        # import here to avoid circular import
        from neuron_explainer.activations.derived_scalars.make_scalar_derivers import (
            make_scalar_deriver,
        )

        layer_indices = dst_config.layer_indices
        if layer_indices is None:
            model_context = dst_config.get_model_context()
            layer_indices = list(range(model_context.n_layers))

        autoencoder_context = dst_config.get_autoencoder_context(node_type)
        assert autoencoder_context is not None

        autoencoder_dst = autoencoder_context.dst
        autoencoder_dst = maybe_convert_autoencoder_dst(autoencoder_dst)
        scalar_deriver = make_scalar_deriver(autoencoder_dst, dst_config)
        apply_autoencoder = make_apply_autoencoder(autoencoder_context)

        def new_tensor_calculate_derived_scalar_fn(
            derived_scalar_tensor: torch.Tensor,
            layer_index: LayerIndex,
            pass_type: PassType,
        ) -> torch.Tensor:
            assert pass_type == PassType.FORWARD
            return apply_autoencoder(derived_scalar_tensor, layer_index)

        output_dst = DerivedScalarType.AUTOENCODER_LATENT.update_from_autoencoder_node_type(
            node_type
        )
        new_scalar_deriver = scalar_deriver.apply_layerwise_transform_fn_to_output(
            layerwise_transform_fn=new_tensor_calculate_derived_scalar_fn,
            pass_type_to_transform=PassType.FORWARD,
            output_dst=output_dst,
        )
        return new_scalar_deriver

    return make_autoencoder_latent_scalar_deriver


def _make_autoencoder_latent_grad_wrt_input_scalar_deriver_helper(
    dst_config: DstConfig,
    output_dst: DerivedScalarType,
    node_type: NodeType | None = None,
) -> ScalarDeriver:
    """Compute the gradient from a particular autoencoder latent, with respect to the
    autoencoder input directions.

    Requires:
    >>> dst_config.layer_indices = [layer_index]
    >>> dst_config.trace_config.tensor_indices = ("All", latent_index)
    """
    # import here to avoid circular import
    from neuron_explainer.activations.derived_scalars.make_scalar_derivers import (
        make_scalar_deriver,
    )

    trace_config = dst_config.trace_config
    assert trace_config is not None
    assert trace_config.node_type in [
        NodeType.AUTOENCODER_LATENT,
        NodeType.MLP_AUTOENCODER_LATENT,
        NodeType.ATTENTION_AUTOENCODER_LATENT,
        NodeType.AUTOENCODER_LATENT_BY_TOKEN_PAIR,
    ]
    assert trace_config.tensor_indices[0] == "All"
    assert isinstance(trace_config.tensor_indices[1], int)
    latent_index_for_grad = trace_config.tensor_indices[1]
    layer_index = trace_config.layer_index
    assert layer_index is not None
    if dst_config.layer_indices is not None:
        assert [layer_index] == dst_config.layer_indices

    autoencoder_context = dst_config.get_autoencoder_context(node_type)
    assert autoencoder_context is not None
    autoencoder_dst = autoencoder_context.dst
    autoencoder_dst = maybe_convert_autoencoder_dst(autoencoder_dst)
    scalar_deriver = make_scalar_deriver(autoencoder_dst, dst_config)
    apply_autoencoder = make_apply_autoencoder(autoencoder_context, use_no_grad=False)

    def new_tensor_calculate_derived_scalar_fn(
        derived_scalar_tensor: torch.Tensor,
    ) -> torch.Tensor:
        derived_scalar_tensor = detach_and_clone(derived_scalar_tensor, requires_grad=True)
        latents = apply_autoencoder(derived_scalar_tensor, layer_index)
        latents[:, latent_index_for_grad].sum(dim=0).backward()  # sum over tokens
        assert derived_scalar_tensor.grad is not None
        return derived_scalar_tensor.grad

    new_scalar_deriver = scalar_deriver.apply_transform_fn_to_output(
        transform_fn=new_tensor_calculate_derived_scalar_fn,
        pass_type_to_transform=PassType.FORWARD,
        output_dst=output_dst,
    )
    return new_scalar_deriver


def make_autoencoder_latent_grad_wrt_residual_input_scalar_deriver(
    dst_config: DstConfig,
    node_type: NodeType | None = None,
) -> ScalarDeriver:
    autoencoder_context = dst_config.get_autoencoder_context(node_type)
    assert autoencoder_context is not None
    assert dst_config.trace_config is not None
    latent_index = dst_config.trace_config.node_index
    assert latent_index is not None
    latent_reconstituter = WriteLatentReconstituter(autoencoder_context)
    return latent_reconstituter.make_gradient_scalar_deriver_for_latent_index(
        latent_index=latent_index,
        dst_config=dst_config,
        output_dst=DerivedScalarType.AUTOENCODER_LATENT_GRAD_WRT_RESIDUAL_INPUT,
    )


def make_autoencoder_latent_grad_wrt_residual_input_scalar_source(
    dst_config: DstConfig,
    node_type: NodeType | None = None,
) -> DerivedScalarSource:
    autoencoder_context = dst_config.get_autoencoder_context(node_type)
    assert autoencoder_context is not None
    assert dst_config.trace_config is not None
    latent_index = dst_config.trace_config.node_index
    assert latent_index is not None
    latent_reconstituter = WriteLatentReconstituter(autoencoder_context)
    return latent_reconstituter.make_gradient_scalar_source_for_latent_index(
        latent_index=latent_index,
        dst_config=dst_config,
        output_dst=DerivedScalarType.AUTOENCODER_LATENT_GRAD_WRT_RESIDUAL_INPUT,
    )


def make_autoencoder_latent_grad_wrt_mlp_post_act_input_scalar_deriver(
    dst_config: DstConfig,
    node_type: NodeType | None = None,
) -> ScalarDeriver:
    """Output shape (n_tokens, n_neurons)"""
    autoencoder_context = dst_config.get_autoencoder_context(node_type)
    assert autoencoder_context is not None
    assert autoencoder_context.dst == DerivedScalarType.MLP_POST_ACT
    return _make_autoencoder_latent_grad_wrt_input_scalar_deriver_helper(
        dst_config,
        output_dst=DerivedScalarType.AUTOENCODER_LATENT_GRAD_WRT_MLP_POST_ACT_INPUT,
        node_type=node_type,
    )


def maybe_convert_autoencoder_dst(autoencoder_dst: DerivedScalarType) -> DerivedScalarType:
    if autoencoder_dst == DerivedScalarType.RESID_DELTA_MLP:
        # TODO: Consider removing this workaround and using RESID_DELTA_MLP directly.
        autoencoder_dst = DerivedScalarType.RESID_DELTA_MLP_FROM_MLP_POST_ACT
    return autoencoder_dst


def make_autoencoder_write_norm_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_autoencoder_write_norm_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        model_context = dst_config.get_model_context()
        dst = DerivedScalarType.AUTOENCODER_WRITE_NORM.update_from_autoencoder_node_type(node_type)

        autoencoder_context = dst_config.get_autoencoder_context(node_type)
        assert autoencoder_context is not None

        write_tensor_by_layer_index = get_autoencoder_write_tensor_by_layer_index(
            autoencoder_context, model_context
        )

        scalar_deriver = make_autoencoder_latent_scalar_deriver_factory(node_type)(dst_config)

        return convert_scalar_deriver_to_write_norm(
            scalar_deriver, write_tensor_by_layer_index, dst
        )

    return make_autoencoder_write_norm_scalar_deriver


def get_autoencoder_alt_from_node_type(node_type: NodeType | None) -> ActivationLocationType:
    """Get the corresponding activation_location_type from a NodeType."""
    return {
        NodeType.AUTOENCODER_LATENT: ActivationLocationType.ONLINE_AUTOENCODER_LATENT,
        NodeType.MLP_AUTOENCODER_LATENT: ActivationLocationType.ONLINE_MLP_AUTOENCODER_LATENT,
        NodeType.ATTENTION_AUTOENCODER_LATENT: ActivationLocationType.ONLINE_ATTENTION_AUTOENCODER_LATENT,
    }[node_type or NodeType.AUTOENCODER_LATENT]


def make_online_autoencoder_latent_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_latent_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        autoencoder_context = dst_config.get_autoencoder_context(node_type)
        assert autoencoder_context is not None
        dst = DerivedScalarType.ONLINE_AUTOENCODER_LATENT.update_from_autoencoder_node_type(
            node_type
        )
        activation_location_type = get_autoencoder_alt_from_node_type(node_type)

        return ScalarDeriver(
            dst=dst,
            dst_config=dst_config,
            sub_scalar_sources=(
                RawScalarSource(
                    activation_location_type=activation_location_type,
                    pass_type=PassType.FORWARD,
                ),
            ),
            tensor_calculate_derived_scalar_fn=no_op_tensor_calculate_derived_scalar_fn,
        )

    return make_online_autoencoder_latent_scalar_deriver


def make_online_autoencoder_write_norm_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_write_norm_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        model_context = dst_config.get_model_context()

        autoencoder_context = dst_config.get_autoencoder_context(node_type)
        assert autoencoder_context is not None
        dst = DerivedScalarType.ONLINE_AUTOENCODER_WRITE_NORM.update_from_autoencoder_node_type(
            node_type
        )

        write_tensor_by_layer_index = get_autoencoder_write_tensor_by_layer_index(
            autoencoder_context, model_context
        )

        scalar_deriver = make_online_autoencoder_latent_scalar_deriver_factory(node_type)(
            dst_config
        )

        return convert_scalar_deriver_to_write_norm(
            scalar_deriver, write_tensor_by_layer_index, dst
        )

    return make_online_autoencoder_write_norm_scalar_deriver


def make_online_autoencoder_latentwise_write_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_latentwise_write_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        model_context = dst_config.get_model_context()

        autoencoder_context = dst_config.get_autoencoder_context(node_type)
        assert autoencoder_context is not None
        dst = DerivedScalarType.ONLINE_AUTOENCODER_WRITE.update_from_autoencoder_node_type(
            node_type
        )

        write_tensor_by_layer_index = get_autoencoder_write_tensor_by_layer_index(
            autoencoder_context, model_context
        )

        scalar_deriver = make_online_autoencoder_latent_scalar_deriver_factory(node_type)(
            dst_config
        )

        return convert_scalar_deriver_to_write_vector(
            scalar_deriver, write_tensor_by_layer_index, dst
        )

    return make_online_autoencoder_latentwise_write_scalar_deriver


def make_online_autoencoder_write_to_final_residual_grad_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_write_to_final_residual_grad_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        scalar_deriver = make_online_autoencoder_latent_scalar_deriver_factory(node_type)(
            dst_config
        )
        dst = DerivedScalarType.ONLINE_AUTOENCODER_WRITE_TO_FINAL_RESIDUAL_GRAD.update_from_autoencoder_node_type(
            node_type
        )
        return convert_scalar_deriver_to_write_to_final_residual_grad(
            scalar_deriver, dst, use_existing_backward_pass_for_final_residual_grad=True
        )

    return make_online_autoencoder_write_to_final_residual_grad_scalar_deriver


def make_online_autoencoder_write_to_final_activation_residual_grad_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_write_to_final_activation_residual_grad_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        scalar_deriver = make_online_autoencoder_latent_scalar_deriver_factory(node_type)(
            dst_config
        )
        dst = DerivedScalarType.ONLINE_AUTOENCODER_WRITE_TO_FINAL_ACTIVATION_RESIDUAL_GRAD.update_from_autoencoder_node_type(
            node_type
        )
        return convert_scalar_deriver_to_write_to_final_residual_grad(
            scalar_deriver, dst, use_existing_backward_pass_for_final_residual_grad=False
        )

    return make_online_autoencoder_write_to_final_activation_residual_grad_scalar_deriver


def make_online_autoencoder_act_times_grad_scalar_deriver_factory(
    node_type: NodeType | None = None,
) -> Callable[[DstConfig], ScalarDeriver]:
    def make_online_autoencoder_act_times_grad_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        act_scalar_deriver = make_online_autoencoder_latent_scalar_deriver_factory(node_type)(
            dst_config
        )
        dst = DerivedScalarType.ONLINE_AUTOENCODER_ACT_TIMES_GRAD.update_from_autoencoder_node_type(
            node_type
        )
        activity_location_type = get_autoencoder_alt_from_node_type(node_type)

        return act_scalar_deriver.apply_layerwise_transform_fn_to_output_and_other_tensor(
            layerwise_transform_fn=lambda act, grad, layer_index, pass_type: act * grad,
            pass_type_to_transform=PassType.FORWARD,  # act
            other_scalar_source=RawScalarSource(
                activation_location_type=activity_location_type,
                pass_type=PassType.BACKWARD,  # grad
            ),
            output_dst=dst,
        )

    return make_online_autoencoder_act_times_grad_scalar_deriver


def make_online_autoencoder_error_scalar_deriver_factory(
    activation_location_type: ActivationLocationType,
) -> Callable[[DstConfig], ScalarDeriver]:
    required_node_type = {
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR: NodeType.MLP_NEURON,
        ActivationLocationType.ONLINE_RESIDUAL_MLP_AUTOENCODER_ERROR: NodeType.RESIDUAL_STREAM_CHANNEL,
        ActivationLocationType.ONLINE_RESIDUAL_ATTENTION_AUTOENCODER_ERROR: NodeType.RESIDUAL_STREAM_CHANNEL,
    }[activation_location_type]

    required_autoencoder_node_type = {
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR: NodeType.MLP_AUTOENCODER_LATENT,
        ActivationLocationType.ONLINE_RESIDUAL_MLP_AUTOENCODER_ERROR: NodeType.MLP_AUTOENCODER_LATENT,
        ActivationLocationType.ONLINE_RESIDUAL_ATTENTION_AUTOENCODER_ERROR: NodeType.ATTENTION_AUTOENCODER_LATENT,
    }[activation_location_type]

    dst = DerivedScalarType.from_activation_location_type(activation_location_type)

    def make_online_autoencoder_error_scalar_deriver(
        dst_config: DstConfig,
    ) -> ScalarDeriver:
        autoencoder_context = dst_config.get_autoencoder_context(required_autoencoder_node_type)
        assert autoencoder_context is not None
        assert autoencoder_context.dst.node_type == required_node_type, (
            autoencoder_context.dst,
            required_node_type,
            activation_location_type,
        )

        return ScalarDeriver(
            dst=dst,
            dst_config=dst_config,
            sub_scalar_sources=(
                RawScalarSource(
                    activation_location_type=activation_location_type,
                    pass_type=PassType.FORWARD,
                ),
            ),
            tensor_calculate_derived_scalar_fn=no_op_tensor_calculate_derived_scalar_fn,
        )

    return make_online_autoencoder_error_scalar_deriver


def make_online_mlp_autoencoder_error_act_times_grad_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    act_scalar_deriver = make_online_autoencoder_error_scalar_deriver_factory(
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR
    )(dst_config)

    return act_scalar_deriver.apply_layerwise_transform_fn_to_output_and_other_tensor(
        layerwise_transform_fn=lambda act, grad, layer_index, pass_type: act * grad,
        pass_type_to_transform=PassType.FORWARD,  # act
        other_scalar_source=RawScalarSource(
            activation_location_type=ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR,
            pass_type=PassType.BACKWARD,  # grad
        ),
        output_dst=DerivedScalarType.ONLINE_MLP_AUTOENCODER_ERROR_ACT_TIMES_GRAD,
    )


def make_online_mlp_autoencoder_error_write_norm_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    model_context = dst_config.get_model_context()

    autoencoder_context = dst_config.get_autoencoder_context(NodeType.MLP_AUTOENCODER_LATENT)
    assert autoencoder_context is not None

    write_tensor_by_layer_index = get_mlp_write_tensor_by_layer_index_with_autoencoder_context(
        autoencoder_context, model_context
    )

    scalar_deriver = make_online_autoencoder_error_scalar_deriver_factory(
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR
    )(dst_config)

    return convert_scalar_deriver_to_write_norm(
        scalar_deriver,
        write_tensor_by_layer_index,
        DerivedScalarType.ONLINE_MLP_AUTOENCODER_ERROR_WRITE_NORM,
    )


def make_online_mlp_autoencoder_error_write_to_final_residual_grad_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    scalar_deriver = make_online_autoencoder_error_scalar_deriver_factory(
        ActivationLocationType.ONLINE_MLP_AUTOENCODER_ERROR
    )(dst_config)
    return convert_scalar_deriver_to_write_to_final_residual_grad(
        scalar_deriver,
        DerivedScalarType.ONLINE_MLP_AUTOENCODER_ERROR_WRITE_TO_FINAL_RESIDUAL_GRAD,
        use_existing_backward_pass_for_final_residual_grad=True,
    )


# helpers for autoencoder gradient


def make_autoencoder_pre_act_encoder_derivative(
    autoencoder_context: AutoencoderContext,
    layer_index: int,
    latent_index: int | None = None,
) -> Callable[[torch.Tensor], torch.Tensor]:
    autoencoder = autoencoder_context.get_autoencoder(layer_index)

    if isinstance(autoencoder.encoder, torch.nn.Linear):
        # if the encoder is linear, then the derivative is just the encoder weight matrix
        encoder = autoencoder.encoder.weight  # shape (n_latents, n_inputs)

        if latent_index is not None:
            encoder = encoder[latent_index : latent_index + 1].clone()
            # ^ need to clone to avoid MPS backend crash

        def pre_act_encoder_derivative(autoencoder_input: torch.Tensor) -> torch.Tensor:
            return autoencoder_input @ encoder.T

        return pre_act_encoder_derivative
    else:
        raise NotImplementedError("Only implemented for linear encoder for now")


def make_autoencoder_activation_fn_derivative(
    autoencoder_context: AutoencoderContext,
    layer_index: int,
) -> Callable[[torch.Tensor], torch.Tensor]:
    autoencoder = autoencoder_context.get_autoencoder(layer_index)

    if _is_relu(autoencoder.activation):
        # if the activation is ReLU, then the derivative is just a step function

        def relu_derivative(post_activations: torch.Tensor) -> torch.Tensor:
            return (post_activations > 0).to(post_activations.dtype)

        return relu_derivative
    else:
        raise NotImplementedError("Only implemented for ReLU activation function for now")


def _is_relu(activation: Callable) -> bool:
    """More robust than isinstance(activation, torch.nn.ReLU), which doesn't always work."""
    if isinstance(activation, torch.nn.ReLU):
        return True
    else:
        test_input = torch.randn(10) * 10 ** torch.randn(10)
        return torch.equal(activation(test_input), torch.nn.ReLU()(test_input))