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>
2020-03-25 16:52:51 -04:00
parent 1de36cdf1e
commit dcf439932a
21 changed files with 1789 additions and 15 deletions
--- a/src/c_api/c_api.cc
+++ b/src/c_api/c_api.cc
@@ -265,6 +265,10 @@ XGB_DLL int XGDMatrixGetFloatInfo(const DMatrixHandle handle,
    vec = &info.weights_.HostVector();
  } else if (!std::strcmp(field, "base_margin")) {
    vec = &info.base_margin_.HostVector();
+  } else if (!std::strcmp(field, "label_lower_bound")) {
+    vec = &info.labels_lower_bound_.HostVector();
+  } else if (!std::strcmp(field, "label_upper_bound")) {
+    vec = &info.labels_upper_bound_.HostVector();
  } else {
    LOG(FATAL) << "Unknown float field name " << field;
  }
@@ -284,8 +288,7 @@ XGB_DLL int XGDMatrixGetUIntInfo(const DMatrixHandle handle,
  if (!std::strcmp(field, "group_ptr")) {
    vec = &info.group_ptr_;
  } else {
-    LOG(FATAL) << "Unknown comp uint field name " << field
-      << " with comparison " << std::strcmp(field, "group_ptr");
+    LOG(FATAL) << "Unknown uint field name " << field;
  }
  *out_len = static_cast<xgboost::bst_ulong>(vec->size());
  *out_dptr = dmlc::BeginPtr(*vec);
--- a/src/common/probability_distribution.cc
+++ b/src/common/probability_distribution.cc
@@ -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
--- a/src/common/probability_distribution.h
+++ b/src/common/probability_distribution.h
@@ -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_
--- a/src/common/survival_util.cc
+++ b/src/common/survival_util.cc
@@ -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
--- a/src/common/survival_util.h
+++ b/src/common/survival_util.h
@@ -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_
--- a/src/data/data.cc
+++ b/src/data/data.cc
@@ -133,15 +133,17 @@ void MetaInfo::Clear() {
 /*
 * Binary serialization format for MetaInfo:
 *
- * | name        | type     | is_scalar | num_row | num_col | value           |
- * |-------------+----------+-----------+---------+---------+-----------------|
- * | num_row     | kUInt64  | True      | NA      |      NA | ${num_row_}     |
- * | num_col     | kUInt64  | True      | NA      |      NA | ${num_col_}     |
- * | num_nonzero | kUInt64  | True      | NA      |      NA | ${num_nonzero_} |
- * | labels      | kFloat32 | False     | ${size} |       1 | ${labels_}      |
- * | group_ptr   | kUInt32  | False     | ${size} |       1 | ${group_ptr_}   |
- * | weights     | kFloat32 | False     | ${size} |       1 | ${weights_}     |
- * | base_margin | kFloat32 | False     | ${size} |       1 | ${base_margin_} |
+ * | name               | type     | is_scalar | num_row | num_col | value                   |
+ * |--------------------+----------+-----------+---------+---------+-------------------------|
+ * | num_row            | kUInt64  | True      | NA      |      NA | ${num_row_}             |
+ * | num_col            | kUInt64  | True      | NA      |      NA | ${num_col_}             |
+ * | num_nonzero        | kUInt64  | True      | NA      |      NA | ${num_nonzero_}         |
+ * | labels             | kFloat32 | False     | ${size} |       1 | ${labels_}              |
+ * | group_ptr          | kUInt32  | False     | ${size} |       1 | ${group_ptr_}           |
+ * | weights            | kFloat32 | False     | ${size} |       1 | ${weights_}             |
+ * | base_margin        | kFloat32 | False     | ${size} |       1 | ${base_margin_}         |
+ * | labels_lower_bound | kFloat32 | False     | ${size} |       1 | ${labels_lower_bound__} |
+ * | labels_upper_bound | kFloat32 | False     | ${size} |       1 | ${labels_upper_bound__} |
 *
 * Note that the scalar fields (is_scalar=True) will have num_row and num_col missing.
 * Also notice the difference between the saved name and the name used in `SetInfo':
@@ -164,6 +166,10 @@ void MetaInfo::SaveBinary(dmlc::Stream *fo) const {
                  {weights_.Size(), 1}, weights_); ++field_cnt;
  SaveVectorField(fo, u8"base_margin", DataType::kFloat32,
                  {base_margin_.Size(), 1}, base_margin_); ++field_cnt;
+  SaveVectorField(fo, u8"labels_lower_bound", DataType::kFloat32,
+                  {labels_lower_bound_.Size(), 1}, labels_lower_bound_); ++field_cnt;
+  SaveVectorField(fo, u8"labels_upper_bound", DataType::kFloat32,
+                  {labels_upper_bound_.Size(), 1}, labels_upper_bound_); ++field_cnt;

  CHECK_EQ(field_cnt, kNumField) << "Wrong number of fields";
 }
@@ -195,6 +201,8 @@ void MetaInfo::LoadBinary(dmlc::Stream *fi) {
  LoadVectorField(fi, u8"group_ptr", DataType::kUInt32, &group_ptr_);
  LoadVectorField(fi, u8"weights", DataType::kFloat32, &weights_);
  LoadVectorField(fi, u8"base_margin", DataType::kFloat32, &base_margin_);
+  LoadVectorField(fi, u8"labels_lower_bound", DataType::kFloat32, &labels_lower_bound_);
+  LoadVectorField(fi, u8"labels_upper_bound", DataType::kFloat32, &labels_upper_bound_);
 }

 // try to load group information from file, if exists
@@ -268,8 +276,18 @@ void MetaInfo::SetInfo(const char* key, const void* dptr, DataType dtype, size_t
    for (size_t i = 1; i < group_ptr_.size(); ++i) {
      group_ptr_[i] = group_ptr_[i - 1] + group_ptr_[i];
    }
+  } else if (!std::strcmp(key, "label_lower_bound")) {
+    auto& labels = labels_lower_bound_.HostVector();
+    labels.resize(num);
+    DISPATCH_CONST_PTR(dtype, dptr, cast_dptr,
+                       std::copy(cast_dptr, cast_dptr + num, labels.begin()));
+  } else if (!std::strcmp(key, "label_upper_bound")) {
+    auto& labels = labels_upper_bound_.HostVector();
+    labels.resize(num);
+    DISPATCH_CONST_PTR(dtype, dptr, cast_dptr,
+                       std::copy(cast_dptr, cast_dptr + num, labels.begin()));
  } else {
-    LOG(FATAL) << "Unknown metainfo: " << key;
+    LOG(FATAL) << "Unknown key for MetaInfo: " << key;
  }
 }

--- a/src/metric/survival_metric.cc
+++ b/src/metric/survival_metric.cc
@@ -0,0 +1,106 @@
+/*!
+ * Copyright 2019 by Contributors
+ * \file survival_metric.cc
+ * \brief Metrics for survival analysis
+ * \author Avinash Barnwal, Hyunsu Cho and Toby Hocking
+ */
+
+#include <rabit/rabit.h>
+#include <xgboost/metric.h>
+#include <xgboost/host_device_vector.h>
+#include <dmlc/registry.h>
+#include <cmath>
+#include <memory>
+#include <vector>
+#include <limits>
+
+#include "xgboost/json.h"
+
+#include "../common/math.h"
+#include "../common/survival_util.h"
+
+using AFTParam = xgboost::common::AFTParam;
+using AFTLoss = xgboost::common::AFTLoss;
+
+namespace xgboost {
+namespace metric {
+// tag the this file, used by force static link later.
+DMLC_REGISTRY_FILE_TAG(survival_metric);
+
+/*! \brief Negative log likelihood of Accelerated Failure Time model */
+struct EvalAFT : public Metric {
+ public:
+  explicit EvalAFT(const char* param) {}
+
+  void Configure(const Args& args) override {
+    param_.UpdateAllowUnknown(args);
+    loss_.reset(new AFTLoss(param_.aft_loss_distribution));
+  }
+
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String(this->Name());
+    out["aft_loss_param"] = toJson(param_);
+  }
+
+  void LoadConfig(Json const& in) override {
+    fromJson(in["aft_loss_param"], &param_);
+  }
+
+  bst_float Eval(const HostDeviceVector<bst_float> &preds,
+                 const MetaInfo &info,
+                 bool distributed) override {
+    CHECK_NE(info.labels_lower_bound_.Size(), 0U)
+      << "y_lower cannot be empty";
+    CHECK_NE(info.labels_upper_bound_.Size(), 0U)
+      << "y_higher cannot be empty";
+    CHECK_EQ(preds.Size(), info.labels_lower_bound_.Size());
+    CHECK_EQ(preds.Size(), info.labels_upper_bound_.Size());
+
+    /* Compute negative log likelihood for each data point and compute weighted average */
+    const auto& yhat = preds.HostVector();
+    const auto& y_lower = info.labels_lower_bound_.HostVector();
+    const auto& y_upper = info.labels_upper_bound_.HostVector();
+    const auto& weights = info.weights_.HostVector();
+    const bool is_null_weight = weights.empty();
+    const float aft_loss_distribution_scale = param_.aft_loss_distribution_scale;
+    CHECK_LE(yhat.size(), static_cast<size_t>(std::numeric_limits<omp_ulong>::max()))
+      << "yhat is too big";
+    const omp_ulong nsize = static_cast<omp_ulong>(yhat.size());
+
+    double nloglik_sum = 0.0;
+    double weight_sum = 0.0;
+    #pragma omp parallel for default(none) \
+     firstprivate(nsize, is_null_weight, aft_loss_distribution_scale) \
+     shared(weights, y_lower, y_upper, yhat) reduction(+:nloglik_sum, weight_sum)
+    for (omp_ulong i = 0; i < nsize; ++i) {
+      // If weights are empty, data is unweighted so we use 1.0 everywhere
+      const double w = is_null_weight ? 1.0 : weights[i];
+      const double loss
+        = loss_->Loss(y_lower[i], y_upper[i], yhat[i], aft_loss_distribution_scale);
+      nloglik_sum += loss;
+      weight_sum += w;
+    }
+
+    double dat[2]{nloglik_sum, weight_sum};
+    if (distributed) {
+      rabit::Allreduce<rabit::op::Sum>(dat, 2);
+    }
+    return static_cast<bst_float>(dat[0] / dat[1]);
+  }
+
+  const char* Name() const override {
+    return "aft-nloglik";
+  }
+
+ private:
+  AFTParam param_;
+  std::unique_ptr<AFTLoss> loss_;
+};
+
+XGBOOST_REGISTER_METRIC(AFT, "aft-nloglik")
+.describe("Negative log likelihood of Accelerated Failure Time model.")
+.set_body([](const char* param) { return new EvalAFT(param); });
+
+}  // namespace metric
+}  // namespace xgboost
--- a/src/objective/aft_obj.cc
+++ b/src/objective/aft_obj.cc
@@ -0,0 +1,119 @@
+/*!
+ * Copyright 2015 by Contributors
+ * \file rank.cc
+ * \brief Definition of aft loss.
+ */
+
+#include <dmlc/omp.h>
+#include <xgboost/logging.h>
+#include <xgboost/objective.h>
+#include <vector>
+#include <limits>
+#include <algorithm>
+#include <memory>
+#include <utility>
+#include <cmath>
+
+#include "xgboost/json.h"
+
+#include "../common/math.h"
+#include "../common/random.h"
+#include "../common/survival_util.h"
+
+using AFTParam = xgboost::common::AFTParam;
+using AFTLoss = xgboost::common::AFTLoss;
+
+namespace xgboost {
+namespace obj {
+
+DMLC_REGISTRY_FILE_TAG(aft_obj);
+
+class AFTObj : public ObjFunction {
+ public:
+  void Configure(const std::vector<std::pair<std::string, std::string> >& args) override {
+    param_.UpdateAllowUnknown(args);
+    loss_.reset(new AFTLoss(param_.aft_loss_distribution));
+  }
+
+  void GetGradient(const HostDeviceVector<bst_float>& preds,
+                   const MetaInfo& info,
+                   int iter,
+                   HostDeviceVector<GradientPair>* out_gpair) override {
+    /* Boilerplate */
+    CHECK_EQ(preds.Size(), info.labels_lower_bound_.Size());
+    CHECK_EQ(preds.Size(), info.labels_upper_bound_.Size());
+
+    const auto& yhat = preds.HostVector();
+    const auto& y_lower = info.labels_lower_bound_.HostVector();
+    const auto& y_upper = info.labels_upper_bound_.HostVector();
+    const auto& weights = info.weights_.HostVector();
+    const bool is_null_weight = weights.empty();
+
+    out_gpair->Resize(yhat.size());
+    std::vector<GradientPair>& gpair = out_gpair->HostVector();
+    CHECK_LE(yhat.size(), static_cast<size_t>(std::numeric_limits<omp_ulong>::max()))
+      << "yhat is too big";
+    const omp_ulong nsize = static_cast<omp_ulong>(yhat.size());
+    const float aft_loss_distribution_scale = param_.aft_loss_distribution_scale;
+
+    #pragma omp parallel for default(none) \
+      firstprivate(nsize, is_null_weight, aft_loss_distribution_scale) \
+      shared(weights, y_lower, y_upper, yhat, gpair)
+    for (omp_ulong i = 0; i < nsize; ++i) {
+      // If weights are empty, data is unweighted so we use 1.0 everywhere
+      const double w = is_null_weight ? 1.0 : weights[i];
+      const double grad = loss_->Gradient(y_lower[i], y_upper[i],
+                                          yhat[i], aft_loss_distribution_scale);
+      const double hess = loss_->Hessian(y_lower[i], y_upper[i],
+                                         yhat[i], aft_loss_distribution_scale);
+      gpair[i] = GradientPair(grad * w, hess * w);
+    }
+  }
+
+  void PredTransform(HostDeviceVector<bst_float> *io_preds) override {
+    // Trees give us a prediction in log scale, so exponentiate
+    std::vector<bst_float> &preds = io_preds->HostVector();
+    const long ndata = static_cast<long>(preds.size()); // NOLINT(*)
+    #pragma omp parallel for default(none) firstprivate(ndata) shared(preds)
+    for (long j = 0; j < ndata; ++j) {  // NOLINT(*)
+      preds[j] = std::exp(preds[j]);
+    }
+  }
+
+  void EvalTransform(HostDeviceVector<bst_float> *io_preds) override {
+    // do nothing here, since the AFT metric expects untransformed prediction score
+  }
+
+  bst_float ProbToMargin(bst_float base_score) const override {
+    return std::log(base_score);
+  }
+
+  const char* DefaultEvalMetric() const override {
+    return "aft-nloglik";
+  }
+
+  void SaveConfig(Json* p_out) const override {
+    auto& out = *p_out;
+    out["name"] = String("survival:aft");
+    out["aft_loss_param"] = toJson(param_);
+  }
+
+  void LoadConfig(Json const& in) override {
+    fromJson(in["aft_loss_param"], &param_);
+    loss_.reset(new AFTLoss(param_.aft_loss_distribution));
+  }
+
+ private:
+  AFTParam param_;
+  std::unique_ptr<AFTLoss> loss_;
+};
+
+// register the objective functions
+XGBOOST_REGISTER_OBJECTIVE(AFTObj, "survival:aft")
+.describe("AFT loss function")
+.set_body([]() { return new AFTObj(); });
+
+}  // namespace obj
+}  // namespace xgboost
+
+