I want to estimate a “density field”, specifically $P(y|x, m)$, for binary labels $y$ associated with 2D points characterized by spatial coordinates $m$ and additional spatio-temporal features $x$. The model should aim to predict a probability at each specific point (a "pointwise probability field"), rather than over the entire space. In other words, the estimated probabilities at these points do not need to sum to 1 over the 2D space; in fact, their sum will likely be significantly larger than 1, except in some unfavourable setup of input 𝑥. I think this problem can also be seen as estimating the 3D surface, where x and y span over 2D space and z is between 0 and 1.

Although my dataset consists of the hundreds of thousands of data points, this might be still limited number to estimate the whole surface. Additionally, I am faced with the challenge of ensuring the smoothness of the model outputs. I wonder what would be the good problem setup, neural network architecture and loss function to promote smoothness in the output.


  1. What type of loss function and neural network architecture could be suitable for modelling a probability field that ensures smooth transitions in probability values across neighbouring points and similar feature sets, especially when only a limited set of points is available?
  2. I am looking for the correct terminology and areas of literature that address similar modelling challenges. While there's abundant research on distribution estimation (e.g., using Gaussian Mixture Models), it does not align with my need to estimate a density field. In fact, I’m not even sure if the “density field” or “pointwise probability field” are the correct terms for what I am looking for. What terms and problem specifications should I search in the literature that deal with similar problems?

I appreciate any guidance or references to similar work that can help frame and solve my problem effectively.

  • $\begingroup$ Not quite sure what the setup. You are talking about 2d, presumably $x,\,m$ and then describe $x$ as spatio-temporal features. $y$ is apparently binary. Why is what you are trying to do not the same as probability density estimation? $\endgroup$
    – Cryo
    Commented May 14 at 13:24
  • $\begingroup$ 2d point coordinates are represented by $m$, spatio-temporal context features of different objects that impact the binary target $y$ are represented by $x$. It's not the same as PDF estimation, because I don't look for a probability distribution over the continuous 2D space, where probabilities over all points belonging to 2D space need to sum up to 1. In fact, in my setup up, for a given $x$ and point $m_0$, the $P(y|x, m_0)$ could be, for example, 0.9, while for some neighbouring point $m_1$, probability $P(y|x, m_1)$ could for example, 0.88. $\endgroup$
    – Xaume
    Commented May 14 at 13:46
  • $\begingroup$ How is your problem different from logistic regression? Just use $x$ and $m$ as regressors and fit your favourite regression model. With many regression models differentiability will also not be a problem. $\endgroup$
    – g g
    Commented May 18 at 10:24
  • $\begingroup$ @gg I believe the relationship between x, m and y are far too complex for logistic regression as it assumes linearity in logit $\endgroup$
    – Xaume
    Commented May 20 at 15:41
  • $\begingroup$ Well, as I said, "fit your favourite model". I was more alluding to the basic structure of the problem, than specifics of the regression function. The loss function is evident: log loss or brier score, the only thing left is how to choose the regressors/features. Sounds straightforward to me! $\endgroup$
    – g g
    Commented May 20 at 17:22

2 Answers 2


I think a model rendering $P(y|x,m)$ can be achieved with a simple classification net. The net takes in the features $x$ and $m$, and outputs a probability that the sample belongs to $y=1$. The loss used to train the net in such cases is often binary cross-entropy.

As with neural nets in general, they are prone to overfitting, leading to irregular decision boundaries that don't generalise well to new samples. This is where regularisation comes in. Constraining the model to produce a smooth output is also a form of regularisation.

One way of penalising the model for a non-smooth probability surface is to compute the Hessian, which carries information about the curvature of the output surface with respect to the inputs. This 'smoothing loss' is added to the original loss, and together they direct the model towards a solution that is both smooth and accurate.

I will demonstrate this with a simple example. For ease of visualisation I assume that both the spatial vector $m$ and the context vector $x$ are 1D. In the classification dataset below, the samples are close together in feature space and aren't linearly separable:

enter image description here

I trained a classification net to fit a sharp decision boundary:

enter image description here

Adding a smoothness loss based on the Laplacian softens the transition:

enter image description here

Reproducible example.

Data for testing:

%matplotlib widget
import numpy as np
from matplotlib import pyplot as plt


# Data for testing
from sklearn.datasets import make_moons

n_samples = 50

