📊 statistics

Partial Identification Learning with Categorical Treatments for Individualized Treatment Rules

This paper proposes a novel partial identification learning framework that extends individualized treatment rule optimization to categorical treatments, outcomes, and instrumental variables by introducing a generalized minimax loss criterion and a symmetric embedding strategy to derive a differentiable surrogate risk function, thereby enabling robust decision-making under unmeasured confounding without relying on strong causal assumptions.

Original authors: Johannes Hruza, Paweł Morzywołek, Jakob Zeitler, Samir Bhatt, Michael C Sachs

Published 2026-08-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Johannes Hruza, Paweł Morzywołek, Jakob Zeitler, Samir Bhatt, Michael C Sachs

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of medicine and economics, making the right choice for an individual is often a matter of life and death, or at least of significant well-being. Doctors must decide which drug to prescribe, and policymakers must choose which intervention to fund, but they rarely have perfect information. The ideal scenario involves knowing exactly how every single person would respond to every possible option if they could try them all at once. In reality, researchers can only observe what happened to people who actually received a specific treatment. To figure out what would have happened otherwise, they must rely on assumptions. The most common assumption is that they have measured every single factor that influences both the choice of treatment and the final outcome. If they have missed even one hidden factor, their conclusions can be dangerously wrong, leading to policies that help some but harm others.

To navigate this uncertainty, scientists sometimes use a special tool called an instrumental variable. Think of this as a natural nudge that influences which treatment a person receives but has no direct power to change their health outcome. A classic example is a genetic variant that makes a person more likely to take a certain medication, or a random assignment in a trial that isn't perfectly followed. Because this nudge is independent of the hidden factors that usually cause confusion, it can help researchers draw a boundary around the truth. Instead of pinpointing a single, exact answer, this method defines a safe zone—a range of possibilities where the true effect must lie. This approach, known as partial identification, accepts that we cannot know everything, but it allows us to make robust decisions that are safe even when we are unsure.

A team of researchers has now taken this concept of partial identification and expanded it to handle a much more complex reality: situations where there are not just two choices, but many. In the real world, doctors often choose between several different drugs, or different dosage levels, rather than simply picking between a treatment and a placebo. Previous methods for handling these multiple options under uncertainty were limited to simple yes-or-no scenarios. The new framework developed by Johannes Hruza and his colleagues bridges this gap, creating a system that can learn the best decision rule when faced with multiple categorical treatments, outcomes, and instrumental variables. Their work provides a way to find the safest, most effective policy even when hidden factors are at play and the number of options is large.

The core challenge the researchers tackled is how to choose the best option when you cannot be certain which one is truly the best. Imagine a doctor looking at a patient's data and seeing that Treatment A might be very good, but could also be mediocre, while Treatment B might be terrible, but could also be excellent. Without knowing the hidden factors, the doctor cannot be sure. The researchers adopted a cautious strategy called the "minimax" approach. This means they look for the decision that minimizes the worst possible mistake. Instead of hoping for the best outcome, they ask: "If I pick this treatment, what is the maximum amount of good I could be missing out on compared to the best alternative?" They then choose the path that keeps this potential loss as small as possible. This ensures that even in the worst-case scenario, the decision remains as close to optimal as the available information allows.

To make this work for many different treatments at once, the team had to solve a difficult geometric problem. Standard methods for comparing multiple options often treat each choice as a separate battle against all the others, which can lead to confusing and contradictory results. The researchers instead mapped all the treatment options onto the corners of a balanced, multi-sided shape in a mathematical space. This shape, known as a simplex, treats every option with perfect symmetry, ensuring that no single choice is favored by the geometry of the math itself. They then trained a computer model to learn how to point toward the best corner of this shape based on a patient's characteristics. By turning the problem into a smooth, continuous search for the right direction, they allowed powerful optimization tools to find the best policy, even though the underlying logic was based on worst-case scenarios.

The researchers tested their new method through a series of rigorous computer simulations. They created virtual worlds where they knew the true answers, allowing them to see how well the new rules performed compared to existing methods. In these tests, they varied the strength of the hidden confounding factors, the power of the instrumental variable, and the number of treatment options available. The results showed that when hidden factors were strong and misleading, traditional methods that assumed no hidden factors failed dramatically, often recommending the wrong treatment. In contrast, the new partial identification method remained stable and reliable. It did not promise to find the absolute perfect answer, but it consistently found a policy that was far safer and closer to the truth than the alternatives when uncertainty was high.

