An orthogonal learner for individualized outcomes In markov decision processes

0:00Welcome back to the Deep Dive. Today we are going to tackle a problem that keeps statisticians awake at night. But, and this is key, it's also a problem that literally determines life or death. It sounds like an exaggeration, but for once it really isn't. We're talking about the mathematics of survival. We're diving into the world of personalized medicine, but not the simple kind. You know, where you take a genetic test and the doctor says, take this pill. We are talking about the hard stuff. The really hard stuff. Long-term treatment strategies. I mean, think about a cancer patient or someone with a chronic autoimmune disease.

0:35As a doctor, you aren't making a single decision. You're designing a protocol, a sequence. Dosage today, check the blood work tomorrow, maybe pause the treatment next week. It all depends. Then you adjust the dosage again the week after. That's a chess game. It is. It's a trajectory. And the tricky part, the part that makes this so hard to model, is that every single decision you make changes the patient's state. And that new state then limits or expands your options for the next decision. It's a feedback loop that just stretches out over time. And here is the kicker. We have mountains of data on this.

1:12I mean, hospitals have petabytes of electronic health records showing exactly how doctors have treated patients in the past. Yeah. So the intuitive thought is, great, let's just feed that into an AI and find the cure. If only it were that simple. Why isn't it? We have the data. Well, we have data on what doctors did do. That's called the behavioral policy. But to improve treatment, we need to know what would happen if we did something different. A new strategy. We want to test a new strategy, an evaluation policy. And here's the snag. We cannot run A-B tests on stage four cancer patients. You can't just say, OK, Group A gets the standard chemo, Group B, well, for you, we're going to try this totally random, unproven dosage schedule just to see what happens.

1:54Exactly. I mean, that would be profoundly unethical. You can't experiment with people's lives like that. So we're stuck. We have to learn off policy. Meaning we have to look at the historical data of the old strategy and try to mathematically predict the outcome of a new one. Yes. And this is where we run into the villain of today's story. It has a very dramatic name, which I love. Uh-huh. What is it? It's called The Curse of the Horizon. The Curse of the Horizon. It sounds like a pirate movie. It does, doesn't it? Yeah. But in statistics, the horizon just means time. It's how far into the future you need to predict.

2:29And the curse is that small errors in our understanding of the present, they don't just stay small. They grow. They compound. As you project them forward step by step, they get bigger and bigger. So if your map is off by like an inch today, you might be off by 10 miles a month from now. Precisely. And in a medical context, being 10 miles off means the algorithm predicts a patient will recover. But in reality, the treatment is toxic. Okay, that's a serious problem. A huge problem. So the mission of this deep dive is to unpack a new method, a breakthrough, really called the DRQ learner. DRQ learner.

3:05It stands for double robust Q learning. And it's a method that combines two heavy hitting fields, causal inference and reinforcement learning. And it does it in a way that is mathematically designed to be immune or orthogonal to those compounding errors. Okay, let's unpack this. Because usually when we talk about AI in medicine, we just think feed data to a neural network, get an answer. But you're saying that brute force approach doesn't work here. Not reliably, no. To understand why, we have to look at the landscape. We are operating in what's called an MDP, a Markov decision process. I remember this from our logic episodes, but give us the refresher.

3:41How does an MDP map onto, say, a hospital room? Sure. An MDP has four main parts. First, you have the state. In our case, that's the patient's health, their vitals, tumor size, blood pressure. Everything that describes how they are right now. Exactly. Second, you have the action. That's the intervention, the drug dosage, the surgery, the decision to wait. Okay, state, action. Third, you have the reward. Did the tumor shrink? Did the patient live longer? And the fourth. The transition. This is the mechanism of the body itself. If I take action A in state B, where do I end up? How does the body respond?

4:17And the goal of all this math is to find something called the Q function. The Q function is the Holy Grail. It's a single number that tells you, If I take this specific action in this specific state and then follow my new policy forever after, what is the total cumulative reward I can expect? So if you know the Q function perfectly, you've effectively solved the disease. You know exactly which move leads to the best long-term outcome. Exactly. You'd have a perfect roadmap. So here's my question again. If we have the history, we have the states, the actions, and the rewards from thousands of past patients, Why is calculating this Q function so hard?

4:56Why can't we just, I don't know, average it out? Because we don't know the ground truth. We don't know the perfect physics of the human body. We only have estimates based on noisy data. The status quo, the way this is usually done, relies on what we call plug-in learners. That's convenient. Just plug it in. It is convenient, and that's why people use it. The idea is you use machine learning to estimate the moving parts, the nuisances, they're called, and you just plug those estimates into an equation to get your Q function. Give me an example of a nuisance. A big one is propensity. What was the probability that the doctor and the historical data gave that specific drug?

5:33You estimate that and then you use it to reweight the data. There's a method called Q regression with inverse propensity weighting or IPTW. Inverse propensity weighting. That basically means if a patient in the database looks like the patient I want to treat, but they were treated rarely, I should weight their data more heavily to make up for it. Right. You are re-weighting the old population to look like the new one. It makes sense, you know, intuitively. But remember, the curse of the horizon. The compounding errors. Exactly. These weights are cumulative. If you're looking just one step ahead, it's probably fine.

