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/tests/cpp/common/test_probability_distribution.cc
+++ b/tests/cpp/common/test_probability_distribution.cc
@@ -0,0 +1,121 @@
+/*!
+ * Copyright (c) by Contributors 2020
+ */
+#include <gtest/gtest.h>
+#include <memory>
+#include <cmath>
+
+#include "xgboost/logging.h"
+#include "../../../src/common/probability_distribution.h"
+
+namespace xgboost {
+namespace common {
+
+TEST(ProbabilityDistribution, DistributionGeneric) {
+  // Assert d/dx CDF = PDF, d/dx PDF = GradPDF, d/dx GradPDF = HessPDF
+  // Do this for every distribution type
+  for (auto type : {ProbabilityDistributionType::kNormal, ProbabilityDistributionType::kLogistic,
+                    ProbabilityDistributionType::kExtreme}) {
+    std::unique_ptr<ProbabilityDistribution> dist{ ProbabilityDistribution::Create(type) };
+    double integral_of_pdf = dist->CDF(-2.0);
+    double integral_of_grad_pdf = dist->PDF(-2.0);
+    double integral_of_hess_pdf = dist->GradPDF(-2.0);
+    // Perform numerical differentiation and integration
+    // Enumerate 4000 grid points in range [-2, 2]
+    for (int i = 0; i <= 4000; ++i) {
+      const double x = static_cast<double>(i) / 1000.0 - 2.0;
+      // Numerical differentiation (p. 246, Numerical Analysis 2nd ed. by Timothy Sauer)
+      EXPECT_NEAR((dist->CDF(x + 1e-5) - dist->CDF(x - 1e-5)) / 2e-5, dist->PDF(x), 6e-11);
+      EXPECT_NEAR((dist->PDF(x + 1e-5) - dist->PDF(x - 1e-5)) / 2e-5, dist->GradPDF(x), 6e-11);
+      EXPECT_NEAR((dist->GradPDF(x + 1e-5) - dist->GradPDF(x - 1e-5)) / 2e-5,
+                  dist->HessPDF(x), 6e-11);
+      // Numerical integration using Trapezoid Rule (p. 257, Sauer)
+      integral_of_pdf += 5e-4 * (dist->PDF(x - 1e-3) + dist->PDF(x));
+      integral_of_grad_pdf += 5e-4 * (dist->GradPDF(x - 1e-3) + dist->GradPDF(x));
+      integral_of_hess_pdf += 5e-4 * (dist->HessPDF(x - 1e-3) + dist->HessPDF(x));
+      EXPECT_NEAR(integral_of_pdf, dist->CDF(x), 2e-4);
+      EXPECT_NEAR(integral_of_grad_pdf, dist->PDF(x), 2e-4);
+      EXPECT_NEAR(integral_of_hess_pdf, dist->GradPDF(x), 2e-4);
+    }
+  }
+}
+
+TEST(ProbabilityDistribution, NormalDist) {
+  std::unique_ptr<ProbabilityDistribution> dist{
+    ProbabilityDistribution::Create(ProbabilityDistributionType::kNormal)
+  };
+
+  // "Three-sigma rule" (https://en.wikipedia.org/wiki/68–95–99.7_rule)
+  //   68% of values are within 1 standard deviation away from the mean
+  //   95% of values are within 2 standard deviation away from the mean
+  // 99.7% of values are within 3 standard deviation away from the mean
+  EXPECT_NEAR(dist->CDF(0.5) - dist->CDF(-0.5), 0.3829, 0.00005);
+  EXPECT_NEAR(dist->CDF(1.0) - dist->CDF(-1.0), 0.6827, 0.00005);
+  EXPECT_NEAR(dist->CDF(1.5) - dist->CDF(-1.5), 0.8664, 0.00005);
+  EXPECT_NEAR(dist->CDF(2.0) - dist->CDF(-2.0), 0.9545, 0.00005);
+  EXPECT_NEAR(dist->CDF(2.5) - dist->CDF(-2.5), 0.9876, 0.00005);
+  EXPECT_NEAR(dist->CDF(3.0) - dist->CDF(-3.0), 0.9973, 0.00005);
+  EXPECT_NEAR(dist->CDF(3.5) - dist->CDF(-3.5), 0.9995, 0.00005);
+  EXPECT_NEAR(dist->CDF(4.0) - dist->CDF(-4.0), 0.9999, 0.00005);
+}
+
+TEST(ProbabilityDistribution, LogisticDist) {
+  std::unique_ptr<ProbabilityDistribution> dist{
+    ProbabilityDistribution::Create(ProbabilityDistributionType::kLogistic)
+  };
+
+  /**
+   * Enforce known properties of the logistic distribution.
+   * (https://en.wikipedia.org/wiki/Logistic_distribution)
+   **/
+
+  // Enumerate 4000 grid points in range [-2, 2]
+  for (int i = 0; i <= 4000; ++i) {
+    const double x = static_cast<double>(i) / 1000.0 - 2.0;
+    // PDF = 1/4 * sech(x/2)**2
+    const double sech_x = 1.0 / std::cosh(x * 0.5);  // hyperbolic secant at x/2
+    EXPECT_NEAR(0.25 * sech_x * sech_x, dist->PDF(x), 1e-15);
+    // CDF = 1/2 + 1/2 * tanh(x/2)
+    EXPECT_NEAR(0.5 + 0.5 * std::tanh(x * 0.5), dist->CDF(x), 1e-15);
+  }
+}
+
+TEST(ProbabilityDistribution, ExtremeDist) {
+  std::unique_ptr<ProbabilityDistribution> dist{
+    ProbabilityDistribution::Create(ProbabilityDistributionType::kExtreme)
+  };
+
+  /**
+   * Enforce known properties of the extreme distribution (also known as Gumbel distribution).
+   * The mean is the negative of the Euler-Mascheroni constant.
+   * The variance is 1/6 * pi**2. (https://mathworld.wolfram.com/GumbelDistribution.html)
+   **/
+
+  // Enumerate 25000 grid points in range [-20, 5].
+  // Compute the mean (expected value) of the distribution using numerical integration.
+  // Nearly all mass of the extreme distribution is concentrated between -20 and 5,
+  // so numerically integrating x*PDF(x) over [-20, 5] gives good estimate of the mean.
+  double mean = 0.0;
+  for (int i = 0; i <= 25000; ++i) {
+    const double x = static_cast<double>(i) / 1000.0 - 20.0;
+    // Numerical integration using Trapezoid Rule (p. 257, Sauer)
+    mean += 5e-4 * ((x - 1e-3) * dist->PDF(x - 1e-3) + x * dist->PDF(x));
+  }
+  EXPECT_NEAR(mean, -probability_constant::kEulerMascheroni, 1e-7);
+
+  // Enumerate 25000 grid points in range [-20, 5].
+  // Compute the variance of the distribution using numerical integration.
+  // Nearly all mass of the extreme distribution is concentrated between -20 and 5,
+  // so numerically integrating (x-mean)*PDF(x) over [-20, 5] gives good estimate of the variance.
+  double variance = 0.0;
+  for (int i = 0; i <= 25000; ++i) {
+    const double x = static_cast<double>(i) / 1000.0 - 20.0;
+    // Numerical integration using Trapezoid Rule (p. 257, Sauer)
+    variance += 5e-4 * ((x - 1e-3 - mean) * (x - 1e-3 - mean) * dist->PDF(x - 1e-3)
+                        + (x - mean) * (x - mean) * dist->PDF(x));
+  }
+  EXPECT_NEAR(variance, probability_constant::kPI * probability_constant::kPI / 6.0, 1e-6);
+}
+
+} // namespace common
+}  // namespace xgboost
--- a/tests/cpp/metric/test_survival_metric.cc
+++ b/tests/cpp/metric/test_survival_metric.cc
@@ -0,0 +1,169 @@
+/*!
+ * Copyright (c) by Contributors 2020
+ */
+#include <gtest/gtest.h>
+#include <memory>
+#include <vector>
+#include <string>
+#include <limits>
+#include <cmath>
+
+#include "xgboost/metric.h"
+#include "xgboost/logging.h"
+#include "../helpers.h"
+#include "../../../src/common/survival_util.h"
+
+namespace xgboost {
+namespace common {
+
+/**
+ * Reference values obtained from
+ * https://github.com/avinashbarnwal/GSOC-2019/blob/master/AFT/R/combined_assignment.R
+ **/
+
+TEST(Metric, AFTNegLogLik) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT metric is CPU only
+
+  /**
+   * Test aggregate output from the AFT metric over a small test data set.
+   * This is unlike AFTLoss.* tests, which verify metric values over individual data points.
+   **/
+  MetaInfo info;
+  info.num_row_ = 4;
+  info.labels_lower_bound_.HostVector()
+    = { 100.0f, -std::numeric_limits<bst_float>::infinity(), 60.0f, 16.0f };
+  info.labels_upper_bound_.HostVector()
+    = { 100.0f, 20.0f, std::numeric_limits<bst_float>::infinity(), 200.0f };
+  info.weights_.HostVector() = std::vector<bst_float>();
+  HostDeviceVector<bst_float> preds(4, std::log(64));
+
+  struct TestCase {
+    std::string dist_type;
+    bst_float reference_value;
+  };
+  for (const auto& test_case : std::vector<TestCase>{ {"normal", 2.1508f}, {"logistic", 2.1804f},
+                                                      {"extreme", 2.0706f} }) {
+    std::unique_ptr<Metric> metric(Metric::Create("aft-nloglik", &lparam));
+    metric->Configure({ {"aft_loss_distribution", test_case.dist_type},
+                        {"aft_loss_distribution_scale", "1.0"} });
+    EXPECT_NEAR(metric->Eval(preds, info, false), test_case.reference_value, 1e-4);
+  }
+}
+
+// Test configuration of AFT metric
+TEST(AFTNegLogLikMetric, Configuration) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT metric is CPU only
+  std::unique_ptr<Metric> metric(Metric::Create("aft-nloglik", &lparam));
+  metric->Configure({{"aft_loss_distribution", "normal"}, {"aft_loss_distribution_scale", "10"}});
+
+  // Configuration round-trip test
+  Json j_obj{ Object() };
+  metric->SaveConfig(&j_obj);
+  auto aft_param_json = j_obj["aft_loss_param"];
+  EXPECT_EQ(get<String>(aft_param_json["aft_loss_distribution"]), "normal");
+  EXPECT_EQ(get<String>(aft_param_json["aft_loss_distribution_scale"]), "10");
+}
+
+/**
+ * AFTLoss.* tests verify metric values over individual data points.
+ **/
+
+// Generate prediction value ranging from 2**1 to 2**15, using grid points in log scale
+// Then check prediction against the reference values
+static inline void CheckLossOverGridPoints(
+                      double true_label_lower_bound,
+                      double true_label_upper_bound,
+                      ProbabilityDistributionType dist_type,
+                      const std::vector<double>& reference_values) {
+  const int num_point = 20;
+  const double log_y_low = 1.0;
+  const double log_y_high = 15.0;
+  std::unique_ptr<AFTLoss> loss(new AFTLoss(dist_type));
+  CHECK_EQ(num_point, reference_values.size());
+  for (int i = 0; i < num_point; ++i) {
+    const double y_pred
+      = std::pow(2.0, i * (log_y_high - log_y_low) / (num_point - 1) + log_y_low);
+    const double loss_val
+      = loss->Loss(true_label_lower_bound, true_label_upper_bound, std::log(y_pred), 1.0);
+    EXPECT_NEAR(loss_val, reference_values[i], 1e-4);
+  }
+}
+
+TEST(AFTLoss, Uncensored) {
+  // Given label 100, compute the AFT loss for various prediction values
+  const double true_label_lower_bound = 100.0;
+  const double true_label_upper_bound = true_label_lower_bound;
+
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kNormal,
+    { 13.1761, 11.3085, 9.7017, 8.3558, 7.2708, 6.4466, 5.8833, 5.5808, 5.5392, 5.7585, 6.2386,
+      6.9795, 7.9813, 9.2440, 10.7675, 12.5519, 14.5971, 16.9032, 19.4702, 22.2980 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kLogistic,
+    { 8.5568, 8.0720, 7.6038, 7.1620, 6.7612, 6.4211, 6.1659, 6.0197, 5.9990, 6.1064, 6.3293,
+      6.6450, 7.0289, 7.4594, 7.9205, 8.4008, 8.8930, 9.3926, 9.8966, 10.4033 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kExtreme,
+    { 27.6310, 27.6310, 19.7177, 13.0281, 9.2183, 7.1365, 6.0916, 5.6688, 5.6195, 5.7941, 6.1031,
+      6.4929, 6.9310, 7.3981, 7.8827, 8.3778, 8.8791, 9.3842, 9.8916, 10.40033 });
+}
+
+TEST(AFTLoss, LeftCensored) {
+  // Given label (-inf, 20], compute the AFT loss for various prediction values
+  const double true_label_lower_bound = -std::numeric_limits<double>::infinity();
+  const double true_label_upper_bound = 20.0;
+
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kNormal,
+    { 0.0107, 0.0373, 0.1054, 0.2492, 0.5068, 0.9141, 1.5003, 2.2869, 3.2897, 4.5196, 5.9846,
+      7.6902, 9.6405, 11.8385, 14.2867, 16.9867, 19.9399, 23.1475, 26.6103, 27.6310 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kLogistic,
+    { 0.0953, 0.1541, 0.2451, 0.3804, 0.5717, 0.8266, 1.1449, 1.5195, 1.9387, 2.3902, 2.8636,
+      3.3512, 3.8479, 4.3500, 4.8556, 5.3632, 5.8721, 6.3817, 6.8918, 7.4021 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kExtreme,
+    { 0.0000, 0.0025, 0.0277, 0.1225, 0.3195, 0.6150, 0.9862, 1.4094, 1.8662, 2.3441, 2.8349,
+      3.3337, 3.8372, 4.3436, 4.8517, 5.3609, 5.8707, 6.3808, 6.8912, 7.4018 });
+}
+
+TEST(AFTLoss, RightCensored) {
+  // Given label [60, +inf), compute the AFT loss for various prediction values
+  const double true_label_lower_bound = 60.0;
+  const double true_label_upper_bound = std::numeric_limits<double>::infinity();
+
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kNormal,
+    { 8.0000, 6.2537, 4.7487, 3.4798, 2.4396, 1.6177, 0.9993, 0.5638, 0.2834, 0.1232, 0.0450,
+      0.0134, 0.0032, 0.0006, 0.0001, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kLogistic,
+    { 3.4340, 2.9445, 2.4683, 2.0125, 1.5871, 1.2041, 0.8756, 0.6099, 0.4083, 0.2643, 0.1668,
+      0.1034, 0.0633, 0.0385, 0.0233, 0.0140, 0.0084, 0.0051, 0.0030, 0.0018 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kExtreme,
+    { 27.6310, 18.0015, 10.8018, 6.4817, 3.8893, 2.3338, 1.4004, 0.8403, 0.5042, 0.3026, 0.1816,
+      0.1089, 0.0654, 0.0392, 0.0235, 0.0141, 0.0085, 0.0051, 0.0031, 0.0018 });
+}
+
+TEST(AFTLoss, IntervalCensored) {
+  // Given label [16, 200], compute the AFT loss for various prediction values
+  const double true_label_lower_bound = 16.0;
+  const double true_label_upper_bound = 200.0;
+  
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kNormal,
+    { 3.9746, 2.8415, 1.9319, 1.2342, 0.7335, 0.4121, 0.2536, 0.2470, 0.3919, 0.6982, 1.1825,
+      1.8622, 2.7526, 3.8656, 5.2102, 6.7928, 8.6183, 10.6901, 13.0108, 15.5826 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kLogistic,
+    { 2.2906, 1.8578, 1.4667, 1.1324, 0.8692, 0.6882, 0.5948, 0.5909, 0.6764, 0.8499, 1.1061,
+      1.4348, 1.8215, 2.2511, 2.7104, 3.1891, 3.6802, 4.1790, 4.6825, 5.1888 });
+  CheckLossOverGridPoints(true_label_lower_bound, true_label_upper_bound,
+    ProbabilityDistributionType::kExtreme,
+    { 8.0000, 4.8004, 2.8805, 1.7284, 1.0372, 0.6231, 0.3872, 0.3031, 0.3740, 0.5839, 0.8995,
+      1.2878, 1.7231, 2.1878, 2.6707, 3.1647, 3.6653, 4.1699, 4.6770, 5.1856 });
+}
+
+}  // namespace common
+}  // namespace xgboost
--- a/tests/cpp/objective/test_aft_obj.cc
+++ b/tests/cpp/objective/test_aft_obj.cc
@@ -0,0 +1,174 @@
+/*!
+ * Copyright (c) by Contributors 2020
+ */
+#include <gtest/gtest.h>
+#include <memory>
+#include <vector>
+#include <limits>
+#include <cmath>
+
+#include "xgboost/objective.h"
+#include "xgboost/logging.h"
+#include "../helpers.h"
+#include "../../../src/common/survival_util.h"
+
+namespace xgboost {
+namespace common {
+
+TEST(Objective, AFTObjConfiguration) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT objective is CPU only
+  std::unique_ptr<ObjFunction> objective(ObjFunction::Create("survival:aft", &lparam));
+  objective->Configure({ {"aft_loss_distribution", "logistic"},
+                          {"aft_loss_distribution_scale", "5"} });
+
+  // Configuration round-trip test
+  Json j_obj{ Object() };
+  objective->SaveConfig(&j_obj);
+  EXPECT_EQ(get<String>(j_obj["name"]), "survival:aft");
+  auto aft_param_json = j_obj["aft_loss_param"];
+  EXPECT_EQ(get<String>(aft_param_json["aft_loss_distribution"]), "logistic");
+  EXPECT_EQ(get<String>(aft_param_json["aft_loss_distribution_scale"]), "5");
+}
+
+/**
+ * Verify that gradient pair (gpair) is computed correctly for various prediction values.
+ * Reference values obtained from
+ * https://github.com/avinashbarnwal/GSOC-2019/blob/master/AFT/R/combined_assignment.R
+ **/
+
+// Generate prediction value ranging from 2**1 to 2**15, using grid points in log scale
+// Then check prediction against the reference values
+static inline void CheckGPairOverGridPoints(
+                      ObjFunction* obj,
+                      bst_float true_label_lower_bound,
+                      bst_float true_label_upper_bound,
+                      const std::string& dist_type,
+                      const std::vector<bst_float>& expected_grad,
+                      const std::vector<bst_float>& expected_hess,
+                      float ftol = 1e-4f) {
+  const int num_point = 20;
+  const double log_y_low = 1.0;
+  const double log_y_high = 15.0;
+  
+  obj->Configure({ {"aft_loss_distribution", dist_type},
+                   {"aft_loss_distribution_scale", "1"} });
+
+  MetaInfo info;
+  info.num_row_ = num_point;
+  info.labels_lower_bound_.HostVector()
+    = std::vector<bst_float>(num_point, true_label_lower_bound);
+  info.labels_upper_bound_.HostVector()
+    = std::vector<bst_float>(num_point, true_label_upper_bound);
+  info.weights_.HostVector() = std::vector<bst_float>();
+  std::vector<bst_float> preds(num_point);
+  for (int i = 0; i < num_point; ++i) {
+    preds[i] = std::log(std::pow(2.0, i * (log_y_high - log_y_low) / (num_point - 1) + log_y_low));
+  }
+
+  HostDeviceVector<GradientPair> out_gpair;
+  obj->GetGradient(HostDeviceVector<bst_float>(preds), info, 1, &out_gpair);
+  const auto& gpair = out_gpair.HostVector();
+  CHECK_EQ(num_point, expected_grad.size());
+  CHECK_EQ(num_point, expected_hess.size());
+  for (int i = 0; i < num_point; ++i) {
+    EXPECT_NEAR(gpair[i].GetGrad(), expected_grad[i], ftol);
+    EXPECT_NEAR(gpair[i].GetHess(), expected_hess[i], ftol);
+  }
+}
+
+TEST(Objective, AFTObjGPairUncensoredLabels) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT objective is CPU only
+  std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &lparam));
+
+  CheckGPairOverGridPoints(obj.get(), 100.0f, 100.0f, "normal",
+    { -3.9120f, -3.4013f, -2.8905f, -2.3798f, -1.8691f, -1.3583f, -0.8476f, -0.3368f, 0.1739f,
+      0.6846f, 1.1954f, 1.7061f, 2.2169f, 2.7276f, 3.2383f, 3.7491f, 4.2598f, 4.7706f, 5.2813f,
+      5.7920f },
+    { 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f,
+      1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f, 1.0000f });
+  CheckGPairOverGridPoints(obj.get(), 100.0f, 100.0f, "logistic",
+    { -0.9608f, -0.9355f, -0.8948f, -0.8305f, -0.7327f, -0.5910f, -0.4001f, -0.1668f, 0.0867f,
+      0.3295f, 0.5354f, 0.6927f, 0.8035f, 0.8773f, 0.9245f, 0.9540f, 0.9721f, 0.9832f, 0.9899f,
+      0.9939f },
+    { 0.0384f, 0.0624f, 0.0997f, 0.1551f, 0.2316f, 0.3254f, 0.4200f, 0.4861f, 0.4962f, 0.4457f,
+      0.3567f, 0.2601f, 0.1772f, 0.1152f, 0.0726f, 0.0449f, 0.0275f, 0.0167f, 0.0101f, 0.0061f });
+  CheckGPairOverGridPoints(obj.get(), 100.0f, 100.0f, "extreme",
+    { -0.0000f, -29.0026f, -17.0031f, -9.8028f, -5.4822f, -2.8897f, -1.3340f, -0.4005f, 0.1596f,
+      0.4957f, 0.6974f, 0.8184f, 0.8910f, 0.9346f, 0.9608f, 0.9765f, 0.9859f, 0.9915f, 0.9949f,
+      0.9969f },
+    { 0.0000f, 30.0026f, 18.0031f, 10.8028f, 6.4822f, 3.8897f, 2.3340f, 1.4005f, 0.8404f, 0.5043f,
+      0.3026f, 0.1816f, 0.1090f, 0.0654f, 0.0392f, 0.0235f, 0.0141f, 0.0085f, 0.0051f, 0.0031f });
+}
+
+TEST(Objective, AFTObjGPairLeftCensoredLabels) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT objective is CPU only
+  std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &lparam));
+
+  CheckGPairOverGridPoints(obj.get(), -std::numeric_limits<float>::infinity(), 20.0f, "normal",
+    { 0.0285f, 0.0832f, 0.1951f, 0.3804f, 0.6403f, 0.9643f, 1.3379f, 1.7475f, 2.1828f, 2.6361f,
+      3.1023f, 3.5779f, 4.0603f, 4.5479f, 5.0394f, 5.5340f, 6.0309f, 6.5298f, 7.0303f, 0.5072f },
+    { 0.0663f, 0.1559f, 0.2881f, 0.4378f, 0.5762f, 0.6878f, 0.7707f, 0.8300f, 0.8719f, 0.9016f,
+      0.9229f, 0.9385f, 0.9501f, 0.9588f, 0.9656f, 0.9709f, 0.9751f, 0.9785f, 0.9812f, 0.0045f },
+    2e-4);
+  CheckGPairOverGridPoints(obj.get(), -std::numeric_limits<float>::infinity(), 20.0f, "logistic",
+    { 0.0909f, 0.1428f, 0.2174f, 0.3164f, 0.4355f, 0.5625f, 0.6818f, 0.7812f, 0.8561f, 0.9084f,
+      0.9429f, 0.9650f, 0.9787f, 0.9871f, 0.9922f, 0.9953f, 0.9972f, 0.9983f, 0.9990f, 0.9994f },
+    { 0.0826f, 0.1224f, 0.1701f, 0.2163f, 0.2458f, 0.2461f, 0.2170f, 0.1709f, 0.1232f, 0.0832f,
+      0.0538f, 0.0338f, 0.0209f, 0.0127f, 0.0077f, 0.0047f, 0.0028f, 0.0017f, 0.0010f, 0.0006f });
+  CheckGPairOverGridPoints(obj.get(), -std::numeric_limits<float>::infinity(), 20.0f, "extreme",
+    { 0.0005f, 0.0149f, 0.1011f, 0.2815f, 0.4881f, 0.6610f, 0.7847f, 0.8665f, 0.9183f, 0.9504f,
+      0.9700f, 0.9820f, 0.9891f, 0.9935f, 0.9961f, 0.9976f, 0.9986f, 0.9992f, 0.9995f, 0.9997f },
+    { 0.0041f, 0.0747f, 0.2731f, 0.4059f, 0.3829f, 0.2901f, 0.1973f, 0.1270f, 0.0793f, 0.0487f,
+      0.0296f, 0.0179f, 0.0108f, 0.0065f, 0.0039f, 0.0024f, 0.0014f, 0.0008f, 0.0005f, 0.0003f });
+}
+
+TEST(Objective, AFTObjGPairRightCensoredLabels) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT objective is CPU only
+  std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &lparam));
+
+  CheckGPairOverGridPoints(obj.get(), 60.0f, std::numeric_limits<float>::infinity(), "normal",
+    { -3.6583f, -3.1815f, -2.7135f, -2.2577f, -1.8190f, -1.4044f, -1.0239f, -0.6905f, -0.4190f,
+      -0.2209f, -0.0973f, -0.0346f, -0.0097f, -0.0021f, -0.0004f, -0.0000f, -0.0000f, -0.0000f,
+      -0.0000f, -0.0000f },
+    { 0.9407f, 0.9259f, 0.9057f, 0.8776f, 0.8381f, 0.7821f, 0.7036f, 0.5970f, 0.4624f, 0.3128f,
+      0.1756f, 0.0780f, 0.0265f, 0.0068f, 0.0013f, 0.0002f, 0.0000f, 0.0000f, 0.0000f, 0.0000f });
+  CheckGPairOverGridPoints(obj.get(), 60.0f, std::numeric_limits<float>::infinity(), "logistic",
+    { -0.9677f, -0.9474f, -0.9153f, -0.8663f, -0.7955f, -0.7000f, -0.5834f, -0.4566f, -0.3352f,
+      -0.2323f, -0.1537f, -0.0982f, -0.0614f, -0.0377f, -0.0230f, -0.0139f, -0.0084f, -0.0051f,
+      -0.0030f, -0.0018f },
+    { 0.0312f, 0.0499f, 0.0776f, 0.1158f, 0.1627f, 0.2100f, 0.2430f, 0.2481f, 0.2228f, 0.1783f,
+      0.1300f, 0.0886f, 0.0576f, 0.0363f, 0.0225f, 0.0137f, 0.0083f, 0.0050f, 0.0030f, 0.0018f });
+  CheckGPairOverGridPoints(obj.get(), 60.0f, std::numeric_limits<float>::infinity(), "extreme",
+    { -2.8073f, -18.0015f, -10.8018f, -6.4817f, -3.8893f, -2.3338f, -1.4004f, -0.8403f, -0.5042f,
+      -0.3026f, -0.1816f, -0.1089f, -0.0654f, -0.0392f, -0.0235f, -0.0141f, -0.0085f, -0.0051f,
+       -0.0031f, -0.0018f },
+    { 0.2614f, 18.0015f, 10.8018f, 6.4817f, 3.8893f, 2.3338f, 1.4004f, 0.8403f, 0.5042f, 0.3026f,
+      0.1816f, 0.1089f, 0.0654f, 0.0392f, 0.0235f, 0.0141f, 0.0085f, 0.0051f, 0.0031f, 0.0018f });
+}
+
+TEST(Objective, AFTObjGPairIntervalCensoredLabels) {
+  auto lparam = CreateEmptyGenericParam(-1);  // currently AFT objective is CPU only
+  std::unique_ptr<ObjFunction> obj(ObjFunction::Create("survival:aft", &lparam));
+
+  CheckGPairOverGridPoints(obj.get(), 16.0f, 200.0f, "normal",
+    { -2.4435f, -1.9965f, -1.5691f, -1.1679f, -0.7990f, -0.4649f, -0.1596f, 0.1336f, 0.4370f,
+      0.7682f, 1.1340f, 1.5326f, 1.9579f, 2.4035f, 2.8639f, 3.3351f, 3.8143f, 4.2995f, 4.7891f,
+      5.2822f },
+    { 0.8909f, 0.8579f, 0.8134f, 0.7557f, 0.6880f, 0.6221f, 0.5789f, 0.5769f, 0.6171f, 0.6818f,
+      0.7500f, 0.8088f, 0.8545f, 0.8884f, 0.9131f, 0.9312f, 0.9446f, 0.9547f, 0.9624f, 0.9684f });
+  CheckGPairOverGridPoints(obj.get(), 16.0f, 200.0f, "logistic",
+    { -0.8790f, -0.8112f, -0.7153f, -0.5893f, -0.4375f, -0.2697f, -0.0955f, 0.0800f, 0.2545f,
+      0.4232f, 0.5768f, 0.7054f, 0.8040f, 0.8740f, 0.9210f, 0.9513f, 0.9703f, 0.9820f, 0.9891f,
+      0.9934f },
+    { 0.1086f, 0.1588f, 0.2176f, 0.2745f, 0.3164f, 0.3374f, 0.3433f, 0.3434f, 0.3384f, 0.3191f,
+      0.2789f, 0.2229f, 0.1637f, 0.1125f, 0.0737f, 0.0467f, 0.0290f, 0.0177f, 0.0108f, 0.0065f });
+  CheckGPairOverGridPoints(obj.get(), 16.0f, 200.0f, "extreme",
+    { -8.0000f, -4.8004f, -2.8805f, -1.7284f, -1.0371f, -0.6168f, -0.3140f, -0.0121f, 0.2841f,
+      0.5261f, 0.6989f, 0.8132f, 0.8857f, 0.9306f, 0.9581f, 0.9747f, 0.9848f, 0.9909f, 0.9945f,
+      0.9967f },
+    { 8.0000f, 4.8004f, 2.8805f, 1.7284f, 1.0380f, 0.6567f, 0.5727f, 0.6033f, 0.5384f, 0.4051f,
+      0.2757f, 0.1776f, 0.1110f, 0.0682f, 0.0415f, 0.0251f, 0.0151f, 0.0091f, 0.0055f, 0.0033f });
+}
+
+}  // namespace common
+}  // namespace xgboost