Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Implement robust regularization in 'survival:aft' objective (#5473)

Browse Source 
* Robust regularization of AFT gradient and hessian

* Fix AFT doc; expose it to tutorial TOC

* Apply robust regularization to uncensored case too

* Revise unit test slightly

* Fix lint

* Update test_survival.py

* Use GradientPairPrecise

* Remove unused variables
This commit is contained in:
Philip Hyunsu Cho
2020-04-04 12:21:24 -07:00
committed by GitHub
parent 939973630d
commit 5fc5ec539d
9 changed files with 205 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

166
src/common/survival_util.cc
View File

@@ -18,6 +18,106 @@
https://github.com/avinashbarnwal/GSOC-2019/blob/master/doc/Accelerated_Failure_Time.pdf
*/
namespace {
// Allowable range for gradient and hessian. Used for regularization
constexpr double kMinGradient = -15.0;
constexpr double kMaxGradient = 15.0;
constexpr double kMinHessian = 1e-16; // Ensure that no data point gets zero hessian
constexpr double kMaxHessian = 15.0;
constexpr double kEps = 1e-12; // A denomitor in a fraction should not be too small
// Clip (limit) x to fit range [x_min, x_max].
// If x < x_min, return x_min; if x > x_max, return x_max; if x_min <= x <= x_max, return x.
// This function assumes x_min < x_max; behavior is undefined if this assumption does not hold.
inline double Clip(double x, double x_min, double x_max) {
if (x < x_min) {
return x_min;
}
if (x > x_max) {
return x_max;
}
return x;
}
using xgboost::common::ProbabilityDistributionType;
enum class CensoringType : uint8_t {
kUncensored, kRightCensored, kLeftCensored, kIntervalCensored
};
using xgboost::GradientPairPrecise;
inline GradientPairPrecise GetLimitAtInfPred(ProbabilityDistributionType dist_type,
CensoringType censor_type,
double sign, double sigma) {
switch (censor_type) {
case CensoringType::kUncensored:
switch (dist_type) {
case ProbabilityDistributionType::kNormal:
return sign ? GradientPairPrecise{ kMinGradient, 1.0 / (sigma * sigma) }
: GradientPairPrecise{ kMaxGradient, 1.0 / (sigma * sigma) };
case ProbabilityDistributionType::kLogistic:
return sign ? GradientPairPrecise{ -1.0 / sigma, kMinHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
case ProbabilityDistributionType::kExtreme:
return sign ? GradientPairPrecise{ kMinGradient, kMaxHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
default:
LOG(FATAL) << "Unknown distribution type";
}
case CensoringType::kRightCensored:
switch (dist_type) {
case ProbabilityDistributionType::kNormal:
return sign ? GradientPairPrecise{ kMinGradient, 1.0 / (sigma * sigma) }
: GradientPairPrecise{ 0.0, kMinHessian };
case ProbabilityDistributionType::kLogistic:
return sign ? GradientPairPrecise{ -1.0 / sigma, kMinHessian }
: GradientPairPrecise{ 0.0, kMinHessian };
case ProbabilityDistributionType::kExtreme:
return sign ? GradientPairPrecise{ kMinGradient, kMaxHessian }
: GradientPairPrecise{ 0.0, kMinHessian };
default:
LOG(FATAL) << "Unknown distribution type";
}
case CensoringType::kLeftCensored:
switch (dist_type) {
case ProbabilityDistributionType::kNormal:
return sign ? GradientPairPrecise{ 0.0, kMinHessian }
: GradientPairPrecise{ kMaxGradient, 1.0 / (sigma * sigma) };
case ProbabilityDistributionType::kLogistic:
return sign ? GradientPairPrecise{ 0.0, kMinHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
case ProbabilityDistributionType::kExtreme:
return sign ? GradientPairPrecise{ 0.0, kMinHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
default:
LOG(FATAL) << "Unknown distribution type";
}
case CensoringType::kIntervalCensored:
switch (dist_type) {
case ProbabilityDistributionType::kNormal:
return sign ? GradientPairPrecise{ kMinGradient, 1.0 / (sigma * sigma) }
: GradientPairPrecise{ kMaxGradient, 1.0 / (sigma * sigma) };
case ProbabilityDistributionType::kLogistic:
return sign ? GradientPairPrecise{ -1.0 / sigma, kMinHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
case ProbabilityDistributionType::kExtreme:
return sign ? GradientPairPrecise{ kMinGradient, kMaxHessian }
: GradientPairPrecise{ 1.0 / sigma, kMinHessian };
default:
LOG(FATAL) << "Unknown distribution type";
}
default:
LOG(FATAL) << "Unknown censoring type";
}
return { 0.0, 0.0 };
}
} // anonymous namespace
namespace xgboost {
namespace common {
@@ -26,14 +126,14 @@ 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));
cost = -std::log(std::max(pdf / (sigma * y_lower), kEps));
} 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
@@ -49,7 +149,7 @@ double AFTLoss::Loss(double y_lower, double y_upper, double y_pred, double 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));
cost = -std::log(std::max(cdf_u - cdf_l, kEps));
}
return cost;
@@ -58,20 +158,25 @@ double AFTLoss::Loss(double y_lower, double y_upper, double y_pred, double sigma
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;
double numerator, denominator, gradient; // numerator and denominator of gradient
CensoringType censor_type;
bool z_sign; // sign of z-score
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));
censor_type = CensoringType::kUncensored;
numerator = grad_pdf;
denominator = sigma * pdf;
z_sign = (z > 0);
} else { // censored; now check what type of censorship we have
double z_u, z_l, pdf_u, pdf_l, cdf_u, cdf_l;
double z_u = 0.0, z_l = 0.0, pdf_u, pdf_l, cdf_u, cdf_l;
censor_type = CensoringType::kIntervalCensored;
if (std::isinf(y_upper)) { // right-censored
pdf_u = 0;
cdf_u = 1;
censor_type = CensoringType::kRightCensored;
} else { // interval-censored or left-censored
z_u = (log_y_upper - y_pred) / sigma;
pdf_u = dist_->PDF(z_u);
@@ -80,38 +185,48 @@ double AFTLoss::Gradient(double y_lower, double y_upper, double y_pred, double s
if (std::isinf(y_lower)) { // left-censored
pdf_l = 0;
cdf_l = 0;
censor_type = CensoringType::kLeftCensored;
} 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));
z_sign = (z_u > 0 || z_l > 0);
numerator = pdf_u - pdf_l;
denominator = sigma * (cdf_u - cdf_l);
}
gradient = numerator / denominator;
if (denominator < kEps && (std::isnan(gradient) || std::isinf(gradient))) {
gradient = GetLimitAtInfPred(dist_type_, censor_type, z_sign, sigma).GetGrad();
}
return gradient;
return Clip(gradient, kMinGradient, kMaxGradient);
}
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;
double numerator, denominator, hessian; // numerator and denominator of hessian
CensoringType censor_type;
bool z_sign; // sign of z-score
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));
censor_type = CensoringType::kUncensored;
numerator = -(pdf * hess_pdf - grad_pdf * grad_pdf);
denominator = sigma * sigma * pdf * pdf;
z_sign = (z > 0);
} 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;
double z_u = 0.0, z_l = 0.0, grad_pdf_u, grad_pdf_l, pdf_u, pdf_l, cdf_u, cdf_l;
censor_type = CensoringType::kIntervalCensored;
if (std::isinf(y_upper)) { // right-censored
pdf_u = 0;
cdf_u = 1;
grad_pdf_u = 0;
censor_type = CensoringType::kRightCensored;
} else { // interval-censored or left-censored
z_u = (log_y_upper - y_pred) / sigma;
pdf_u = dist_->PDF(z_u);
@@ -122,6 +237,7 @@ double AFTLoss::Hessian(double y_lower, double y_upper, double y_pred, double si
pdf_l = 0;
cdf_l = 0;
grad_pdf_l = 0;
censor_type = CensoringType::kLeftCensored;
} else { // interval-censored or right-censored
z_l = (log_y_lower - y_pred) / sigma;
pdf_l = dist_->PDF(z_l);
@@ -131,15 +247,17 @@ double AFTLoss::Hessian(double y_lower, double y_upper, double y_pred, double si
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;
const double sqrt_denominator = sigma * cdf_diff;
z_sign = (z_u > 0 || z_l > 0);
numerator = -(cdf_diff * grad_diff - pdf_diff * pdf_diff);
denominator = sqrt_denominator * sqrt_denominator;
}
hessian = numerator / denominator;
if (denominator < kEps && (std::isnan(hessian) || std::isinf(hessian))) {
hessian = GetLimitAtInfPred(dist_type_, censor_type, z_sign, sigma).GetHess();
}
return hessian;
return Clip(hessian, kMinHessian, kMaxHessian);
}
} // namespace common

9
src/common/survival_util.h
View File

@@ -42,15 +42,16 @@ struct AFTParam : public XGBoostParameter<AFTParam> {
class AFTLoss {
private:
std::unique_ptr<ProbabilityDistribution> dist_;
ProbabilityDistributionType dist_type_;
public:
/*!
* \brief Constructor for AFT loss function
* \param dist Choice of probability distribution for the noise term in AFT
* \param dist_type Choice of probability distribution for the noise term in AFT
*/
explicit AFTLoss(ProbabilityDistributionType dist) {
dist_.reset(ProbabilityDistribution::Create(dist));
}
explicit AFTLoss(ProbabilityDistributionType dist_type)
: dist_(ProbabilityDistribution::Create(dist_type)),
dist_type_(dist_type) {}
public:
/*!