Spiked-CFR: Causal Representation Learning from LLMs via Wasserstein Projection Pursuit

Estimating treatment effects from observational text is increasingly practical with Large Language Models (LLMs). However, applying causal representation learning directly to high-dimensional LLM embeddings faces a fundamental barrier: empirical Wasserstein matching suffers from the curse of dimensionality, rendering standard generalization guarantees effectively vacuous. We propose SPIKED-CFR, a framework bridging this gap by assuming a Spiked Structure, where treatment selection bias is assumed to manifest primarily as a low-dimensional treated--control discrepancy in the semantic representation. We develop Wasserstein Projection Pursuit, a minimax objective that adversarially learns an orthogonal projection on the Stiefel manifold to identify and balance only this subspace while preserving prognostic information. Under a spiked structure, we show the projected discrepancy can be estimated at a rate governed by the intrinsic dimension $k \ll D$, and we derive a tighter PEHE generalization bound that depends on $k$ rather than the ambient embedding dimension. Experiments on four semi-synthetic benchmarks and four real-world clinical benchmarks demonstrate improved accuracy and robustness over strong baselines.