However, the simulations also revealed a clear limit to what this method can achieve. The quality of the decision depends heavily on how much information the instrumental variable provides. When the researchers used a strong instrument that gave clear signals, the new method performed exceptionally well, narrowing the gap between the safe choice and the perfect choice. But when the instrument was weak, or when the number of treatment options exceeded the number of levels the instrument could distinguish, the safety zone became too wide to be useful. In these cases, the method became overly conservative, unable to distinguish between the options because the data simply did not contain enough leverage to separate them. This finding highlights a fundamental truth: no amount of clever math can overcome a lack of information.

The work represents a significant step forward in the field of personalized decision-making. By extending partial identification to multiple choices, the researchers have provided a tool that is better suited for the messy, multi-option reality of clinical and economic life. They have shown that it is possible to learn robust policies without making unrealistic assumptions about hidden factors. While the method is not a magic bullet that solves every problem, and while it relies on the validity of the underlying causal model, it offers a principled way to navigate uncertainty. For doctors and policymakers facing complex choices with incomplete data, this framework offers a way to proceed with confidence, knowing that their decisions are protected against the worst-case scenarios that hidden confounders might otherwise create.

Technical Summary: Partial Identification Learning with Categorical Treatments for Individualized Treatment Rules

1. Problem Statement

The paper addresses the challenge of learning Individualized Treatment Rules (ITRs) in settings where unmeasured confounding is present, and the treatment, outcome, and instrumental variables (IVs) are categorical (multiclass).

Standard ITR optimization relies on the assumption of "no unmeasured confounding" to point-identify the conditional expected potential outcomes, $\mu_a(x) = E[Y(a) | X=x]$ . When this assumption is violated (e.g., due to noncompliance in trials or unobserved factors in observational studies), standard estimators become biased. While Instrumental Variables (IVs) can mitigate this, traditional IV approaches often require strong structural assumptions (e.g., monotonicity) to achieve point identification, which may be untestable or implausible.

Existing partial identification methods for ITRs, such as those by Pu and Zhang [24], are largely restricted to binary treatment settings. This paper identifies a gap in extending these robust, partial identification frameworks to multicategorical treatments (e.g., choosing among multiple drugs or dosage levels), where the decision space is inherently more complex and standard binary extensions (like One-vs-Rest) introduce geometric inconsistencies.

2. Methodology

The proposed framework, Categorical Partial Identification Learning, consists of three main components:

A. Partial Identification Bounds

Instead of attempting to point-identify $\mu_a(x)$ , the framework characterizes a feasible region for the conditional mean using lower and upper bounds, $L_a(x)$ and $U_a(x)$ , derived from the observed data and a causal model (e.g., a Directed Acyclic Graph). These bounds satisfy:
$L_a(x) \leq \mu_a(x) \leq U_a(x)$
The paper utilizes the causaloptim package (Sachs et al. [26]) to derive sharp symbolic bounds for categorical treatments and IVs under general causal structures, translating the Nonparametric Structural Equation Model (NPSEM) constraints into a linear optimization problem.

B. Minimax Decision Criterion

To select a treatment under uncertainty, the authors adopt an Opportunist (Minimax) criterion. Rather than maximizing the best-case scenario (Optimist) or the worst-case lower bound (Pessimist), the goal is to minimize the maximum possible loss relative to the unknown optimal treatment.

For a chosen action $a$ , the worst-case loss is defined as:
$L(x, a) := \max_{b \in A \setminus \{a\}} U_b(x) - L_a(x)$
This quantifies the maximum regret if action $a$ is chosen while the true optimal action is $b$ . The objective is to find a policy $g$ that minimizes the expected worst-case risk $R(g) = E[L(X, g(X))]$ .

C. Symmetric Multicategory Embedding and Surrogate Loss

Directly minimizing the empirical worst-case risk is computationally intractable due to the discontinuity of the $\max$ operator and the discrete nature of the action space. To overcome this, the paper introduces:

