Prediction of regression coefficients with XGBoost

Question

I am doing survival analysis. There is a dataset of items (id, group_id, observed lifetime, censorship status), each item belongs to a certain group. Each item is unique, but this uniqueness is unobservable, unpredictable and uncontollable, so we cannot base our predictions on it. Let's call that dataset items dataset. Each group has features easily obtainable as a separate dataset of groups (group_id, group_features). Let's call that dataset groups dataset.

One of the main goals of the research is to be able to predict expected lifetime of items of a certain group from its features and to be able to explain why (I utilise SHAP for that). In fact order predictions make more sense, but it is the next stage of the research. At the current stage I just want pretty accurate lifetime expectations.

I have combined the data from items dataset by gid and got a set of subdatasets of failure times belonging to each group. I have fitted a Weibull regression against each of those subdatasets of and got ρ and λ modelling survival of each group. I call this operation an aggregation. It maps items dataset to the dataset (gid, ρ, λ, #items with that gid in items dataset). Then I add there the data from groups dataset.

Then I try to fit a model predicting ρ and λ with XGBoost using the objective reg:squarederror and count of items in each group as a weight.

Then I predict the expected lifetime of an item of a group using the formula

$$E(Weibull(ρ=XGB<ρ>(...), λ=XGB<λ>(...)))$$,

where XGB<ρ>(...) means calling the trained XGBoost model for the column ρ passing it the group's feature vector.

Unfortunately,

• CV error (only XGBoost one) is tremendous even after 10000 iterations of hyperparameter optimization

• the composite model predicts values which are completely mad. For example typical non-censored lifetime of an item is about 1500 days, but the composite model predicts 2 days as an expectation.

  1. Am I doing the right thing? Is it possible at all to predict regression coefficients and then use them to predict expectation? Should I use count of items in groups as weights?

  2. Should I modify XGBoost objective plugging there hessian and gradient of expectation of Weibull distribution and try to predict the expectation directly by feeding it directly with items dataset augmented with the data from groups dataset and having a gradient for each item individually, rather than just predicting Weibull distribution coefficients and then expectation from them? Can I expect them to have the comparable performance?

Answer 1

As I understand the question, you have information on survival of individuals as a function of group membership, but don't have detailed information about the features associated with each individual. You do, however, have information about covariates associated with each of the groups.

First, you've done separate Weibull models on each group, which (as @adibender noted in comments) is simply modeling the baseline hazard for each group. Although Weibull modeling can often be useful, it's not clear whether you have even documented that Weibull models are appropriate for your data. For example, Weibull models don't allow for non-monotonic hazards as a function of time, so one potential explanation for your difficulty is that Weibull modeling is not appropriate for your data.

Second, parametric Weibull modeling of survival data, even when appropriate, is typically done with all cases at once. The model would identify a value of $\lambda$ common to all groups along with the differences among groups with respect to ρ. I can't prove that your approach (separate Weibull models for each group and then modeling the Weibull coefficients as a function of covariate value associated with groups) is necessarily wrong, but it doesn't seem to comport with the way that survival analysis in typically done.

If your interest is in the relations of the covariates to outcome, I see two ways to proceed.

One way would be modeling survival as a function of group first, then modeling group membership as a function of covariates. If all the information you have about survival is directly related to group membership, then the most direct survival model would be based on group membership. You then identify covariates most associated with the longest/shortest-survival groups.

The second way to proceed would be to accept the limitations of your data and simply assign to each individual the covariate values associated with that individual's group. That would remove the group membership from the analysis and allow direct (although clearly imprecise) modeling of how the covariates are related to outcome. If your group dataset has information on the probabilities of features associated with each group, you could treat this as a type of missing-data problem, using multiple imputation to obtain several sets of imputed data and thus incorporate the uncertainty in the individual-covariate associations.

With either of those approaches you would not have to use a parametric model, Weibull or otherwise. You would be able to test the appropriateness of the models (e.g, test for proportional hazards) in a way that your present approach doesn't seem to allow.