6:09But if you're planning a treatment protocol for 12 weeks, you are multiplying 12 weights together. And those probabilities are small. Right. Like, say, 0.5 times pure 5 times fewer 0.5. It gets tiny very, very fast. And in the math, you end up dividing by these tiny numbers. The weights explode. You get massive variance. The equation becomes incredibly unstable. So the plug-in approach fails because it assumes our initial estimates are good enough to support a 12-story building. but actually they crumble under the weight. That's a great analogy. It suffers from what's called plug-in bias. It's a first-order error.

6:44If your initial guess about the propensity is just slightly wrong, that error passes directly into your final answer. And over time, it snowballs into a disaster. Which brings us to the hero of the day, the DRQ learner. Now, is this just a smarter neural network? No, and that's a really important distinction. The DRQ learner isn't a model like, you know, GPT-4 or a ResNet. It's a meta-learner. a meta-learner. Think of it as a wrapper or a manager. A manager. Yeah, it manages the learning process. You can put whatever machine learning model you want inside it, neural networks, random forests, anything.

7:16The DRQ learner provides the mathematical rules for how those models should learn to make sure the final result is valid. So it's the rules of engagement, not the weapon itself. Precisely. And it claims to solve the problem by bridging two worlds, causal inference and MDPs. Let's talk about why it works. The research highlights three pillars of reliability. I want to go through these because they sound like fancy buzzwords, but the concepts are actually really cool. The first one is double robustness. Right. This is your safety net. To solve this problem using causal inference, you essentially have to model two things.

7:54First, the propensity. How likely was the doctor to take that action? Right. And second, the outcome. What actually happens to the patient. Okay, propensity and outcome. Two models. Usually if either of those models is wrong, your final answer is wrong. It's a fragile chain. But with double robustness, the math is set up so that if you mess up the propensity model but get the outcome model right, you still get the correct answer. And vice versa. Yes. If your outcome model is a total mess but you nailed the propensity, you're still good. You only need to be right about one of the two. That seems incredibly forgiving.

8:26Usually in math, if you get one variable wrong, The whole thing's wrong. It is forgiving. It takes the pressure off having a perfect model for everything. But the second pillar is the real heavy hitter. It's what breaks the curse of the horizon. It's called Naaman orthogonality. Naaman orthogonality. That sounds like something I'd hear in a quantum physics lecture. Huh. It does sound intimidating, but think of it as a shield. A shield against what? Against the shakes. Imagine you were trying to tune an old school analog radio to a specific station. You know, the ones with the dial. Okay, yeah, I'm turning the dial.

9:01They're static. With the standard plug-in learners we talked about, that radio is incredibly sensitive. If your hand shakes even a microscopic amount, and in statistics, estimation error is that shaking hand, you lose the signal completely. The error in your hand movement translates directly to error in the audio. Yes. And since we are estimating things from messy medical data, our hand is always shaking. We never have perfect data. Right. We're always shaking. Name and orthogonality is like having a stabilizer on that dial. Even if your hand shakes a little, the frequency doesn't drift. How is that even possible?

9:37Well, the method is mathematically designed so that the gradient of the loss, with respect to those errors, is zero. That's the orthogonal part. So the direction of the error is perpendicular to the direction of the answer. You got it. So local perturbations, small mistakes in the setup, don't throw off the final result. It shields the final answer from the noise of the initial steps. That is huge. That is huge when you're dealing with noisy data like EHRs. Absolutely. Okay, so we have a safety net and a shield. Then there's the third pillar, quasi-oracle efficiency. I love that name. It sounds like a character in The Matrix.

10:13I am the quasi-oracle. It basically is a superpower for an algorithm. In statistics, an oracle is a theoretical thing that knows the ground truth. truth. The oracle knows the perfect propensity and the perfect outcome functions without ever needing to look at data. So we wouldn't need to do any of this if we were oracles. Right. But we aren't. We have to estimate those functions using data, which usually slows us down. It makes our final answer less precise. Quasi-oracle efficiency means that this method converges to the truth at the same rate as if we were the oracle. Wait, so we aren't penalized for guessing?

10:50Antsymptautically, no. We aren't penalized for having to estimate the messy parts first. The algorithm learns just as fast and just as accurately as if we had been given the cheat codes to the universe. That is a bold claim. So how does it actually do this? What is the machine doing step by step? It operates in a two-stage process. And this structure is really the key. Walk us through stage one. Stage one is the rough draft. The system looks at the data and it estimates those nuisances we talked about. It looks at the behavioral policy, what happened in the data. It estimates the density ratio, how different is our new plan from the old plan.

11:23Okay. And it makes a rough first guess of the Q function. So stage one is basically doing what the old methods did. Just do your best with standard tools. Yeah, correct. And these estimates might be biased. They might be shaky. That's expected. Because then comes stage two, the refinement. This is where the magic happens. This is where they use something called the efficient influence function, or EIF. Unpack that. Think of the EIF as a sophisticated error detector. It looks at the estimates from stage one and it calculates exactly how biased they likely are. Then it derives a specific loss function that essentially subtracts that expected error.