features, y = make_moons(n_samples, noise=0.1, random_state=1)
x, m = features.T

#View data
ax = plt.figure(figsize=(4, 4)).add_subplot(projection='3d')
ax.scatter(x, m, y, c=y, cmap='seismic', marker='^')
ax.set(xlabel='m', ylabel='x', zlabel='y')
ax.view_init(azim=0, elev=90)

Model fitting:

import torch
from torch import nn
from torch.utils.data import DataLoader

#To tensors, and batchify
features_t = torch.tensor(features).float()
y_t = torch.tensor(y).float()

train_loader = DataLoader(list(zip(features_t, y_t)), shuffle=True, batch_size=8)

#Define model

n_features = features_t.shape[1]

hidden_size = 12
model = nn.Sequential(
    nn.Linear(n_features, hidden_size),
    nn.Linear(hidden_size, hidden_size),
    nn.Linear(hidden_size, hidden_size),
    nn.Linear(hidden_size, hidden_size),

    nn.Linear(hidden_size, hidden_size),
    nn.Linear(hidden_size, hidden_size),
    nn.Linear(hidden_size, hidden_size),
    nn.Linear(hidden_size, hidden_size),

    nn.Linear(hidden_size, 1),

model_with_prob_out = nn.Sequential(model, nn.Sigmoid())

    'Model size is',
    sum([p.numel() for p in model.parameters() if p.requires_grad]),
    'trainable parameters'

from torch.autograd.functional import hessian as torch_hessian
def calc_smoothness_loss(model_with_prob_out, sample_mb):
    abs_laplacian = torch.zeros(sample_mb.shape[0])
    for i, sample in enumerate(sample_mb):
        hess = torch_hessian(model_with_prob_out, sample.unsqueeze(0), create_graph=True)
        abs_laplacian[i] = torch.diagonal(hess.squeeze()).abs().sum()
    smoothness_loss = abs_laplacian.mean()
    return smoothness_loss

#Training loop
optimiser = torch.optim.NAdam(model.parameters())

for epoch in range(n_epochs := 200):

    for sample_mb, y_mb in train_loader:
        logits = model(sample_mb).ravel()
        label_loss = nn.BCEWithLogitsLoss()(logits, y_mb)
        smoothness_loss = calc_smoothness_loss(model_with_prob_out, sample_mb)
        weighted_loss = label_loss + 0.2 * smoothness_loss

        #Step optimiser

    if not (epoch==0 or (epoch + 1) % 25==0):
    with torch.no_grad():
        logits = model(features_t).ravel()
        smoothness_loss = calc_smoothness_loss(model_with_prob_out, features_t)
    label_loss = nn.BCEWithLogitsLoss()(logits, y_t)
    acc = ((logits > 0) == y_t.bool()).float().mean() * 100
        f'[epoch {epoch + 1:>3d}/{n_epochs:>3d}]',
        f'weighted loss: {weighted_loss:>7.3f} |',
        f'acc: {acc:5.1f}% |',
        f'label loss: {label_loss:>5.3f} |',
        f'smoothness loss: {smoothness_loss:>7.3f}',

Visualise results:

sample_res = 1000
x_points = np.random.uniform(x.min(), x.max(), sample_res)
m_points = np.random.uniform(m.min(), m.max(), sample_res)
xm_points = np.column_stack([x_points, m_points])

with torch.no_grad(): logits = model( torch.tensor(xm_points).float() )
xm_probas = nn.Sigmoid()(logits).squeeze()

# Plot
ax = plt.figure(figsize=(5, 5)).add_subplot(projection='3d', proj_type='persp', focal_length=0.12)

if True:
    #Coloured wireframe plot
    X, M = np.meshgrid(np.linspace(x.min(), x.max()), np.linspace(m.min(), m.max()))
    XM = np.column_stack([X.ravel(), M.ravel()])
    with torch.no_grad():
        XM_preds = nn.Sigmoid()(model( torch.tensor(XM).float() )).reshape(X.shape)
    im = ax.plot_surface(
        X, M, XM_preds, cmap='coolwarm',
        facecolors=plt.colormaps['coolwarm'](XM_preds), facecolor=[0]*4,
        antialiased=True, shade=False, rcount=45, ccount=45, linewidth=0.2
    #(Tri)surface plot
    im = ax.plot_trisurf(x_points, m_points, xm_probas, cmap='coolwarm', antialiased=False)

#Instersecting projection
# ax.tricontourf(x_points, m_points, xm_probas, cmap='coolwarm', zdir='z', offset=0.5, alpha=0.1)
ax.scatter(x, m, y, c=y, cmap='seismic', depthshade=False, marker='^', s=30)
ax.view_init(azim=106, elev=47)
ax.set(xlabel='x', ylabel='m', zlabel='y')
ax.set_zticks([0, 1])
    im, shrink=0.5, aspect=10, label='output probability', ticks=[0, 0.5, 1]

f, ax = plt.subplots(figsize=(7, 3), layout='tight')
im = ax.tricontourf(x_points, m_points, xm_probas, cmap='coolwarm', alpha=0.3)
ax.tricontour(x_points, m_points, xm_probas, colors='black', linestyles='--', linewidths=3, levels=1, alpha=0.3)
ax.scatter(x, m, c=y, marker='^', cmap='seismic')
ax.set(xlabel='x', ylabel='m')
# ax.set(xticks=[-1, 0, 1], yticks=[-1, 0, 1])
f.colorbar(im, label='output probability', aspect=8, pad=0.05, ticks=[0, 0.5, 1])

A more efficient implementation of the Hessian operation is discussed here.

  • $\begingroup$ +1 for proposition for smoothness loss. However, due to sparsity of data sample, I still hope there's some work out there that introduce neural network architecture/setup that promotes smooth classification probability surface due to some form of inductive bias $\endgroup$
    – Xaume
    Commented May 20 at 15:45
  • $\begingroup$ What sort of dimensionality are $m$ and $x$, and approximately how many samples are there? That would help me understand the degree of sparsity. $\endgroup$ Commented May 20 at 16:02
  • 1
    $\begingroup$ $m$ is 2-d while x is ~1000-d. The training dataset contains ~200k samples, but in the end we are challenged with learning all the weights with binary labels only. $\endgroup$
    – Xaume
    Commented May 20 at 16:38
  • $\begingroup$ That sounds like a relatively high-dimensional space, for which a linear model might work well. In high-dimensional spaces the samples will be far apart and the vastness of the space can make linear separation feasible. I'd consider trying a simple model like logistic regression, and seeing how that scores on an unseen validation set. If it overfits, you can use dimensionality reduction to try for a more generalisable solution. If a sample or synthetic dataset is available please consider including it in the question. $\endgroup$ Commented May 20 at 17:05

I could try to axiomatize it as follows. Your sample space is $Y\times M\times X$, with binary outcomes of $m$, possible values of $m$ and $x$.

What you seem to seek is:

$$ P[Y=1\,|\,x,m]=\frac{P\left[Y=1,x,m\right]}{P[x,m]}=\frac{P\left[Y=1,x,m\right]}{P\left[Y=1,x,m\right]+P\left[Y=0,x,m\right]} $$


$$ P\left[Y=1,x,m\right]=P\left[x,\,m\,|\,Y=1\right]\cdot P[Y=1] $$

You should be able to extract $P\left[Y=1\right]$ from your data - it is simply the proportion of Y=1 events over all time and space. Then as long as you can estimate the density of Y=1 events, i.e. $P\left[x,\,m\,|\,Y=1\right]$ and the same for $Y=0$, you should be able to get $P\left[Y=1\,|\,x,m\right]$.

In the end, I think it will boil down to probability density estimation.

If you can assume independence of events between times and places, you can try to model it Bernoulli random variables indexed by your position and time:

$$ Y_{x,m}\sim Bernoulli\left(\lambda\left(x,\,m\right)\right) $$

Then the problem will be translated into finding the distribution of $\lambda$. You can try to parametrize $\lambda\left(x,m\right)$, subject to $\lambda\in [0,1]$ and estimate the parameters on maximum likelihood of your data

  • 1
    $\begingroup$ Interesting proposition. However, I'm afraid the space of contextual features $x$ is far too complex to model $P(x, m|Y=1)$ $\endgroup$
    – Xaume
    Commented May 20 at 15:46
  • $\begingroup$ @Xaume It were helpful, if you made clear what your issues are. As your question is phrased now, it sounds more like you want to understand the basic structure. But your comment here refers to complexity of the features, which you have not mentioned in the question. $\endgroup$
    – g g
    Commented May 20 at 17:26

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.