Add Accelerated Failure Time loss for survival analysis task (#4763)
* [WIP] Add lower and upper bounds on the label for survival analysis * Update test MetaInfo.SaveLoadBinary to account for extra two fields * Don't clear qids_ for version 2 of MetaInfo * Add SetInfo() and GetInfo() method for lower and upper bounds * changes to aft * Add parameter class for AFT; use enum's to represent distribution and event type * Add AFT metric * changes to neg grad to grad * changes to binomial loss * changes to overflow * changes to eps * changes to code refactoring * changes to code refactoring * changes to code refactoring * Re-factor survival analysis * Remove aft namespace * Move function bodies out of AFTNormal and AFTLogistic, to reduce clutter * Move function bodies out of AFTLoss, to reduce clutter * Use smart pointer to store AFTDistribution and AFTLoss * Rename AFTNoiseDistribution enum to AFTDistributionType for clarity The enum class was not a distribution itself but a distribution type * Add AFTDistribution::Create() method for convenience * changes to extreme distribution * changes to extreme distribution * changes to extreme * changes to extreme distribution * changes to left censored * deleted cout * changes to x,mu and sd and code refactoring * changes to print * changes to hessian formula in censored and uncensored * changes to variable names and pow * changes to Logistic Pdf * changes to parameter * Expose lower and upper bound labels to R package * Use example weights; normalize log likelihood metric * changes to CHECK * changes to logistic hessian to standard formula * changes to logistic formula * Comply with coding style guideline * Revert back Rabit submodule * Revert dmlc-core submodule * Comply with coding style guideline (clang-tidy) * Fix an error in AFTLoss::Gradient() * Add missing files to amalgamation * Address @RAMitchell's comment: minimize future change in MetaInfo interface * Fix lint * Fix compilation error on 32-bit target, when size_t == bst_uint * Allocate sufficient memory to hold extra label info * Use OpenMP to speed up * Fix compilation on Windows * Address reviewer's feedback * Add unit tests for probability distributions * Make Metric subclass of Configurable * Address reviewer's feedback: Configure() AFT metric * Add a dummy test for AFT metric configuration * Complete AFT configuration test; remove debugging print * Rename AFT parameters * Clarify test comment * Add a dummy test for AFT loss for uncensored case * Fix a bug in AFT loss for uncensored labels * Complete unit test for AFT loss metric * Simplify unit tests for AFT metric * Add unit test to verify aggregate output from AFT metric * Use EXPECT_* instead of ASSERT_*, so that we run all unit tests * Use aft_loss_param when serializing AFTObj This is to be consistent with AFT metric * Add unit tests for AFT Objective * Fix OpenMP bug; clarify semantics for shared variables used in OpenMP loops * Add comments * Remove AFT prefix from probability distribution; put probability distribution in separate source file * Add comments * Define kPI and kEulerMascheroni in probability_distribution.h * Add probability_distribution.cc to amalgamation * Remove unnecessary diff * Address reviewer's feedback: define variables where they're used * Eliminate all INFs and NANs from AFT loss and gradient * Add demo * Add tutorial * Fix lint * Use 'survival:aft' to be consistent with 'survival:cox' * Move sample data to demo/data * Add visual demo with 1D toy data * Add Python tests Co-authored-by: Philip Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Avinash Barnwal
GitHub
107
src/common/probability_distribution.cc
Normal file
107
src/common/probability_distribution.cc
Normal file
@@ -0,0 +1,107 @@
|
||||
/*!
|
||||
* Copyright 2020 by Contributors
|
||||
* \file probability_distribution.cc
|
||||
* \brief Implementation of a few useful probability distributions
|
||||
* \author Avinash Barnwal and Hyunsu Cho
|
||||
*/
|
||||
|
||||
#include <xgboost/logging.h>
|
||||
#include <cmath>
|
||||
#include "probability_distribution.h"
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
|
||||
ProbabilityDistribution* ProbabilityDistribution::Create(ProbabilityDistributionType dist) {
|
||||
switch (dist) {
|
||||
case ProbabilityDistributionType::kNormal:
|
||||
return new NormalDist;
|
||||
case ProbabilityDistributionType::kLogistic:
|
||||
return new LogisticDist;
|
||||
case ProbabilityDistributionType::kExtreme:
|
||||
return new ExtremeDist;
|
||||
default:
|
||||
LOG(FATAL) << "Unknown distribution";
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
double NormalDist::PDF(double z) {
|
||||
const double pdf = std::exp(-z * z / 2) / std::sqrt(2 * probability_constant::kPI);
|
||||
return pdf;
|
||||
}
|
||||
|
||||
double NormalDist::CDF(double z) {
|
||||
const double cdf = 0.5 * (1 + std::erf(z / std::sqrt(2)));
|
||||
return cdf;
|
||||
}
|
||||
|
||||
double NormalDist::GradPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double grad = -1 * z * pdf;
|
||||
return grad;
|
||||
}
|
||||
|
||||
double NormalDist::HessPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double hess = (z * z - 1) * pdf;
|
||||
return hess;
|
||||
}
|
||||
|
||||
double LogisticDist::PDF(double z) {
|
||||
const double w = std::exp(z);
|
||||
const double sqrt_denominator = 1 + w;
|
||||
const double pdf
|
||||
= (std::isinf(w) || std::isinf(w * w)) ? 0.0 : (w / (sqrt_denominator * sqrt_denominator));
|
||||
return pdf;
|
||||
}
|
||||
|
||||
double LogisticDist::CDF(double z) {
|
||||
const double w = std::exp(z);
|
||||
const double cdf = std::isinf(w) ? 1.0 : (w / (1 + w));
|
||||
return cdf;
|
||||
}
|
||||
|
||||
double LogisticDist::GradPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double w = std::exp(z);
|
||||
const double grad = std::isinf(w) ? 0.0 : pdf * (1 - w) / (1 + w);
|
||||
return grad;
|
||||
}
|
||||
|
||||
double LogisticDist::HessPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double w = std::exp(z);
|
||||
const double hess
|
||||
= (std::isinf(w) || std::isinf(w * w)) ? 0.0 : pdf * (w * w - 4 * w + 1) / ((1 + w) * (1 + w));
|
||||
return hess;
|
||||
}
|
||||
|
||||
double ExtremeDist::PDF(double z) {
|
||||
const double w = std::exp(z);
|
||||
const double pdf = std::isinf(w) ? 0.0 : (w * std::exp(-w));
|
||||
return pdf;
|
||||
}
|
||||
|
||||
double ExtremeDist::CDF(double z) {
|
||||
const double w = std::exp(z);
|
||||
const double cdf = 1 - std::exp(-w);
|
||||
return cdf;
|
||||
}
|
||||
|
||||
double ExtremeDist::GradPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double w = std::exp(z);
|
||||
const double grad = std::isinf(w) ? 0.0 : ((1 - w) * pdf);
|
||||
return grad;
|
||||
}
|
||||
|
||||
double ExtremeDist::HessPDF(double z) {
|
||||
const double pdf = this->PDF(z);
|
||||
const double w = std::exp(z);
|
||||
const double hess = (std::isinf(w) || std::isinf(w * w)) ? 0.0 : ((w * w - 3 * w + 1) * pdf);
|
||||
return hess;
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
94
src/common/probability_distribution.h
Normal file
94
src/common/probability_distribution.h
Normal file
@@ -0,0 +1,94 @@
|
||||
/*!
|
||||
* Copyright 2020 by Contributors
|
||||
* \file probability_distribution.h
|
||||
* \brief Implementation of a few useful probability distributions
|
||||
* \author Avinash Barnwal and Hyunsu Cho
|
||||
*/
|
||||
|
||||
#ifndef XGBOOST_COMMON_PROBABILITY_DISTRIBUTION_H_
|
||||
#define XGBOOST_COMMON_PROBABILITY_DISTRIBUTION_H_
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
|
||||
namespace probability_constant {
|
||||
|
||||
/*! \brief Constant PI */
|
||||
const double kPI = 3.14159265358979323846;
|
||||
/*! \brief The Euler-Mascheroni_constant */
|
||||
const double kEulerMascheroni = 0.57721566490153286060651209008240243104215933593992;
|
||||
|
||||
} // namespace probability_constant
|
||||
|
||||
/*! \brief Enum encoding possible choices of probability distribution */
|
||||
enum class ProbabilityDistributionType : int {
|
||||
kNormal = 0, kLogistic = 1, kExtreme = 2
|
||||
};
|
||||
|
||||
/*! \brief Interface for a probability distribution */
|
||||
class ProbabilityDistribution {
|
||||
public:
|
||||
/*!
|
||||
* \brief Evaluate Probability Density Function (PDF) at a particular point
|
||||
* \param z point at which to evaluate PDF
|
||||
* \return Value of PDF evaluated
|
||||
*/
|
||||
virtual double PDF(double z) = 0;
|
||||
/*!
|
||||
* \brief Evaluate Cumulative Distribution Function (CDF) at a particular point
|
||||
* \param z point at which to evaluate CDF
|
||||
* \return Value of CDF evaluated
|
||||
*/
|
||||
virtual double CDF(double z) = 0;
|
||||
/*!
|
||||
* \brief Evaluate first derivative of PDF at a particular point
|
||||
* \param z point at which to evaluate first derivative of PDF
|
||||
* \return Value of first derivative of PDF evaluated
|
||||
*/
|
||||
virtual double GradPDF(double z) = 0;
|
||||
/*!
|
||||
* \brief Evaluate second derivative of PDF at a particular point
|
||||
* \param z point at which to evaluate second derivative of PDF
|
||||
* \return Value of second derivative of PDF evaluated
|
||||
*/
|
||||
virtual double HessPDF(double z) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Factory function to instantiate a new probability distribution object
|
||||
* \param dist kind of probability distribution
|
||||
* \return Reference to the newly created probability distribution object
|
||||
*/
|
||||
static ProbabilityDistribution* Create(ProbabilityDistributionType dist);
|
||||
};
|
||||
|
||||
/*! \brief The (standard) normal distribution */
|
||||
class NormalDist : public ProbabilityDistribution {
|
||||
public:
|
||||
double PDF(double z) override;
|
||||
double CDF(double z) override;
|
||||
double GradPDF(double z) override;
|
||||
double HessPDF(double z) override;
|
||||
};
|
||||
|
||||
/*! \brief The (standard) logistic distribution */
|
||||
class LogisticDist : public ProbabilityDistribution {
|
||||
public:
|
||||
double PDF(double z) override;
|
||||
double CDF(double z) override;
|
||||
double GradPDF(double z) override;
|
||||
double HessPDF(double z) override;
|
||||
};
|
||||
|
||||
/*! \brief The extreme distribution, also known as the Gumbel (minimum) distribution */
|
||||
class ExtremeDist : public ProbabilityDistribution {
|
||||
public:
|
||||
double PDF(double z) override;
|
||||
double CDF(double z) override;
|
||||
double GradPDF(double z) override;
|
||||
double HessPDF(double z) override;
|
||||
};
|
||||
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
|
||||
#endif // XGBOOST_COMMON_PROBABILITY_DISTRIBUTION_H_
|
||||
146
src/common/survival_util.cc
Normal file
146
src/common/survival_util.cc
Normal file
@@ -0,0 +1,146 @@
|
||||
/*!
|
||||
* Copyright 2019 by Contributors
|
||||
* \file survival_util.cc
|
||||
* \brief Utility functions, useful for implementing objective and metric functions for survival
|
||||
* analysis
|
||||
* \author Avinash Barnwal, Hyunsu Cho and Toby Hocking
|
||||
*/
|
||||
|
||||
#include <dmlc/registry.h>
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include "survival_util.h"
|
||||
|
||||
/*
|
||||
- Formulas are motivated from document -
|
||||
http://members.cbio.mines-paristech.fr/~thocking/survival.pdf
|
||||
- Detailed Derivation of Loss/Gradient/Hessian -
|
||||
https://github.com/avinashbarnwal/GSOC-2019/blob/master/doc/Accelerated_Failure_Time.pdf
|
||||
*/
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
|
||||
DMLC_REGISTER_PARAMETER(AFTParam);
|
||||
|
||||
double AFTLoss::Loss(double y_lower, double y_upper, double y_pred, double sigma) {
|
||||
const double log_y_lower = std::log(y_lower);
|
||||
const double log_y_upper = std::log(y_upper);
|
||||
const double eps = 1e-12;
|
||||
double cost;
|
||||
|
||||
if (y_lower == y_upper) { // uncensored
|
||||
const double z = (log_y_lower - y_pred) / sigma;
|
||||
const double pdf = dist_->PDF(z);
|
||||
// Regularize the denominator with eps, to avoid INF or NAN
|
||||
cost = -std::log(std::max(pdf / (sigma * y_lower), eps));
|
||||
} else { // censored; now check what type of censorship we have
|
||||
double z_u, z_l, cdf_u, cdf_l;
|
||||
if (std::isinf(y_upper)) { // right-censored
|
||||
cdf_u = 1;
|
||||
} else { // left-censored or interval-censored
|
||||
z_u = (log_y_upper - y_pred) / sigma;
|
||||
cdf_u = dist_->CDF(z_u);
|
||||
}
|
||||
if (std::isinf(y_lower)) { // left-censored
|
||||
cdf_l = 0;
|
||||
} else { // right-censored or interval-censored
|
||||
z_l = (log_y_lower - y_pred) / sigma;
|
||||
cdf_l = dist_->CDF(z_l);
|
||||
}
|
||||
// Regularize the denominator with eps, to avoid INF or NAN
|
||||
cost = -std::log(std::max(cdf_u - cdf_l, eps));
|
||||
}
|
||||
|
||||
return cost;
|
||||
}
|
||||
|
||||
double AFTLoss::Gradient(double y_lower, double y_upper, double y_pred, double sigma) {
|
||||
const double log_y_lower = std::log(y_lower);
|
||||
const double log_y_upper = std::log(y_upper);
|
||||
double gradient;
|
||||
const double eps = 1e-12;
|
||||
|
||||
if (y_lower == y_upper) { // uncensored
|
||||
const double z = (log_y_lower - y_pred) / sigma;
|
||||
const double pdf = dist_->PDF(z);
|
||||
const double grad_pdf = dist_->GradPDF(z);
|
||||
// Regularize the denominator with eps, so that gradient doesn't get too big
|
||||
gradient = grad_pdf / (sigma * std::max(pdf, eps));
|
||||
} else { // censored; now check what type of censorship we have
|
||||
double z_u, z_l, pdf_u, pdf_l, cdf_u, cdf_l;
|
||||
if (std::isinf(y_upper)) { // right-censored
|
||||
pdf_u = 0;
|
||||
cdf_u = 1;
|
||||
} else { // interval-censored or left-censored
|
||||
z_u = (log_y_upper - y_pred) / sigma;
|
||||
pdf_u = dist_->PDF(z_u);
|
||||
cdf_u = dist_->CDF(z_u);
|
||||
}
|
||||
if (std::isinf(y_lower)) { // left-censored
|
||||
pdf_l = 0;
|
||||
cdf_l = 0;
|
||||
} else { // interval-censored or right-censored
|
||||
z_l = (log_y_lower - y_pred) / sigma;
|
||||
pdf_l = dist_->PDF(z_l);
|
||||
cdf_l = dist_->CDF(z_l);
|
||||
}
|
||||
// Regularize the denominator with eps, so that gradient doesn't get too big
|
||||
gradient = (pdf_u - pdf_l) / (sigma * std::max(cdf_u - cdf_l, eps));
|
||||
}
|
||||
|
||||
return gradient;
|
||||
}
|
||||
|
||||
double AFTLoss::Hessian(double y_lower, double y_upper, double y_pred, double sigma) {
|
||||
const double log_y_lower = std::log(y_lower);
|
||||
const double log_y_upper = std::log(y_upper);
|
||||
const double eps = 1e-12;
|
||||
double hessian;
|
||||
|
||||
if (y_lower == y_upper) { // uncensored
|
||||
const double z = (log_y_lower - y_pred) / sigma;
|
||||
const double pdf = dist_->PDF(z);
|
||||
const double grad_pdf = dist_->GradPDF(z);
|
||||
const double hess_pdf = dist_->HessPDF(z);
|
||||
// Regularize the denominator with eps, so that gradient doesn't get too big
|
||||
hessian = -(pdf * hess_pdf - std::pow(grad_pdf, 2))
|
||||
/ (std::pow(sigma, 2) * std::pow(std::max(pdf, eps), 2));
|
||||
} else { // censored; now check what type of censorship we have
|
||||
double z_u, z_l, grad_pdf_u, grad_pdf_l, pdf_u, pdf_l, cdf_u, cdf_l;
|
||||
if (std::isinf(y_upper)) { // right-censored
|
||||
pdf_u = 0;
|
||||
cdf_u = 1;
|
||||
grad_pdf_u = 0;
|
||||
} else { // interval-censored or left-censored
|
||||
z_u = (log_y_upper - y_pred) / sigma;
|
||||
pdf_u = dist_->PDF(z_u);
|
||||
cdf_u = dist_->CDF(z_u);
|
||||
grad_pdf_u = dist_->GradPDF(z_u);
|
||||
}
|
||||
if (std::isinf(y_lower)) { // left-censored
|
||||
pdf_l = 0;
|
||||
cdf_l = 0;
|
||||
grad_pdf_l = 0;
|
||||
} else { // interval-censored or right-censored
|
||||
z_l = (log_y_lower - y_pred) / sigma;
|
||||
pdf_l = dist_->PDF(z_l);
|
||||
cdf_l = dist_->CDF(z_l);
|
||||
grad_pdf_l = dist_->GradPDF(z_l);
|
||||
}
|
||||
const double cdf_diff = cdf_u - cdf_l;
|
||||
const double pdf_diff = pdf_u - pdf_l;
|
||||
const double grad_diff = grad_pdf_u - grad_pdf_l;
|
||||
// Regularize the denominator with eps, so that gradient doesn't get too big
|
||||
const double cdf_diff_thresh = std::max(cdf_diff, eps);
|
||||
const double numerator = -(cdf_diff * grad_diff - pdf_diff * pdf_diff);
|
||||
const double sqrt_denominator = sigma * cdf_diff_thresh;
|
||||
const double denominator = sqrt_denominator * sqrt_denominator;
|
||||
hessian = numerator / denominator;
|
||||
}
|
||||
|
||||
return hessian;
|
||||
}
|
||||
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
85
src/common/survival_util.h
Normal file
85
src/common/survival_util.h
Normal file
@@ -0,0 +1,85 @@
|
||||
/*!
|
||||
* Copyright 2019 by Contributors
|
||||
* \file survival_util.h
|
||||
* \brief Utility functions, useful for implementing objective and metric functions for survival
|
||||
* analysis
|
||||
* \author Avinash Barnwal, Hyunsu Cho and Toby Hocking
|
||||
*/
|
||||
#ifndef XGBOOST_COMMON_SURVIVAL_UTIL_H_
|
||||
#define XGBOOST_COMMON_SURVIVAL_UTIL_H_
|
||||
|
||||
#include <xgboost/parameter.h>
|
||||
#include <memory>
|
||||
#include "probability_distribution.h"
|
||||
|
||||
DECLARE_FIELD_ENUM_CLASS(xgboost::common::ProbabilityDistributionType);
|
||||
|
||||
namespace xgboost {
|
||||
namespace common {
|
||||
|
||||
/*! \brief Parameter structure for AFT loss and metric */
|
||||
struct AFTParam : public XGBoostParameter<AFTParam> {
|
||||
/*! \brief Choice of probability distribution for the noise term in AFT */
|
||||
ProbabilityDistributionType aft_loss_distribution;
|
||||
/*! \brief Scaling factor to be applied to the distribution */
|
||||
float aft_loss_distribution_scale;
|
||||
DMLC_DECLARE_PARAMETER(AFTParam) {
|
||||
DMLC_DECLARE_FIELD(aft_loss_distribution)
|
||||
.set_default(ProbabilityDistributionType::kNormal)
|
||||
.add_enum("normal", ProbabilityDistributionType::kNormal)
|
||||
.add_enum("logistic", ProbabilityDistributionType::kLogistic)
|
||||
.add_enum("extreme", ProbabilityDistributionType::kExtreme)
|
||||
.describe("Choice of distribution for the noise term in "
|
||||
"Accelerated Failure Time model");
|
||||
DMLC_DECLARE_FIELD(aft_loss_distribution_scale)
|
||||
.set_default(1.0f)
|
||||
.describe("Scaling factor used to scale the distribution in "
|
||||
"Accelerated Failure Time model");
|
||||
}
|
||||
};
|
||||
|
||||
/*! \brief The AFT loss function */
|
||||
class AFTLoss {
|
||||
private:
|
||||
std::unique_ptr<ProbabilityDistribution> dist_;
|
||||
|
||||
public:
|
||||
/*!
|
||||
* \brief Constructor for AFT loss function
|
||||
* \param dist Choice of probability distribution for the noise term in AFT
|
||||
*/
|
||||
explicit AFTLoss(ProbabilityDistributionType dist) {
|
||||
dist_.reset(ProbabilityDistribution::Create(dist));
|
||||
}
|
||||
|
||||
public:
|
||||
/*!
|
||||
* \brief Compute the AFT loss
|
||||
* \param y_lower Lower bound for the true label
|
||||
* \param y_upper Upper bound for the true label
|
||||
* \param y_pred Predicted label
|
||||
* \param sigma Scaling factor to be applied to the distribution of the noise term
|
||||
*/
|
||||
double Loss(double y_lower, double y_upper, double y_pred, double sigma);
|
||||
/*!
|
||||
* \brief Compute the gradient of the AFT loss
|
||||
* \param y_lower Lower bound for the true label
|
||||
* \param y_upper Upper bound for the true label
|
||||
* \param y_pred Predicted label
|
||||
* \param sigma Scaling factor to be applied to the distribution of the noise term
|
||||
*/
|
||||
double Gradient(double y_lower, double y_upper, double y_pred, double sigma);
|
||||
/*!
|
||||
* \brief Compute the hessian of the AFT loss
|
||||
* \param y_lower Lower bound for the true label
|
||||
* \param y_upper Upper bound for the true label
|
||||
* \param y_pred Predicted label
|
||||
* \param sigma Scaling factor to be applied to the distribution of the noise term
|
||||
*/
|
||||
double Hessian(double y_lower, double y_upper, double y_pred, double sigma);
|
||||
};
|
||||
|
||||
} // namespace common
|
||||
} // namespace xgboost
|
||||
|
||||
#endif // XGBOOST_COMMON_SURVIVAL_UTIL_H_
|
||||
Reference in New Issue
Block a user