11:59It's like noise-canceling headphones. That's a perfect analogy. Noise-canceling headphones listen to the ambient noise and generate an anti-noise sound wave to cancel it out. The EIF allows the DRQ learner to calculate the anti-bias. It constructs a new target that's mathematically orthogonal to the errors of the first stage. Exactly. So even if the inputs from stage one were shaky, the stage two process cancels out the shake. And the goal is to get to the individualized potential outcome. We aren't just asking, does this drug work on average? We are calculating the value for this specific patient.

12:32Okay, this all sounds theoretically airtight. But theory is one thing. I want to know if it works when the rubber meets the road. Did they test this? They did. They used a classic validation environment called the taxi environment. Taxi environment. Like a yellow cab. Yeah, it's a standard benchmark in AI. Imagine an agent, a taxi navigating a grid world. It has to pick up passengers and drop them off. It sounds simple, but it's a sequential decision problem. You move north, south, east, west. And every move changes your state. Exactly. It's a simplified version of the doctor's dilemma, navigating a path to a goal.

13:08So how did they test the curse of the horizon with a taxi? They made the simulation incredibly difficult in two ways. First, they varied the overlap. What's overlap? Overlap measures how similar the new policy you want to learn is to the old data you have. High overlap is easy. It's like predicting traffic on a Tuesday when you have data from last Tuesday. Low overlap is hard. So low overlap is like saying, I want you to drive like a Formula One driver, but I'm only showing you data of a grandma driving to church. Yes, exactly. And second, they varied the horizon. They forced the taxi to plan 20 moves ahead instead of just three.

13:44So hard data and a long timeline. How did the DRQ learner stack up against the standard methods? It crushed them. They compared it against the standard tools. Q regression, FQE, minimax, Q learning. When the overlap was low, when the new policy was radically different from the old data, the standard plug-in methods basically crashed. They couldn't handle the gram-out-of-formula-one conversion? No. Their error rates skyrocketed because those probability weights we talked about earlier just exploded. But the DRQ learner remained stable. The error rates stayed low. That's the name in orthogonality kicking in.

14:19The shield held up. It did. And the same thing happened with the horizon. As they forced the algorithms to look further into the future, the gap between DRQ and the others just got wider and wider. The other methods succumbed to the curse of the horizon. And DRQ broke it. It broke it. Now, does this only work for grid worlds? Because real life isn't a chessboard. It's continuous. Blood pressure isn't just high or low, it's 120 over 80. That's a key feature of the method. The DRQ learner works for both discrete states, like the taxi grid, and continuous states, like real-world medical measurements.

14:54And because it's a meta learner, you can put deep neural networks inside it to handle that complexity. So let's zoom out. We started with cancer treatment. What does the existence of this tool actually unlock for us? It means we might finally be able to utilize the history we have. Hospitals are sitting on mountains of data records of what happened to millions of patients. Until now, using that data to discover new, better long-term treatment protocols was just incredibly risky. We were afraid the algorithms would hallucinate a cure because of a math error five steps back. We were afraid of the shaking hand on the radio dial causing us to hear a signal that wasn't there.

15:32And this fixes the shaking hand. This method gives us a level of mathematical safety we didn't have before. It allows us to ask what if, questions about long-term treatments, complex multi-year strategies, and get answers that are statistically valid without having to run dangerous experiments on patients first. It turns observational data into a laboratory. That is a massive shift. It really is. It's the difference between guessing and knowing. Here is where it gets really interesting for me. If we can now trust algorithms to learn individualized long-term outcomes from messy, biased historical data, this isn't just about medicine, is it?

16:13Oh, definitely not. I'm thinking about economics, policymaking. That's exactly where my mind goes. Think about setting interest rates or tax policy. Those are sequential decisions. A decision the Federal Reserve makes today impacts inflation next year, which then impacts the decision they make then. And just like with patients, you can't run A-B tests on the national economy. Let's crash the GDP of Ohio just to see if this stimulus package works. Yeah, you can't do that. We only have historical data. And that data is biased by whatever policy was in place at the time. Exactly. The curse of the horizon is everywhere in climate policy, too.

16:46We need to make decisions now about carbon emissions that target outcomes 50 years from now. That is the ultimate long horizon. If the DRQ learner can break the curse for a cancer patient's trajectory, maybe it can break the curse for the planet's trajectory. It raises a provocative question. What other fields have been stalled because they couldn't trust their historical data? Maybe the bottleneck wasn't the data itself, but the lack of a tool like this. A tool that knows how to cancel out the noise. A mathematical shield against our own uncertainty. I like that image. It's powerful. It changes the way we look at the past.

17:21The past isn't just a record anymore. It's a simulation engine for the future. Well, that is a lot to process, but I feel like I understand the stakes and the solution much better. Always more to learn. That's it for this deep dive. We'll see you in the next one.