Sampling a non-trivial survival geometry

This note can be considered as a complement to the survival analysis walkthrough, commenting on the non-trivial sampling procedure.

The setup #

We consider here the exact same setup as in the survival hard drive analysis. We consider a fleet of hard drives, in the context of survival analysis. Each hard drive ¹ has either an event (failure/death) or is censored. Many entries are left/right truncated. We call each continuous episode of “observed” drive a spell.

The model in question is given by

$$ \begin{split} a_m &\sim \mu_j + \tau u_m \end{split} $$

with $j$ - the manufacturer’s index. So $\mu_j$ - the model’s grand mean, tight to the manufacturer, from which models are allowed to deviate. $u_m$, $\tau$ determine how far they’re allowed to deviate.

Similarly, for $b_m$, we have

$$ b_m = \nu + \rho v_m $$

with $\rho v_m$ playing the same role as $\tau u_m$ - how much the per-model parameter can deviate from the grand mean, in this case, $\nu$. Note that there’s is no manufacturer’s level sharing. The reason is the actual meaning of the $k$ and $\lambda$. $\lambda$ - the scale parameter represents how fast/late the object fails.

The Problem #

The model does not contain a high number of levels. The problem is the exact likelihood fitting/describing the model and the model in question - the Weibull model, containing an exponential and two degrees of freedom, giving much room to the parameter space and uncertainty, in particular, with the long-living and highly truncated hard drives’ spells. Meaning drives live for a long time, fail rarely with a short observing time window. Consequently, this (but not only) leads to sampling problem, in particular to high values than r-hat $\hat{R}$.

Diagnostics - $\hat{R}$ #

When the specified model has a complex geometry, with typical examples including hierarchical models (inducing funnel-like geometry) and/or Gaussian multimodelity, the sampler fails explore the whole posterior. The typical indicator for that is the $\hat{R}$. It is defined (unformally) to be the ratio of two variances:

$$ \hat{R} \sim \frac{\text{between chain Var.}}{\text{within chain Var.}} $$

