FastOnlineSupervisedMStep.hpp

#ifndef _LDAPLUSPLUS_EM_FASTONLINESUPERVISEDMSTEP_HPP_
#define _LDAPLUSPLUS_EM_FASTONLINESUPERVISEDMSTEP_HPP_

#include "ldaplusplus/em/MStepInterface.hpp"

namespace ldaplusplus {
namespace em {


template <typename Scalar>
class FastOnlineSupervisedMStep : public MStepInterface<Scalar>
{
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> VectorX;

    public:
        FastOnlineSupervisedMStep(
            VectorX class_weights,
            Scalar regularization_penalty = 1e-2,
            size_t minibatch_size = 128,
            Scalar eta_momentum = 0.9,
            Scalar eta_learning_rate = 0.01,
            Scalar beta_weight = 0.9
        );
        FastOnlineSupervisedMStep(
            size_t num_classes,
            Scalar regularization_penalty = 1e-2,
            size_t minibatch_size = 128,
            Scalar eta_momentum = 0.9,
            Scalar eta_learning_rate = 0.01,
            Scalar beta_weight = 0.9
        );

        virtual void m_step(
            std::shared_ptr<parameters::Parameters> parameters
        ) override;

        virtual void doc_m_step(
            const std::shared_ptr<corpus::Document> doc,
            const std::shared_ptr<parameters::Parameters> v_parameters,
            std::shared_ptr<parameters::Parameters> m_parameters
        ) override;

    private:
        // Number of classes
        VectorX class_weights_;
        size_t num_classes_;

        // Minibatch size and portion (the portion of the corpus)
        size_t minibatch_size_;

        // The regularization penalty for the multinomial logistic regression
        // Mind that it should account for the minibatch size
        Scalar regularization_penalty_;

        // The suff stats and data needed to optimize the ELBO w.r.t. model
        // parameters
        MatrixX b_;
        Scalar beta_weight_;
        MatrixX expected_z_bar_;
        Eigen::VectorXi y_;
        MatrixX eta_velocity_;
        MatrixX eta_gradient_;
        Scalar eta_momentum_;
        Scalar eta_learning_rate_;

        // The number of document's seen so far
        size_t docs_seen_so_far_;
};

}  // namespace em
}  // namespace ldaplusplus

#endif  // _LDAPLUSPLUS_EM_FASTONLINESUPERVISEDMSTEP_HPP_