linear_probe.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cmath>
#include <cstdint>
#include <random>
#include <stdexcept>
#include <string>
#include <vector>
 
#include "pulsatrix/determinism.hpp"
#include "pulsatrix/bce_with_logits_loss.hpp"
#include "pulsatrix/linear_module.hpp"
 
namespace pulsatrix {
 
class LinearProbe {
public:
    LinearProbe(int64_t activation_dim, DeviceBackend* backend)
        : LinearProbe(activation_dim, backend, static_cast<unsigned>(next_seed())) {}
 
    LinearProbe(int64_t activation_dim, DeviceBackend* backend, unsigned seed)
        : activation_dim_(activation_dim),
          // Clamped only so a rejected dimension can't reach LinearModule/Shape and throw
          // *their* message before the check below throws this class's own, clearer one --
          // member initialization necessarily runs before the constructor body.
          classifier_(activation_dim > 0 ? activation_dim : 1, 1, backend),
          loss_(backend) {
        if (activation_dim <= 0) {
            throw std::invalid_argument("LinearProbe: activation_dim must be positive");
        }
        std::mt19937 rng(seed);
        std::vector<float> weights(static_cast<size_t>(activation_dim));
        for (float& w : weights) {
            w = uniform_symmetric(rng, 0.01f);
        }
        classifier_.set_weight(weights);
        classifier_.set_bias(std::vector<float>{uniform_symmetric(rng, 0.01f)});
    }
 
    template <typename OptimizerT>
    float train_step(const Tensor& activation_batch, const Tensor& label_batch, OptimizerT& optimizer) {
        validate_batch(activation_batch, label_batch, "LinearProbe::train_step");
 
        optimizer.zero_grad(classifier_);
 
        Tensor logits = classifier_.forward(activation_batch);
        const float loss_value = loss_.forward(logits, label_batch);
 
        Tensor grad_logits = loss_.backward();
        (void)classifier_.backward(grad_logits);  // the activations have no upstream to receive this
 
        optimizer.step(classifier_);
        return loss_value;
    }
 
    [[nodiscard]] float accuracy(const Tensor& activation_batch, const Tensor& label_batch) const {
        validate_batch(activation_batch, label_batch, "LinearProbe::accuracy");
 
        const Tensor logits = classifier_.forward(activation_batch);
        const int64_t n = logits.numel();
        int64_t correct = 0;
        for (int64_t i = 0; i < n; ++i) {
            const bool predicted_positive = logits.data()[i] >= 0.0f;
            const bool labeled_positive = label_batch.data()[i] >= 0.5f;
            if (predicted_positive == labeled_positive) {
                ++correct;
            }
        }
        return static_cast<float>(correct) / static_cast<float>(n);
    }
 
    [[nodiscard]] int64_t activation_dim() const { return activation_dim_; }
 
    [[nodiscard]] LinearModule& classifier() { return classifier_; }
 
    [[nodiscard]] const LinearModule& classifier() const { return classifier_; }
 
private:
    static float uniform_symmetric(std::mt19937& rng, float scale) {
        // std::uniform_real_distribution's mapping is implementation-defined, so identical
        // seeds would not give identical weights across standard libraries. This mapping is
        // fixed here, making a seeded probe reproducible everywhere.
        const float unit = static_cast<float>(rng() - std::mt19937::min()) /
                            static_cast<float>(std::mt19937::max() - std::mt19937::min() + 1ull);
        return (unit * 2.0f - 1.0f) * scale;
    }
 
    void validate_batch(const Tensor& activation_batch, const Tensor& label_batch, const char* method) const {
        const std::string where(method);
        if (activation_batch.rank() != 2) {
            throw std::invalid_argument(where + ": activation_batch must be rank-2 (N, activation_dim)");
        }
        if (label_batch.rank() != 2 || label_batch.shape().dim(1) != 1) {
            throw std::invalid_argument(where + ": label_batch must be rank-2 (N, 1)");
        }
        if (activation_batch.shape().dim(1) != activation_dim_) {
            throw std::invalid_argument(where + ": activation_batch width must equal activation_dim");
        }
        if (activation_batch.shape().dim(0) != label_batch.shape().dim(0)) {
            throw std::invalid_argument(where + ": activation_batch and label_batch must have the same batch size");
        }
        if (activation_batch.shape().dim(0) <= 0) {
            throw std::invalid_argument(where + ": batch must not be empty");
        }
    }
 
    int64_t activation_dim_;
    mutable LinearModule classifier_;
    BCEWithLogitsLoss loss_;
};
 
}  // namespace pulsatrix