explainer_context.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cmath>
#include <cstddef>
#include <cstdlib>
#include <optional>
#include <stdexcept>
#include <string>
#include <typeinfo>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "pulsatrix/activation_snapshot.hpp"
#include "pulsatrix/assert.hpp"
#include "pulsatrix/autograd.hpp"
#include "pulsatrix/circuit_graph.hpp"
#include "pulsatrix/computation_graph.hpp"
#include "pulsatrix/lrp_rule_config.hpp"
#include "pulsatrix/module.hpp"
#include "pulsatrix/multihead_attention_module.hpp"
#include "pulsatrix/tensor.hpp"
 
#if defined(__GNUG__)
#include <cxxabi.h>
#endif
 
namespace pulsatrix {
 
namespace detail {
inline std::string module_type_name(const Module& module) {
    const char* raw = typeid(module).name();
#if defined(__GNUG__)
    int status = 0;
    char* demangled = abi::__cxa_demangle(raw, nullptr, nullptr, &status);
    std::string name = (status == 0 && demangled != nullptr) ? demangled : raw;
    std::free(demangled);
    return name;
#else
    return raw;
#endif
}
}  // namespace detail
 
class ExplainerContext {
public:
    explicit ExplainerContext(std::vector<Module*> modules) : modules_(std::move(modules)) {
        if (modules_.empty()) {
            throw std::invalid_argument("ExplainerContext: modules must not be empty");
        }
        for (Module* m : modules_) {
            if (m == nullptr) {
                throw std::invalid_argument("ExplainerContext: modules must not contain a null pointer");
            }
        }
    }
 
    Tensor forward_pass(const Tensor& input) {
        last_forward_was_patched_ = false;
        graph_ = ComputationGraph{};
        autograd_ = Autograd{};
        activations_.clear();
 
        input_node_ = graph_.add_node(OpType::Elementwise, input.shape(), "input");
        activations_.emplace(input_node_, Tensor(input));
 
        Tensor current = input;
        NodeId current_node = input_node_;
        for (Module* module : modules_) {
            auto [output, node_id] = module->forward_traced(current, current_node, graph_, autograd_);
            activations_.emplace(node_id, Tensor(output));
            current = std::move(output);
            current_node = node_id;
        }
        output_node_ = current_node;
        return current;
    }
 
    Tensor forward_pass_with_patch(const Tensor& input, NodeId patch_node_id, const Tensor& patch_value) {
        last_forward_was_patched_ = true;
        if (patch_node_id > modules_.size()) {
            throw std::invalid_argument("ExplainerContext::forward_pass_with_patch: patch_node_id out of range");
        }
 
        ComputationGraph graph;
        Autograd autograd;
        std::unordered_map<NodeId, Tensor> activations;
 
        NodeId input_node = graph.add_node(OpType::Elementwise, input.shape(), "input");
 
        Tensor current = input;
        NodeId current_node = input_node;
        if (patch_node_id == input_node) {
            current = check_patch_shape(current, patch_value);
        }
        activations.emplace(input_node, Tensor(current));
 
        for (Module* module : modules_) {
            auto [output, node_id] = module->forward_traced(current, current_node, graph, autograd);
            if (node_id == patch_node_id) {
                output = check_patch_shape(output, patch_value);
            }
            activations.emplace(node_id, Tensor(output));
            current = std::move(output);
            current_node = node_id;
        }
 
        graph_ = std::move(graph);
        autograd_ = std::move(autograd);
        activations_ = std::move(activations);
        input_node_ = input_node;
        output_node_ = current_node;
        return current;
    }
 
    Tensor backward_pass(const Tensor& output_grad) {
        if (last_forward_was_patched_) {
            throw std::logic_error(
                "ExplainerContext::backward_pass: the most recent forward pass was "
                "forward_pass_with_patch(); a patched activation is a constant substitution, not a "
                "differentiable function of the input, so gradients through it would be silently wrong. "
                "Run an unpatched forward_pass() before backward_pass().");
        }
        autograd_.backward(graph_, output_node_, output_grad);
        return Tensor(autograd_.gradient(input_node_));
    }
 
    Tensor relevance_pass(const Tensor& output_relevance, const LRPRuleConfig& config) {
        return relevance_pass(output_relevance, std::vector<LRPRuleConfig>(modules_.size(), config));
    }
 
    Tensor relevance_pass(const Tensor& output_relevance, const std::vector<LRPRuleConfig>& configs) {
        if (last_forward_was_patched_) {
            throw std::logic_error(
                "ExplainerContext::relevance_pass: the most recent forward pass was "
                "forward_pass_with_patch(); relevance through a patched activation would describe a "
                "computation the input did not produce. Run an unpatched forward_pass() first.");
        }
        if (configs.size() != modules_.size()) {
            throw std::invalid_argument("ExplainerContext::relevance_pass: need exactly one LRPRuleConfig per module");
        }
        for (size_t i = 0; i < modules_.size(); ++i) {
            if (!modules_[i]->supports_lrp_rule(configs[i].rule)) {
                throw std::invalid_argument("ExplainerContext::relevance_pass: module " + std::to_string(i) + " (" +
                                            detail::module_type_name(*modules_[i]) + ") does not implement the " +
                                            lrp_rule_name(configs[i].rule) + " LRP rule");
            }
        }
        Tensor relevance = output_relevance;
        for (size_t i = modules_.size(); i-- > 0;) {
            relevance = modules_[i]->propagate_relevance(relevance, configs[i]);
        }
        return relevance;
    }
 