It measures how much the in-between chain variances exceeds what you expect from the within-chain variance. This means, that if chains have properly mixed (chain #i mixed well with chain #j, i$\neq$j), they are indistinguishable from the bulk of these chains themselves. Thus the ratio is around 1. On the other hand, if one chain has not mixed with the rest, the variance in the numerator will be greater than in the denominator, leading to $\hat{R}$’s significantly higher than 1.

We clearly see a bulk, where chains are mixed, and an “outlier” chain, that looks like independent. This is exactly what the r-hat is used for diagnostics.

Identifying the cause #

Let’s recall the building blocks of the Weibull model. The main relations “generating functions” (in the sence that they fully specify the model) of the model are given by

$$ S(t) = \exp\Biggl[ -\biggl( \frac{t}{\lambda} \biggr)^k \Biggr]\quad ; \qquad H(t) = \biggl( \frac{t}{\lambda} \biggr)^{k} $$

• $\lambda$ - the scale dictating how fast the items “die”.
• $k$ - the shape describing how the hazard changes with item’s age.

We note the physical meaning of the two parameters. The scale $k$ varies around $1$, with $k<1$ meaning a decreasing hazard rate (infant mortality phenomenon), and $k>1$ suggesting a increase of hazard rate with time (aging process). $\lambda$ tells us how long we expect for a module to live on average $S(\lambda) = e^{-1} \approx 0.37$. So at $\lambda$, we expect that there is $e^{-1}$ probability that a unit survives past $\lambda$. This means $\lambda$ has time units.
For ease of work, we also define transformed quantities $a\coloneqq \log \lambda $ and $b \coloneqq \log k$.
As the end goal, we would like to obtain the parameters’ estimates to later infer e.g. the survival time for some time $S(t_0)$.

Although the main “level” of the data are models, the data is highly heterogeneous at this exact level. Meaning that different models have a very different number of datapoints, sometimes orders of difference, which for survival data, that is in addition truncated, is a heavy bottleneck. For example, the model HGST HUH721212ALE604 has around $14\text{ k}$ spells, with around $900$ events fed into the model. In contrast, HGST HUS728T8TALE6L4 has 2 spells with 0 events occured. This is true that these two particular models share a common prior that enables pooling. First, this is not always the case, as there are many different manufacturers and, sometimes, such a significant difference cannot always be compensated by a moderately tight prior.

The backblaze date is sparse in failure events. Thus, if we see only some 0.05 to have actually, fail the actual value of the $\lambda$ we want to estimate is very far away, as remember, $\lambda$ is where only $\approx 0.37$ are still alive! This is what we will be referring to as the “extrapolation” issue. Indeed, a model never reaching the actual value of $\lambda$, will need to “guess” or “extrapolate” based on small number of evidence, with failure rates of e.g. $<0.1$.

Consider some given $t_0 = 365 \text{ days}$, at which we’re probing the experiment, and the survival rate at that time is e.g. $0.95$. Then, the following combinations $(k,\lambda)$ fullfill our “probing”:

(k)	Approximate $\lambda$	(S(4))
0.5	380 years	0.95
0.75	52 years	0.95
1.00	19.5 years	0.95
2.00	4 years	0.95

The scales have magnitudes of difference yet all agree at 1y. A small uncertainty in $k$ creates a huge uncertainty in $\lambda$. In addition, if we have few observed events, we have limited info about the shape $k$ (as $k$ dictates how likely a unit can fail throughout its lifetime). As a result, many uncertainty around $k$ (due to few datapoints) $\rightarrow$ huge uncertainty in $\lambda$².

Due to the nature of our dataset (that has been truncated and has not been studied for all possible edge cases), we’re often studying hard drives within a short time window. Then the likelihood cannot properly distinguish them at the heavily extrapolated period of time.

Attempting to fix #

The Hazard curve in the Weibull model is actually a line on the log-log scale:

$$ \log H(t) = \log -\log [S(t)] = k\log t - k\log \lambda $$

where $k$ is the slope and $ - k\log \lambda$ - the intercept. They can be used to define the line through two points ($x\equiv 0$ and $y\equiv 0$).

Based on the previously introduced assumptions, let’s define

$$ q = \log[ - \log S(t_0)] = k\log(t_0) - k\log \lambda $$

Then our line goes through the point $(\log(t_0), q)$, so $q$ becomes the height of the line at some “common” $t_0$.

This line is the same, but we have identified some “common” reference time $t_0$, that intersects with the likelihood. By rearranging, we can confirm, what we just saw about the uncertainty:

$$ \log \lambda = \log t_0 - q/k $$

meaning when $k$ is highly uncertain, the inferred $\lambda$ becomes highly unstable - $k \rightarrow 0$ $\implies$ $-q/k \rightarrow \infty$. Confirming the two are highly related. In fact, in the sampled $(a,b)$ space, this creates a banana-like sampling shape, given by

$$ a = \log t_0 - qe^{-b} $$

and as a consequence smaller $b$’s are compensated by larger $a$’s, producing heavy dependence and as a result cuved, banana-like ridge.

In order to attempt to control this high-uncertainty in extrapolated $\lambda$, we include the $q$ instead of the $\lambda$ in the Weibull equation for the cumulative hazard rate,

$$ \log H(t) = q + e^{b} \log(t/t_0) $$

By replacing $q$ with its definition via $\lambda$, we can reconstruct the equation for the initial Weibull model. With that, we confirm that we haven’t changed anything in the model - we introduced a new variable (from which we can reconstruct the initial variable $\lambda$), which we will be sampling instead. $\lambda$ is then reconstructed from the $q$ and the (arbitrarily) introduced $t_0$.

Thus, to summarize, the idea is to replace $\lambda$ with another, related (scaled) variable $q$, so that the main bulk (main part) of the data has high likelihood around $q$, which then becomes the sampled variable. This does not remove the pair $(\lambda, k)$ uncertainty, yet it does partially remove the issue with extrapolation/uncertainty propagation. As a result, partially removes chains getting stucked in “distant maxima”.

Verifying the results #

The main problem that we’re tackling are different chains exploring different regions. As explained before, the fix, that introduced the new sampled parameter $q$ does not change the model. Instead, it re-writes the whole model using $q$, which is more identifiable with our data (no significant extrapolation needed anymore).

The initial model ($\lambda$ - only, no $q$) is not stable. That is, sometimes all of the chains successfully converge towards the same target.

It is true, that if we remove some of the spells, the overall diagnostics will be healthier, as lower number of spells is often but not always associated with high $\hat{R}$.

A good (not the only) example is the model ST12000NM001G, which, although has many entries and events, has high $\hat{R}$ value, that comes from one chain being stuck in a local maxima.

Let’s illustrate this effect

Conclusion #

We departed from a highly censored (multilevel) survival setup, assumed to be Weibull-like. We omitted the (lengthy but vital) steps of model validation via the prior-predictive check. We, however, focused on the sampling issue, which turned out to potentially be universal for problems of similar class. In short, the two issues are the extrapolation of $\lambda$ and the compensation of $k$, $\lambda$. Using the introduced trick, we were able to partly solve the problem, without changing or simplifying the model.