Simplex Embedding: Categorical treatments are mapped to the vertices of a regular simplex in $\mathbb{R}^{k-1}$ . This symmetric representation avoids the geometric redundancies and ambiguities of One-vs-Rest approaches. A policy is induced by a function $f: \mathcal{X} \to \mathbb{R}^{k-1}$ that maps covariates to this latent space, selecting the treatment $a$ that maximizes the inner product $\langle w_a, f(x) \rangle$ , where $w_a$ is the vertex vector for treatment $a$ .
Weighted Surrogate Risk: A differentiable surrogate loss is constructed using a softmax weighting scheme. The surrogate risk is defined as:
$\hat{R}_s(f) = \frac{1}{N} \sum_{i=1}^N \sum_{a=1}^k p_a(X_i, f) \hat{L}(X_i, a) + \lambda \Omega(f)$
where $p_a(X, f)$ is a softmax weight based on the alignment between $f(X)$ and the treatment vector $w_a$ . This allows the use of gradient-based optimization (e.g., L-BFGS or Stochastic Gradient Descent).
Implementation: The framework supports two implementations:
- Kernelized: Using Reproducing Kernel Hilbert Spaces (RKHS) with the Representer Theorem for theoretical guarantees.
- Neural Network: Using Multi-Layer Perceptrons (MLPs) for scalability to large datasets.

3. Key Contributions

Extension to Categorical Treatments: The paper generalizes partial identification learning from binary to finite multicategorical treatment settings, a significant gap in existing literature.
Geometrically Consistent Embedding: It introduces a simplex-based embedding strategy that treats all treatment options symmetrically, resolving the geometric inconsistencies found in standard multiclass classification extensions for ITRs.
Differentiable Optimization: The derivation of a smooth, weighted surrogate risk function enables the use of efficient gradient-based optimization methods for ITR learning under partial identification, moving beyond the limitations of indicator-based losses.
Theoretical Guarantees: The authors provide finite-sample convergence rates via an oracle inequality. They prove that minimizing the surrogate risk converges to the minimax optimal policy and establish conditions under which the estimation error of the bounds and the learning error are controlled.
Empirical Validation: Through extensive simulations, the paper demonstrates that the proposed methods (Kernelized and Neural) yield ITRs significantly closer to the oracle ITR than existing alternatives (outcome regression, MOML) in the presence of unmeasured confounding.

4. Experimental Results

The authors conducted four simulation studies varying sample size, confounding strength, IV strength, and the number of treatments ( $k$ ):

Robustness to Confounding: In the presence of unmeasured confounding, standard outcome regression methods (Tree, Neural) and point-identified methods (MOML) degrade rapidly as confounding strength increases. In contrast, the proposed partial identification methods remain stable and superior, staying close to the theoretical minimax bound.
Instrument Strength: Performance of the proposed methods is positively correlated with IV strength. Stronger instruments narrow the identification bounds, reducing the gap between the minimax policy and the oracle policy.
Scalability and Structural Limits: The methods perform well when the number of treatments $k$ is less than or equal to the number of instrument levels $|Z|$ . However, when $k > |Z|$ , performance degrades sharply. The authors attribute this to a structural limitation: an instrument with $|Z|$ levels cannot induce enough distinct conditional treatment distributions to tightly bound $k > |Z|$ potential outcomes simultaneously.
Convergence: In the absence of confounding, the proposed methods converge to the minimax rule (which is suboptimal compared to the oracle in this specific case), while outcome regression methods converge to the oracle. This confirms the framework correctly targets the minimax objective.

5. Significance and Claims

The paper claims to provide a robust framework for decision-making under uncertainty when strong causal assumptions cannot be justified. By leveraging partial identification bounds and a symmetric embedding strategy, it enables the learning of optimal treatment rules for multicategorical choices in the face of unmeasured confounding.

The authors emphasize that the framework's effectiveness is contingent on the informativeness of the bounds. Stronger instruments lead to narrower bounds and better performance, while weak instruments or scenarios where the number of treatments exceeds the instrument's cardinality lead to wider bounds and more conservative (potentially less accurate) decisions.

The work is presented as a foundational extension of partial identification learning, moving beyond binary settings to address real-world complexities in biomedical and econometric decision problems. The authors note that while the current focus is on categorical variables, the framework could be extended to ordinal or continuous treatments and outcomes in future work.