Answering my own question: apparently this is indeed doable and already researched. I kind of stumbled upon things over time, but have not tried any in production as I have moved on to doing other things.
Quantile Random Forests
The simplest answer to my question is probably Quantile Random Forests. Here is a readable blog post with references to papers, as well as R and Python implementations. The rough idea is to modify the random forest a little to remember the target values in the leaves and to derive quantiles from them at prediction time. Very nice!
Uncertainty estimating neural nets
The less simple but lovely (at least for my taste) idea, is to train two neural networks - one for mean and one for variance - that fit the conditional distribution. This one I had to try to believe that it works, so here's how it goes briefly:
Take some synthetic data with $x \sim \text{Uniform}[-1, 1]$, and $y | x \sim \text{Normal}(x, |x|)$:
And try to fit with two neural networks. They need not to be special, and here is a pytorch example with both combined into one for comfort:
class TwoOutputNet(nn.Module):
def __init__(self):
super().__init__()
self.m1 = nn.Linear(1, 10)
self.m2 = nn.Linear(10, 1)
self.s1 = nn.Linear(1, 10)
self.s2 = nn.Linear(10, 1)
def forward(self, x):
m = F.tanh(self.m1(x))
s = F.tanh(self.s1(x))
return self.m2(m), torch.clamp(self.s2(s), min=-5, max=5)
The key is the loss function that takes true value (y), estimated mean (m), estimated standard deviation (s) and returns a loss that penalizes both error weighted by uncertainty and the uncertainty. For example, the negative log likelihood of the normal distribution:
class NegativeLogLikelihood(nn.Module):
def __init__(self):
super().__init__()
def forward(self, y, m, s):
sq_part = ((y - m) / torch.exp(s)) ** 2
lg_part = s
return (1/2 * sq_part + lg_part).mean()
Another point to note, is that fitting for standard deviation directly seems to be hard (convergence is erratic for me), so it is worth experimenting with different parameterizations. The pasted code takes the exponent (that's why the second output of the net is clamped between -5 and 5 -- the standard deviation gets clamped between 0.005 and 150).
In the end we get the both the estimates of the mean and standard deviation at each value of x:
This idea and different parameterizations are explored more in a recent paper. It also notes, that the parametric assumptions, such as used above, are not necessary, and the second net in general can output the expected loss of the first one, which I think is very nice.
Bayesian neural nets
And then there is the whole Bayesian neural networks thing, which I did not investigate further. I quite liked this blog post for inspiration, and I guess edwardlib and this thesis are good starting points.
