Distributional Mediation Analysis

Data Generating Process

• \(A \in \{0, 1\}\) (Binary Treatment)
• \(M | A \sim N(A, 0.5^2)\) (Mediator shifts its mean with treatment)
• \(Y = M^2 + W + \epsilon\) (Outcome is a squared function of the mediator)
• \(W \sim N(0, 1)\) (Confounder)

set.seed(42)
n <- 2000
W <- rnorm(n)
A <- rbinom(n, 1, 0.5)
M <- rnorm(n, A, 0.5)
Y <- M^2 + W + rnorm(n, 0, 0.5)
data <- data.frame(A, M, Y, W)

ggplot(data, aes(x = M, fill = factor(A))) +
  geom_density(alpha = 0.5) +
  labs(title = "Mediator Distribution by Treatment", fill = "Treatment (A)") +
  theme_minimal()

Oracle Counterfactual Outcome Distributions

Since we know the DGP, we can compute the true potential outcome distributions under each treatment regime. This is what plot_counterfactual_density() aims to recover from data.

Under \(do(A=a)\): \(M \sim N(a, 0.25)\) and \(Y = M^2 + W + \epsilon\), so the outcome distribution changes not just in location but in shape — the squared mediator introduces right-skew when \(A=1\).

set.seed(42)
n_oracle <- 10000
W_o <- rnorm(n_oracle)
eps_o <- rnorm(n_oracle, 0, 0.5)

# Y(0): do(A=0), M ~ N(0, 0.25)
M0 <- rnorm(n_oracle, 0, 0.5)
Y0 <- M0^2 + W_o + eps_o

# Y(1): do(A=1), M ~ N(1, 0.25)
M1 <- rnorm(n_oracle, 1, 0.5)
Y1 <- M1^2 + W_o + eps_o

# Y(1, M(0)): direct effect — treat but keep mediator at control level
Y10 <- M0^2 + W_o + eps_o  # same as Y0 (no direct effect of A on Y here)

# Y(0, M(1)): indirect effect — don't treat but shift mediator as if treated
Y01 <- M1^2 + W_o + eps_o  # same as Y1

oracle_df <- rbind(
  data.frame(Y = Y0, Regime = "Y(0, M(0)): Control"),
  data.frame(Y = Y1, Regime = "Y(1, M(1)): Treated"),
  data.frame(Y = Y10, Regime = "Y(1, M(0)): Direct only")
)

ggplot(oracle_df, aes(x = Y, fill = Regime)) +
  geom_density(alpha = 0.4) +
  labs(title = "Oracle Counterfactual Outcome Distributions",
       subtitle = "Y = M² + W + eps: treatment shifts distribution shape, not just location",
       x = "Outcome (Y)", y = "Density") +
  theme_minimal()

Key observations:

• Y(0, M(0)) (control): concentrated near zero, since \(M \approx 0\) and \(M^2 \approx 0\)
• Y(1, M(1)) (treated): shifted right with heavier right tail, since \(M \approx 1\) and \(M^2 \approx 1\)
• Y(1, M(0)) (direct only) overlaps exactly with control — because \(A\) does not enter the outcome equation (\(Y = M^2 + W + \epsilon\)), treating while holding \(M\) at its control level changes nothing. The NDE is zero.
• Since NDE = 0, the entire treatment effect is indirect: NIE \(= E[Y(1,M(1))] - E[Y(1,M(0))] \approx\) 1, equal to the ATE. But this single number hides the distributional shift from a symmetric, concentrated distribution to a right-skewed one.

Below we fit dma and use plot_counterfactual_density() to recover this from data, without knowing the true DGP.

Coefficient Plot

The plot() method shows point estimates and confidence intervals for all estimated effects:

plot(res)

Marginal Counterfactual Densities

The default plot_counterfactual_density() shifts the treatment variable and predicts Y from the trained engression models. This shows the marginal outcome distributions \(P(Y \mid do(A=0))\) and \(P(Y \mid do(A=1))\):

plot_counterfactual_density(res, n_samples = 5)

Mediation Counterfactual Densities

For mediation analysis, we care about the cross-world counterfactuals that define the NDE and NIE:

• \(Y(0, M(0))\): outcome under control
• \(Y(1, M(0))\): outcome under treatment but mediator at control level (defines NDE)
• \(Y(1, M(1))\): outcome under full treatment (NIE = difference from \(Y(1,M(0))\))

Passing use_weights = TRUE uses the learned Riesz representers (density ratios from the mediation cascade) to reweight the observed \(Y\) to each counterfactual regime. This is importance sampling: if \(\alpha_k\) is the density ratio for regime \(k\), then \(E[\alpha_k \cdot f(Y)] = E[f(Y(k))]\).

plot_counterfactual_density(res, use_weights = TRUE)

Estimated vs Oracle Comparison

We can overlay the weighted estimates with the oracle distributions for all three mediation regimes:

# Weighted estimates
alphas <- res$alpha_n
W <- alphas[["alpha3"]]

df_est <- do.call(rbind, lapply(colnames(W), function(nm) {
  w <- pmax(W[, nm], 0)
  w <- w / sum(w) * length(w)
  data.frame(Y = data$Y, Weight = w, Regime = nm)
}))

regime_map <- c("000" = "Y(0, M(0))", "100" = "Y(1, M(0))", "111" = "Y(1, M(1))")
df_est$Regime <- regime_map[df_est$Regime]

df_oracle <- rbind(
  data.frame(Y = Y0, Regime = "Y(0, M(0))"),
  data.frame(Y = Y10, Regime = "Y(1, M(0))"),
  data.frame(Y = Y1, Regime = "Y(1, M(1))")
)

ggplot() +
  geom_density(data = df_oracle, aes(x = Y, color = Regime, linetype = "Oracle"), linewidth = 0.9) +
  geom_density(data = df_est, aes(x = Y, color = Regime, weight = Weight, linetype = "Estimated"), linewidth = 0.9) +
  scale_color_manual(values = c(
    "Y(0, M(0))" = "#377eb8", "Y(1, M(0))" = "#984ea3", "Y(1, M(1))" = "#e41a1c"
  )) +
  scale_linetype_manual(values = c("Oracle" = "solid", "Estimated" = "dashed")) +
  labs(title = "Estimated vs Oracle Counterfactual Distributions",
       subtitle = "Solid = oracle (n=10,000), Dashed = weighted via Riesz representers (n=2,000)",
       x = "Outcome (Y)", y = "Density", linetype = "Source") +
  theme_minimal() +
  theme(legend.position = "bottom")

Analysis

In this DGP with no direct effect (\(Y = M^2 + W + \epsilon\)):

• Y(0, M(0)) and Y(1, M(0)) overlap (NDE = 0) — treatment without shifting the mediator changes nothing.
• Y(1, M(1)) shifts right with a heavier tail — the entire effect operates through \(M\), changing the shape of the outcome distribution, not just its mean.
• The weighted density (dashed) recovers all three counterfactual distributions from the observed data using the Riesz representer cascade, including the cross-world \(Y(1, M(0))\) that cannot be obtained by simple treatment shifting.

Traditional methods report the average shift (NIE), but dma lets you see how the entire probability mass of the outcome moves — invaluable for understanding treatment mechanisms in non-linear systems.