egan_mutation.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <cmath>
#include <limits>
#include <stdexcept>
#include <vector>
 
#include "pulsatrix/bce_with_logits_loss.hpp"
#include "pulsatrix/module.hpp"
#include "pulsatrix/mse_loss.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
enum class MutationObjective { Minimax, Heuristic, LeastSquares };
 
class MutationLoss {
public:
    explicit MutationLoss(DeviceBackend* backend) : backend_(backend), bce_(backend), mse_(backend) {}
 
    float Forward(MutationObjective objective, const Tensor& logits) {
        objective_ = objective;
        Tensor ones(logits.shape(), backend_);
        ones.fill(1.0f);
        switch (objective) {
            case MutationObjective::Heuristic:
                return bce_.forward(logits, ones);
            case MutationObjective::LeastSquares:
                return mse_.forward(logits, ones);
            case MutationObjective::Minimax: {
                Tensor zeros(logits.shape(), backend_);
                zeros.fill(0.0f);
                return -bce_.forward(logits, zeros);
            }
        }
        throw std::invalid_argument("MutationLoss::Forward: unknown MutationObjective");
    }
 
    [[nodiscard]] Tensor Backward() const {
        switch (objective_) {
            case MutationObjective::Heuristic:
                return bce_.backward();
            case MutationObjective::LeastSquares:
                return mse_.backward();
            case MutationObjective::Minimax:
                return Negate(bce_.backward());
        }
        throw std::invalid_argument("MutationLoss::Backward: unknown MutationObjective");
    }
 
private:
    [[nodiscard]] Tensor Negate(const Tensor& t) const {
        std::vector<float> values(static_cast<size_t>(t.numel()));
        for (int64_t i = 0; i < t.numel(); ++i) {
            values[static_cast<size_t>(i)] = -t.data()[i];
        }
        return Tensor(t.shape(), backend_, values, t.device());
    }
 
    DeviceBackend* backend_;
    BCEWithLogitsLoss bce_;
    MSELoss mse_;
    MutationObjective objective_ = MutationObjective::Heuristic;
};
 
inline float QualityFitness(const Tensor& logits) {
    double total = 0.0;
    for (int64_t i = 0; i < logits.numel(); ++i) {
        total += 1.0 / (1.0 + std::exp(-static_cast<double>(logits.data()[i])));
    }
    return static_cast<float>(total / static_cast<double>(logits.numel()));
}
 
inline float DiversityFitness(Module& discriminator, const Tensor& fake, DeviceBackend* backend) {
    auto params = discriminator.parameters();
    if (params.empty()) {
        throw std::invalid_argument("DiversityFitness: discriminator must have at least one parameter");
    }
 
    BCEWithLogitsLoss bce(backend);
    Tensor logits = discriminator.forward(fake);
    Tensor zeros(logits.shape(), backend);
    zeros.fill(0.0f);
    (void)bce.forward(logits, zeros);
    (void)discriminator.backward(bce.backward());
 
    double sum_squares = 0.0;
    for (const auto& p : params) {
        for (int64_t i = 0; i < p.grad->numel(); ++i) {
            double g = p.grad->data()[i];
            sum_squares += g * g;
        }
    }
    double norm = std::sqrt(sum_squares);
    // An exactly-zero gradient (e.g. against a freshly-constructed, all-zero-weight
    // discriminator) maps to +infinity fitness -- treated as "maximally diverse" by explicit
    // convention rather than left as an undefined log(0), a documented edge case, not a
    // silent bug.
    if (norm <= 0.0) {
        return std::numeric_limits<float>::infinity();
    }
    return static_cast<float>(-std::log(norm));
}
 
inline float CombinedFitness(float quality, float diversity, float gamma = 0.05f) {
    return quality + gamma * diversity;
}
 
template <typename OptimizerT>
float RunMutationStep(Module& offspring, Module& discriminator, const Tensor& noise, MutationObjective objective,
                       OptimizerT& g_optimizer, DeviceBackend* backend, float gamma = 0.05f) {
    Tensor fake = offspring.forward(noise);
    Tensor logits = discriminator.forward(fake);
    MutationLoss loss(backend);
    (void)loss.Forward(objective, logits);
    Tensor grad_fake = discriminator.backward(loss.Backward());
    (void)offspring.backward(grad_fake);
    g_optimizer.step(offspring);
    g_optimizer.zero_grad(offspring);
 
    Tensor post_fake = offspring.forward(noise);
    Tensor post_logits = discriminator.forward(post_fake);
    float quality = QualityFitness(post_logits);
    float diversity = DiversityFitness(discriminator, post_fake, backend);
    return CombinedFitness(quality, diversity, gamma);
}
 
}  // namespace pulsatrix