    [[nodiscard]] const std::vector<Module*>& modules() const { return modules_; }
 
    [[nodiscard]] const ComputationGraph& graph() const { return graph_; }
 
    [[nodiscard]] const Tensor& activation(NodeId id) const {
        auto it = activations_.find(id);
        PULSATRIX_ASSERT(it != activations_.end());
        return it->second;
    }
 
    [[nodiscard]] const Tensor& gradient(NodeId id) const { return autograd_.gradient(id); }
 
    [[nodiscard]] ActivationSnapshot activation_snapshot() const {
        std::vector<NodeId> node_ids = graph_.topological_order();
 
        std::unordered_map<NodeId, Tensor> activations;
        std::unordered_map<NodeId, ActivationSnapshot::NodeMetadata> metadata;
        for (NodeId id : node_ids) {
            auto it = activations_.find(id);
            PULSATRIX_ASSERT(it != activations_.end());
            activations.emplace(id, Tensor(it->second));
            const Node& node = graph_.node(id);
            metadata.emplace(id, ActivationSnapshot::NodeMetadata{node.op_type(), node.label()});
        }
 
        return ActivationSnapshot(std::move(activations), std::move(node_ids), std::move(metadata));
    }
 
    [[nodiscard]] Tensor logit_lens(NodeId node_id) const {
        auto it = activations_.find(node_id);
        if (it == activations_.end()) {
            throw std::invalid_argument(
                "ExplainerContext::logit_lens: node_id has no cached activation from the most recent "
                "forward pass");
        }
 
        const NodeId head_input_node = static_cast<NodeId>(modules_.size() - 1);
        auto head_input = activations_.find(head_input_node);
        if (head_input == activations_.end()) {
            throw std::invalid_argument(
                "ExplainerContext::logit_lens: no forward pass has run, so the final module's input "
                "shape is unknown");
        }
        if (!(it->second.shape() == head_input->second.shape())) {
            throw std::invalid_argument(
                "ExplainerContext::logit_lens: the node's cached activation shape does not match the "
                "final module's input shape from the most recent forward pass");
        }
 
        Module* head = modules_.back();
        return head->forward(it->second);
    }
 
    [[nodiscard]] Tensor attention_weights(NodeId node_id) const {
        if (node_id >= graph_.node_count() || node_id > modules_.size()) {
            throw std::invalid_argument(
                "ExplainerContext::attention_weights: node_id is outside the range of nodes the most "
                "recent forward pass produced (or no forward pass has run yet)");
        }
        if (graph_.node(node_id).op_type() != OpType::Attention) {
            throw std::invalid_argument(
                "ExplainerContext::attention_weights: node_id does not name an attention layer");
        }
 
        auto* mha = dynamic_cast<MultiHeadAttentionModule*>(modules_[node_id - 1]);
        PULSATRIX_ASSERT(mha != nullptr);
        return Tensor(mha->last_attention_weights());
    }
 
    [[nodiscard]] CircuitGraph build_circuit_graph(const Tensor& input) {
        const Tensor baseline = forward_pass(input);
        const NodeId output_node = output_node_;
        // Captured before any patching: every patched pass replaces graph_/activations_
        // wholesale, so the baseline's shapes and per-node metadata are read from this
        // self-contained copy rather than from state the loop itself overwrites.
        const ActivationSnapshot clean = activation_snapshot();
 
        std::vector<CircuitNode> nodes;
        nodes.reserve(clean.node_ids().size());
        for (NodeId id : clean.node_ids()) {
            float ablation_effect = 0.0f;
            if (id != output_node) {
                Tensor zero_patch(clean.activation(id));
                zero_patch.fill(0.0f);
                const Tensor patched = forward_pass_with_patch(input, id, zero_patch);
                ablation_effect = l2_distance(baseline, patched);
            }
            nodes.push_back(CircuitNode{id, clean.op_type(id), clean.label(id), ablation_effect});
        }
 
        std::vector<CircuitEdge> edges;
        if (!nodes.empty()) {
            edges.reserve(nodes.size() - 1);
            for (size_t i = 0; i + 1 < nodes.size(); ++i) {
                edges.push_back(CircuitEdge{nodes[i].id, nodes[i + 1].id, nodes[i].ablation_effect});
            }
        }
 
        (void)forward_pass(input);
        return CircuitGraph(std::move(nodes), std::move(edges));
    }
 
    [[nodiscard]] std::optional<std::string> layer_label(NodeId id) const { return graph_.node(id).label(); }
 
private:
    [[nodiscard]] static Tensor check_patch_shape(const Tensor& natural, const Tensor& patch_value) {
        if (!(patch_value.shape() == natural.shape())) {
            throw std::invalid_argument(
                "ExplainerContext::forward_pass_with_patch: patch_value shape must match the patched node's "
                "natural output shape");
        }
        return Tensor(patch_value);
    }
 
    [[nodiscard]] static float l2_distance(const Tensor& a, const Tensor& b) {
        PULSATRIX_ASSERT(a.numel() == b.numel());
        float sum_of_squares = 0.0f;
        for (int64_t i = 0; i < a.numel(); ++i) {
            const float diff = a.data()[i] - b.data()[i];
            sum_of_squares += diff * diff;
        }
        return std::sqrt(sum_of_squares);
    }
 
    std::vector<Module*> modules_;
    ComputationGraph graph_;
    Autograd autograd_;
    std::unordered_map<NodeId, Tensor> activations_;
    NodeId input_node_ = 0;
    NodeId output_node_ = 0;
    bool last_forward_was_patched_ = false;
};
 
}  // namespace pulsatrix