Should we keep storing draws from non-free parameters in the trace by default?

[tomicapretto]

See the following simple model and the output

import bambi as bmb

my_data = bmb.load_data("my_data")

formula = bmb.Formula("y ~ x")
model = bmb.Model(formula, my_data)
idata = model.fit()
idata.posterior.data_vars

Data variables:
    Intercept  (chain, draw) float64 16kB 0.3069 0.04518 ... 0.2771 0.06184
    x          (chain, draw) float64 16kB 0.1171 0.05828 ... 0.1202 -0.09098
    y_sigma    (chain, draw) float64 16kB 1.467 0.7531 0.8906 ... 0.972 1.156

The model has 3 free parameters. Intercept and x determine the linear predictor for the mean mu, and y_sigma is the standard deviation of the error term, a scalar.

Now look at the following

formula = bmb.Formula("y ~ x", "sigma ~ x")
model = bmb.Model(formula, my_data)
idata = model.fit()
idata.posterior.data_vars

Data variables:
    Intercept        (chain, draw) float64 16kB 0.07661 0.1871 ... 0.3599 0.1323
    x                (chain, draw) float64 16kB 0.0148 0.121 ... -0.1719
    sigma_Intercept  (chain, draw) float64 16kB 0.07594 -0.06736 ... 0.1496
    sigma_x          (chain, draw) float64 16kB -0.02853 -0.04591 ... -0.1696
    sigma            (chain, draw, y_obs) float64 800kB 1.065 1.11 ... 1.278

The model has 4 free parameters: Intercept and x for the linear predictor of the mean, and sigma_Intercept and sigma_x for the linear predictor of the standard deviation of the error term. However, here we are also storing sigma, which is not a free parameter as its values are determined by the values from sigma_Intercept, sigma_x, and the predictor x.

When I implemented distributional models, I decided to keep sigma in the trace by default. It seemed the natural choice, as sigma was already being included in the non-distributional approach.

However, now I realize (thanks to @AlexanderFengler) this approach is both inconsistent and inefficient. What's more, we're not giving the user the option to disable it for computational issues.

• Inconsistent: In this second model both mu and sigma are of the same kind. They're both the result of some linear predictor that depends on other free parameters. Why do we keep non-parent parameters and not parent parameters?. I don't find a satisfying answer to this question.
• Inefficient: The number of values to store increase with the number of observations. Also, this is implemented as a Deterministic in PyMC, which is also not the best.

The only positive aspect I see of creating Deterministics (and I would consider it for both mu and sigma) is the nicer graph viz we would get.

If we don't keep these draws by default, we would still be able to compute them when it's needed.

[ahartikainen]

I would argue we keep them by default but enable users to disable them (manually selecting the parameters to keep? or not to keep)

[AlexanderFengler]

Maybe even allow saving of trace of the parent as well to complete the option set?
I am most often interested in turning all of them off, but can see how the exact opposite may apply for other users who use distributional models.

Thanks @tomicapretto for raising this.

[tjburch]

I've been doing some distributional models in brms lately, they don't return the additional parameter.

r$> set.seed(123)
    n <- 100
    x <- rnorm(n, mean = 0, sd = 1)
    y <- 2 + 1.5 * x + rnorm(n, mean = 0, sd = abs(0.5 * x))
    my_data <- data.frame(x, y)

r$> model_fit <- brm(bf(y~x, sigma~x), 
                     data = my_data, 
                     family = gaussian(),
                     chains = 2, 
                     iter = 2000, 
                     seed = 123)
    print(model_fit)
    
    ...Removed fitting output here...
    
 Family: gaussian 
  Links: mu = identity; sigma = log 
Formula: y ~ x
         sigma ~ x
   Data: my_data (Number of observations: 100)
  Draws: 2 chains, each with iter = 2000; warmup = 1000; thin = 1;
         total post-warmup draws = 2000

Population-Level Effects:
                Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
