Learning the Optimal Composite Mediator: Closed-Form Solution and Inference

1. Problem Statement

The paper addresses the challenge of high-dimensional mediation analysis, where an exposure (treatment) $A$ affects an outcome $Y$ through a large set of candidate mediators $X \in \mathbb{R}^{n \times p}$ (e.g., gene expression or protein abundances).

• The Goal: Instead of analyzing individual mediators, the authors seek a composite mediator $M = Xw$ (a weighted linear combination of mediators) that maximizes the indirect effect (Natural Indirect Effect, NIE) of $A$ on $Y$.
• The Objective: Find the weight vector $w^*$ that maximizes the product of the path coefficients:
  $$w^* = \arg\max_{w} \hat{\alpha}(w)\hat{\beta}(w)$$
  where $\hat{\alpha}(w)$ is the effect of $A$ on $M$, and $\hat{\beta}(w)$ is the effect of $M$ on $Y$ (controlling for $A$).
• The Challenge: The objective function is non-convex in $w$, and standard numerical optimization methods are computationally expensive and prone to local optima. Furthermore, existing methods (like Principal Component Mediation Analysis) optimize variance or likelihood rather than the specific mediation criterion, and iterative methods (like Directions of Mediation) lack closed-form solutions and global optimality guarantees.

2. Methodology: MaxIE Algorithm

The authors propose MaxIE, a method that derives a closed-form global solution to the non-convex optimization problem using a geometric argument.

A. Geometric Reduction

The problem is reformulated in a weighted inner product space defined by the residual Gram matrix $V = Z^\top Z$ (where $Z$ are mediators residualized on $A$).

1. Sufficient Statistics: The method computes three statistics: $a = X^\top A$, $z = X^\top Y - \dots$, and $V$.
2. Path Vectors: Two key vectors are defined:
   • $p = V^{-1}a$ (direction maximizing the treatment-to-mediator path).
   • $q = V^{-1}z$ (direction maximizing the mediator-to-outcome path).
3. The Bisector Solution: The objective function factorizes into a "path strength" term and an "alignment" term. The alignment term $\cos(\angle(w, p)) \cdot \cos(\angle(w, q))$ is maximized when $w$ bisects the angle between $p$ and $q$.
   • The optimal weight vector is:
     $$w^* \propto \frac{p}{\|p\|_V} + \frac{q}{\|q\|_V}$$
   • This solution is obtained via two linear solves (to find $p$ and $q$) and a simple vector addition.

B. Primal and Dual Regimes

• Primal ($n > p$): Directly solves for $w \in \mathbb{R}^p$ with complexity $O(np + p^3/3)$.
• Dual ($p \ge n$): When mediators outnumber observations, the primal matrix $V$ is singular. The authors derive a dual formulation working in $\mathbb{R}^n$ using a kernel matrix $K_Z = ZZ^\top$. This recovers the optimal weight in the row space of $X$ with complexity $O(n^2p + n^3/3)$.

C. Global Null Test (Cosine Test)

The same geometric framework yields a statistical test for the existence of any composite mediator ($H_0: \max_w |\alpha(w)\beta(w)| = 0$).

• Statistic: The test is based on the cosine of the angle $\phi$ between the two path vectors $p$ and $q$.
• Distribution: Under the null, the statistic $T = \cos\phi \cdot \sqrt{(p-1)/(1-\cos^2\phi)}$ follows a $t(p-1)$ distribution in the primal regime and $t(n-2)$ in the dual regime.
• Significance: This test is a "free by-product" of the solver, requiring no additional computation. It detects both concordant (positive indirect effect) and suppression (negative indirect effect) mediation.

3. Key Contributions

1. Closed-Form Global Optimality: Proves that the non-convex problem of maximizing the indirect effect has a unique global solution found via a geometric bisector, eliminating the need for iterative optimization.
2. Computational Efficiency: The algorithm runs at the same cost as Ordinary Least Squares (OLS), orders of magnitude faster than numerical optimization (e.g., L-BFGS-B) or iterative eigenvalue methods.
3. Dual Implementation: Provides the first computationally feasible solution for the high-dimensional regime where $p \ge n$.
4. Inference Framework: Derives an asymptotic global test for the existence of composite mediation, characterizing power as a function of the population path angle.
5. Handling Suppression: The method naturally identifies both concordant and suppression mediators by solving for the bisector of $p$ and $q$ (or $p$ and $-q$).

4. Experimental Results

• Optimization Performance:
  • In simulations, MaxIE consistently achieved higher indirect effects ($\hat{\alpha}\hat{\beta}$) than single-path OLS baselines (regressing $Y$ on $X$ or $A$ on $X$) and the iterative "Directions of Mediation" (DM) method.
  • MaxIE was ~200x to 400x faster than numerical solvers while achieving identical objective values, confirming global optimality.
• Inference Performance:
  • The cosine test controlled Type I error rates at the nominal level across various null scenarios.
  • It demonstrated significantly higher power than the Intersection-Union Test (IUT), especially in high-dimensional settings ($p \gg n$) and with weak signals.
  • Power was shown to be a U-shaped function of the angle between path vectors, peaking when paths are aligned or opposed, and dropping to zero when orthogonal.
• Real-World Application (UK Biobank):
  • Applied to $n=38,383$ individuals with $p=2,916$ proteins to study the mediation of age on dementia.
  • Result: MaxIE rejected the global null ($p = 6.4 \times 10^{-9}$).
  • Comparison: Unlike standard aging clocks (high age tracking, low disease prediction) or disease predictors (high disease prediction, low age tracking), the MaxIE composite jointly optimized both: it tracked biological aging ($r_{MA} \approx 0.90$) and predicted dementia ($AUC \approx 0.84$) better than any single-criterion baseline.

5. Significance and Implications

• Theoretical: Resolves a long-standing gap in mediation analysis by providing a tractable, optimal solution for composite mediators in high dimensions without relying on heuristics or iterative approximations.
• Practical: Enables the discovery of "biological clocks" or composite biomarkers that are not just correlated with an exposure (e.g., age) but are mechanistically linked to the outcome (e.g., dementia).
• Federated Learning: The method relies only on cross-product summaries ($X^\top A, X^\top Y, X^\top X$), making it ideal for federated data infrastructures where individual-level data cannot be shared across institutions.
• Limitations & Future Work: The current theory assumes Gaussian errors and linear structural models. Future work aims to extend this to sparse extensions (L1-penalization), nonlinear kernels, and semiparametric models.

In summary, this paper introduces a mathematically elegant and computationally efficient framework for identifying optimal composite mediators, bridging the gap between high-dimensional data analysis and causal inference.