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.

Many important questions in healthcare, education, and public policy ask whether an intervention truly changes an outcome. In practice, the information needed to answer these questions is often written in free-form text, such as clinical notes, online posts, or personal reports, rather than stored in clean tables. Large language models can turn such text into useful numerical representations, but these representations are very high-dimensional, which makes it difficult to reliably compare people who received an intervention with those who did not. This paper proposes Spiked-CFR, a method for estimating treatment effects from text. Instead of forcing the two groups to match in every direction of a large representation space, our method searches for a smaller part of the representation where the most relevant group differences appear, and focuses balancing there while keeping the full text representation for prediction. We evaluate the method on text-based benchmarks derived from semi-synthetic datasets and real-world clinical treatment comparisons. The results show more accurate and stable treatment-effect estimates than strong text-based baselines. This work may help researchers use unstructured text more reliably for causal analysis, while still requiring care because observational text can contain missing information, hidden biases, and extraction errors.