Intercept           1.93      0.05     1.84     2.02 1.00     1968     1510
sigma_Intercept    -0.84      0.07    -0.98    -0.69 1.00     1950     1417
x                   1.47      0.05     1.38     1.57 1.00     1736     1519
sigma_x             0.04      0.06    -0.07     0.15 1.00     1879     1646

We're certainly not bound to what they do by any means but I just wanted to put that here as a data point from another library that has tackled this issue.

Now onto my opinion as a personal user, when I write pymc models, most times I don't put observation-level deterministic in my models, so I will haphazardly chuck az.plot_trace(fit) as soon as it's done sampling without selecting just the parent-level parameters. This goes bananas for the observation-level parameters, and I've crashed more than my fair share of sessions in this way. So because of the inconsistency point you mention, I'm a "no," however, it might be nice to have on as an option just for the sake of the graphviz plot.

[AlexanderFengler]

On top of the az.plot_trace() issue (can related there as well), there is also the aspect that this can severely blow up memory during sampling / storing the traces, for large datasets?

[tomicapretto]

I will chat with some PyMC folks to see how feasible it is to adapt pm.sample to not include deterministics. That way we would get a good graph plus avoid extra memory requirement

[tomicapretto]

@ricardoV94 shared the other day a trick that we could exploit (don't know if we would be causing undesired problems).

edit with this, we would have the best of all worlds

• We can show the deterministics in the graph
• Avoid storing things in the trace
• Allow users to specify if they want to store things in the trace

Look at this code

import pymc as pm

# Simulate data
x = pm.draw(pm.Normal.dist(), 100)
y = 3 + 2 * pm.draw(pm.Normal.dist(sigma=2), 100)

# Create model, including the mean as a deterministic
with pm.Model() as model:
    alpha = pm.Normal("alpha")
    beta = pm.Normal("beta")
    sigma = pm.HalfNormal("sigma")
    mu = pm.Deterministic("mu", alpha + beta * x)
    _ = pm.Normal("y", mu=mu, sigma=sigma, observed=y)

# Get the graph, showing "mu"
model.to_graphviz()

# The deterministics are stored in a list in the attribute `deterministics`.
# Store them for later
deterministics = model.deterministics

# Prints [mu]
print(deterministics)

# Set the model deterministics to an empty list
model.deterministics = []

# Sample
with model:
    idata = pm.sample()

# Check the deterministics are not in the trace
print(idata.posterior.data_vars)

# Bring deterministics back to the model, so the graph includes them again
model.deterministics = deterministics
model.to_graphviz()

[mu]

Data variables:
    alpha    (chain, draw) float64 16kB 2.677 2.239 2.397 ... 2.785 1.615 3.087
    beta     (chain, draw) float64 16kB 0.7562 -0.06486 0.4914 ... 0.3969 0.8051
    sigma    (chain, draw) float64 16kB 3.321 4.025 3.572 ... 3.897 3.789 3.623

[tomicapretto]

PyMC 5.12 has been released and I think we can have the best of both worlds https://github.com/pymc-devs/pymc/releases/tag/v5.12.0

[tomicapretto]

Haha, I think it's a coincidence. Well, we've talked about it in the past with Ricardo, but I think I was not the only one asking for such a feature.

[ahartikainen]

So, if we don't store things on trace we can still easily calculate the resulting value?

Does this work when data is saved and then later loaded?

[ricardoV94]

Yes Deterministics can always be computed later if you can recreate the original model. We'll add a specific helper for that in the next few days

[ahartikainen]

Then I think we can drop them from bambi too?

How do we save models or information of model in bambi?

[tomicapretto]

Then I think we can drop them from bambi too?

Yup, this is what motivated the discussion. We can always compute the deterministics after inference in Bambi, even without the helper from PyMC, as we already had the internal tooling to do so.

How do we save models or information of model in bambi?

Do you refer to the model, the trace, or both?

For the model, I'm not sure if pickle or similar approaches work. But if you save the trace as an InferenceData, you can just recreate the model with the same code and use the InferenceData loaded from disk.