The Advantages of Energy Regression for Mediation

# Optional: enable parallel cross-fitting
library(future)
plan(multisession)

Introduction

Energy regression (engression) learns the full conditional distribution \(P(Y|X)\), whereas standard mediation tools like crumble (via SuperLearner) typically target only the conditional mean \(E[Y|X]\). This vignette demonstrates scenarios where this difference matters for obtaining correct causal estimates, comparing dma against crumble with a GLM learner on each scenario.

Mediation decomposition

In natural effects mediation, we decompose the ATE into:

Standard methods estimate these as mean contrasts. But the full potential outcome distributions \(P(Y(a, M(a')))\) reveal much more: how treatment reshapes the entire outcome distribution, not just its center.

1. The Bimodal Mediator Challenge

In many settings, a behavior (the mediator) is bimodal: individuals are either “high” or “low,” but rarely at the average.

Scenario:

library(dma)
library(crumble)

set.seed(123)
n <- 1000
A <- rbinom(n, 1, 0.5)
M <- ifelse(A == 0, rnorm(n, 0, 0.5),
            ifelse(runif(n) > 0.5, rnorm(n, -2, 0.5), rnorm(n, 2, 0.5)))
Y <- M^2 + rnorm(n, 0, 0.5)
data <- data.frame(A, M, Y, W = rnorm(n))

d0 <- \(data, trt) 0
d1 <- \(data, trt) 1

# dma (Distributional)
res_dma <- dma(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  control = dma_control(crossfit_folds = 1L, num_epochs = 200L, riesz_epochs = 50L)
)

# crumble (Mean-based GLM)
res_crumble <- crumble(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  learners = "glm",
  control = crumble_control(crossfit_folds = 1L, epochs = 20L)
)
Estimand Truth dma crumble (GLM)
NIE 4.0 4.992 -0.240

Distributional decomposition

set.seed(42)
n_o <- 10000
eps <- rnorm(n_o, 0, 0.5)

# Draw mediators under each regime
M_ctrl <- rnorm(n_o, 0, 0.5)                    # M(0)
M_trt  <- ifelse(runif(n_o) > 0.5,              # M(1)
                 rnorm(n_o, -2, 0.5), rnorm(n_o, 2, 0.5))

# Four potential outcomes: Y = M^2 + eps (no direct A effect)
Y00 <- M_ctrl^2 + eps   # Y(0, M(0))
Y11 <- M_trt^2  + eps   # Y(1, M(1))
Y10 <- M_ctrl^2 + eps   # Y(1, M(0)) = Y(0, M(0)) since no direct effect
Y01 <- M_trt^2  + eps   # Y(0, M(1)) = Y(1, M(1)) since no direct effect

knitr::kable(print_decomposition(Y00, Y11, Y10, Y01), align = "lr")
Effect Value
ATE = E[Y(1,M(1))] - E[Y(0,M(0))] 3.999
NDE = E[Y(1,M(0))] - E[Y(0,M(0))] 0.000
NIE = E[Y(1,M(1))] - E[Y(1,M(0))] 3.999 
plot_po_distributions(Y00, Y11, Y10, Y01,
  title = "1. Bimodal Mediator: Potential Outcome Distributions",
  subtitle = "Y = M² (no direct effect): ATE = NIE, NDE = 0. Mean-based methods see NIE ≈ 0.")

Key insight: Because \(A\) has no direct effect on \(Y\), \(Y(1,M(0)) = Y(0,M(0))\) and \(Y(0,M(1)) = Y(1,M(1))\) — the NDE is zero and the entire ATE flows through the mediator. The treated and indirect-only distributions are identical: a heavy right tail from \(M^2 \approx 4\). A mean-based learner that only sees \(E[M|A=0] = E[M|A=1] = 0\) concludes NIE \(\approx 0\).

2. Heteroskedasticity (Variance as Signal)

Treatment often changes the variance or spread of a behavior without shifting its mean.

Scenario:

set.seed(123)
n <- 1000
A <- rbinom(n, 1, 0.5)
M <- rnorm(n, 0, sd = ifelse(A == 0, 1, 3))
Y <- M^2 + rnorm(n, 0, 0.5)
data <- data.frame(A, M, Y, W = rnorm(n))

# dma
res_dma_het <- dma(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  control = dma_control(crossfit_folds = 1L, num_epochs = 200L, riesz_epochs = 50L)
)

# crumble
res_crumble_het <- crumble(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  learners = "glm",
  control = crumble_control(crossfit_folds = 1L, epochs = 20L)
)
Estimand Truth dma crumble (GLM)
NIE 8.0 7.721 -0.137

Distributional decomposition

set.seed(42)
n_o <- 10000
eps <- rnorm(n_o, 0, 0.5)

M_ctrl <- rnorm(n_o, 0, 1)   # M(0): sd = 1
M_trt  <- rnorm(n_o, 0, 3)   # M(1): sd = 3

# Y = M² + eps (no direct effect)
Y00 <- M_ctrl^2 + eps
Y11 <- M_trt^2  + eps
Y10 <- M_ctrl^2 + eps  # = Y00
Y01 <- M_trt^2  + eps  # = Y11

knitr::kable(print_decomposition(Y00, Y11, Y10, Y01), align = "lr")
Effect Value
ATE = E[Y(1,M(1))] - E[Y(0,M(0))] 8.238
NDE = E[Y(1,M(0))] - E[Y(0,M(0))] 0.000
NIE = E[Y(1,M(1))] - E[Y(1,M(0))] 8.238 
plot_po_distributions(Y00, Y11, Y10, Y01,
  title = "2. Heteroskedasticity: Potential Outcome Distributions",
  subtitle = "Y = M²: tripled sd(M) creates extreme right tail. ATE = NIE = 8, NDE = 0",
  xlim = c(-2, 40))

Key insight: Treatment triples \(\text{sd}(M)\) without changing \(E[M]\). Since \(E[M^2] = \text{Var}(M) + E[M]^2\), the entire effect operates through the variance channel. The treated outcome distribution has a massively heavier right tail. Mean-based methods cannot detect this.

3. Robustness to Heavy-Tailed Noise (Outliers)

Energy regression is more robust to extreme outliers than standard mean squared error.

Scenario:

set.seed(123)
n <- 500
W <- rcauchy(n)
W <- pmax(pmin(W, 20), -20)
A <- rbinom(n, 1, plogis(0.1 * W))
M <- A + 0.5 * W + rnorm(n)
Y <- A + M + 0.5 * W + rnorm(n)
data <- data.frame(A, M, Y, W)

# dma
res_dma_heavy <- dma(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  control = dma_control(crossfit_folds = 1L, num_epochs = 200L, riesz_epochs = 50L)
)

# crumble
res_crumble_heavy <- crumble(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = d0, d1 = d1, effect = "N",
  learners = "glm",
  control = crumble_control(crossfit_folds = 1L, epochs = 20L)
)
Estimand Truth dma crumble (GLM)
NIE 1.0 1.067 1.049

Distributional decomposition

set.seed(42)
n_o <- 10000
W_o <- rcauchy(n_o)
W_o <- pmax(pmin(W_o, 20), -20)
eps_m <- rnorm(n_o)
eps_y <- rnorm(n_o)

# M(0) and M(1)
M_ctrl <- 0 + 0.5 * W_o + eps_m       # A=0
M_trt  <- 1 + 0.5 * W_o + eps_m       # A=1

# Y(a, M(a')) = a + M(a') + 0.5W + eps_y
Y00 <- 0 + M_ctrl + 0.5 * W_o + eps_y  # Y(0, M(0))
Y11 <- 1 + M_trt  + 0.5 * W_o + eps_y  # Y(1, M(1))
Y10 <- 1 + M_ctrl + 0.5 * W_o + eps_y  # Y(1, M(0)) — direct effect
Y01 <- 0 + M_trt  + 0.5 * W_o + eps_y  # Y(0, M(1)) — indirect effect

knitr::kable(print_decomposition(Y00, Y11, Y10, Y01), align = "lr")
Effect Value
ATE = E[Y(1,M(1))] - E[Y(0,M(0))] 2
NDE = E[Y(1,M(0))] - E[Y(0,M(0))] 1
NIE = E[Y(1,M(1))] - E[Y(1,M(0))] 1 
plot_po_distributions(Y00, Y11, Y10, Y01,
  title = "3. Heavy Tails: Potential Outcome Distributions",
  subtitle = "Linear DGP with Cauchy confounders: NDE = NIE = 1, all distributions heavy-tailed",
  xlim = c(-25, 30))

Key insight: Here all four distributions are visible as distinct curves. The NDE shifts from \(Y(0,M(0))\) to \(Y(1,M(0))\) by 1 (direct \(A\) effect), and the NIE shifts from \(Y(1,M(0))\) to \(Y(1,M(1))\) by 1 (mediator pathway). The heavy Cauchy tails don’t change the decomposition — both methods handle this linear case well.

4. Extreme Treatment Shifts (Extrapolation)

Engression maintains the structural relationship of the data better than many ML learners when shifting far outside the training range.

Scenario: Shift \(A\) by +5 units (standard deviation is 1).

set.seed(123)
n <- 1000
A <- rnorm(n)
M <- A + rnorm(n, 0, 0.5)
Y <- M + rnorm(n, 0, 0.5)
data <- data.frame(A, M, Y, W = rnorm(n))

d_extrap <- \(data, trt) data[[trt]] + 5

# dma
res_dma_extrap <- dma(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = \(data, trt) data[[trt]],
  d1 = d_extrap, effect = "N",
  control = dma_control(crossfit_folds = 1L, num_epochs = 200L, riesz_epochs = 50L)
)

# crumble
res_crumble_extrap <- crumble(
  data = data, trt = "A", outcome = "Y", mediators = "M", covar = "W",
  d0 = \(data, trt) data[[trt]],
  d1 = d_extrap, effect = "N",
  learners = "glm",
  control = crumble_control(crossfit_folds = 1L, epochs = 20L)
)
Estimand Truth dma crumble (GLM)
ATE 5.0 4.582 9.814

Distributional decomposition

set.seed(42)
n_o <- 10000
A_o <- rnorm(n_o)
eps_m <- rnorm(n_o, 0, 0.5)
eps_y <- rnorm(n_o, 0, 0.5)

# M under each regime
M_d0 <- A_o + eps_m           # d0: M(A)
M_d1 <- (A_o + 5) + eps_m     # d1: M(A+5)

# Y = M + eps_y (no direct A effect on Y)
Y00 <- M_d0 + eps_y   # Y(d0, M(d0))
Y11 <- M_d1 + eps_y   # Y(d1, M(d1))
Y10 <- M_d0 + eps_y   # Y(d1, M(d0)) = Y(d0, M(d0)) since no direct effect
Y01 <- M_d1 + eps_y   # Y(d0, M(d1)) = Y(d1, M(d1)) since no direct effect

knitr::kable(print_decomposition(Y00, Y11, Y10, Y01), align = "lr")
Effect Value
ATE = E[Y(1,M(1))] - E[Y(0,M(0))] 5
NDE = E[Y(1,M(0))] - E[Y(0,M(0))] 0
NIE = E[Y(1,M(1))] - E[Y(1,M(0))] 5 
plot_po_distributions(Y00, Y11, Y10, Y01,
  title = "4. Extrapolation: Potential Outcome Distributions",
  subtitle = "Y = M + eps, M = A + eps: +5 shift is entirely indirect (NIE = 5, NDE = 0)")

Key insight: The +5 shift moves the entire distribution by 5 units through the mediator. Since there’s no direct \(A\) effect on \(Y\), the NDE is zero and \(Y(1,M(0)) = Y(0,M(0))\). The challenge is extrapolation: the shifted regime is 5 SDs outside training range. Engression preserves the linear structural relationship; crumble’s GLM nearly doubles the effect to 9.8.

Summary

Scenario NDE (Truth) NIE (Truth) ATE (Truth) dma crumble Key distributional feature
Bimodal mediator 0.0 4.0 4.0 4.99 -0.24 Bimodal \(Y\) under treatment
Heteroskedasticity 0.0 8.0 8.0 7.72 -0.14 Heavy right tail from variance channel
Heavy-tailed 1.0 1.0 2.0 1.07 1.05 Cauchy tails (both methods OK)
Extrapolation 0.0 5.0 5.0 4.58 9.81 Distribution shifted 5 SDs out of range

The potential outcome distributions above show what mean-based estimates compress into a single number. In scenarios 1–2, the treatment operates entirely through the mediator (\(\text{NDE} = 0\)), but the mechanism is distributional — changing variance or modality rather than the mean. These are precisely the cases where dma’s distributional regression recovers effects that conditional-mean methods miss.

References