sparse_autoencoder.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/linear_module.hpp"
#include "pulsatrix/mse_loss.hpp"
#include "pulsatrix/relu_module.hpp"
 
namespace pulsatrix {
 
class SparseAutoencoder {
public:
    SparseAutoencoder(int64_t dim, int64_t hidden_dim, float l1_lambda, DeviceBackend* backend)
        : SparseAutoencoder(dim, hidden_dim, l1_lambda, backend, static_cast<unsigned>(next_seed())) {}
 
    SparseAutoencoder(int64_t dim, int64_t hidden_dim, float l1_lambda, DeviceBackend* backend, unsigned seed)
        : dim_(dim),
          hidden_dim_(hidden_dim),
          l1_lambda_(l1_lambda),
          backend_(backend),
          // 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.
          encoder_(dim > 0 ? dim : 1, hidden_dim > 0 ? hidden_dim : 1, backend),
          relu_(backend),
          decoder_(hidden_dim > 0 ? hidden_dim : 1, dim > 0 ? dim : 1, backend),
          loss_(backend) {
        if (dim <= 0) {
            throw std::invalid_argument("SparseAutoencoder: dim must be positive");
        }
        if (hidden_dim <= 0) {
            throw std::invalid_argument("SparseAutoencoder: hidden_dim must be positive");
        }
        if (l1_lambda < 0.0f) {
            throw std::invalid_argument("SparseAutoencoder: l1_lambda must be non-negative");
        }
 
        std::mt19937 rng(seed);
        // Fan-in scaling: each layer's init range shrinks with its own input width, so the
        // pre-activation magnitude at step 0 stays O(1) regardless of dim/hidden_dim rather
        // than growing with the layer's width (which for a deliberately overcomplete SAE
        // would otherwise saturate the decoder's input from the first step).
        init_layer(encoder_, rng, 1.0f / std::sqrt(static_cast<float>(dim)));
        init_layer(decoder_, rng, 1.0f / std::sqrt(static_cast<float>(hidden_dim)));
    }
 
    template <typename OptimizerT>
    float train_step(const Tensor& input_batch, OptimizerT& optimizer) {
        validate_batch(input_batch, "SparseAutoencoder::train_step");
 
        optimizer.zero_grad(encoder_);
        optimizer.zero_grad(decoder_);
 
        const Tensor hidden = relu_.forward(encoder_.forward(input_batch));
        const Tensor reconstruction = decoder_.forward(hidden);
        const float reconstruction_loss = loss_.forward(reconstruction, input_batch);
 
        const Tensor grad_reconstruction = loss_.backward();
        Tensor grad_hidden = decoder_.backward(grad_reconstruction);
 
        const float batch_size = static_cast<float>(input_batch.shape().dim(0));
        Tensor l1_grad(grad_hidden.shape(), backend_, grad_hidden.device());
        l1_grad.fill(l1_lambda_ / batch_size);
        grad_hidden.accumulate(l1_grad);
 
        const Tensor grad_pre_activation = relu_.backward(grad_hidden);
        (void)encoder_.backward(grad_pre_activation);  // the input has no upstream to receive this
 
        optimizer.step(encoder_);
        optimizer.step(decoder_);
        // relu_ has no parameters (Module::parameters() default) -- stepping it would be a
        // safe no-op but would imply there is something to update.
        return reconstruction_loss;
    }
 
    [[nodiscard]] Tensor reconstruct(const Tensor& input_batch) const {
        validate_batch(input_batch, "SparseAutoencoder::reconstruct");
        return decoder_.forward(encode(input_batch));
    }
 
    [[nodiscard]] float reconstruction_error(const Tensor& input_batch) const {
        validate_batch(input_batch, "SparseAutoencoder::reconstruction_error");
 
        const Tensor reconstruction = reconstruct(input_batch);
        const int64_t n = reconstruction.numel();
        float sum_squared = 0.0f;
        for (int64_t i = 0; i < n; ++i) {
            const float diff = reconstruction.data()[i] - input_batch.data()[i];
            sum_squared += diff * diff;
        }
        return sum_squared / static_cast<float>(n);
    }
 
    [[nodiscard]] float mean_hidden_activation(const Tensor& input_batch) const {
        validate_batch(input_batch, "SparseAutoencoder::mean_hidden_activation");
 
        const Tensor hidden = encode(input_batch);
        const int64_t n = hidden.numel();
        float sum = 0.0f;
        for (int64_t i = 0; i < n; ++i) {
            sum += hidden.data()[i];
        }
        return sum / static_cast<float>(n);
    }
 
    [[nodiscard]] int64_t dim() const { return dim_; }
 
    [[nodiscard]] int64_t hidden_dim() const { return hidden_dim_; }
 
    [[nodiscard]] float l1_lambda() const { return l1_lambda_; }
 
    [[nodiscard]] LinearModule& encoder() { return encoder_; }
 
    [[nodiscard]] const LinearModule& encoder() const { return encoder_; }
 
    [[nodiscard]] LinearModule& decoder() { return decoder_; }
 
    [[nodiscard]] const LinearModule& decoder() const { return decoder_; }
 
private:
    [[nodiscard]] Tensor encode(const Tensor& input_batch) const { return relu_.forward(encoder_.forward(input_batch)); }
 
    static void init_layer(LinearModule& layer, std::mt19937& rng, float scale) {
        std::vector<float> weights(static_cast<size_t>(layer.weight().numel()));
        for (float& w : weights) {
            w = uniform_symmetric(rng, scale);
        }
        layer.set_weight(weights);
        layer.set_bias(std::vector<float>(static_cast<size_t>(layer.bias().numel()), 0.0f));
    }
 
    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 SAE reproducible everywhere -- the same convention
        // LinearProbe already uses.
        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& input_batch, const char* method) const {
        const std::string where(method);
        if (input_batch.rank() != 2) {
            throw std::invalid_argument(where + ": input_batch must be rank-2 (N, dim)");
        }
        if (input_batch.shape().dim(1) != dim_) {
            throw std::invalid_argument(where + ": input_batch width must equal dim");
        }
        if (input_batch.shape().dim(0) <= 0) {
            throw std::invalid_argument(where + ": batch must not be empty");
        }
    }
 
    int64_t dim_;
    int64_t hidden_dim_;
    float l1_lambda_;
    DeviceBackend* backend_;
    mutable LinearModule encoder_;
    mutable ReluModule relu_;
    mutable LinearModule decoder_;
    MSELoss loss_;
};
 
}  // namespace pulsatrix