Dr Claude: Monte Carlo tree search plans the diagnostic interview whilst Claude conducts it

Code: github.com/fadynakhla/dr-claude · Devpost: Dr Claude · Thread: x.com/WianStipp

Third place at Anthropic's Claude 2 Hackathon (Shack15, San Francisco, 29–30 July 2023). Built in 24 hours with Fady Nakhla, Arthur Böök and Sukru Kiymaci.

The problem

Given a chief complaint, how do you build a system that asks a sequence of questions to reach a diagnosis with maximum accuracy in as few steps as possible?

This is not a single-shot prediction problem. Each question changes what the next best question is, the answers are uncertain until you ask, and every question has a cost: patient time, clinician time, and the risk of anchoring on a wrong hypothesis. Misdiagnosis is not a rare failure mode. Roughly one in twenty US adults, about 12 million people, experiences an outpatient diagnostic error each year¹. For the dangerous diseases where it matters most (major vascular events, infections and cancers) the error rate is about 11% per case, and the downstream harm is an estimated 795,000 Americans killed or permanently disabled annually². The National Academies put it plainly: most people will experience at least one diagnostic error in their lifetime³. A tool that pathfinds toward the right differential in fewer steps is directly useful.

The two obvious approaches each fail on their own.

• An LLM alone is a strong system 1 thinker. Claude 2 will happily ask a plausible next question, but it has no explicit model of how much each question would actually narrow the differential, no calibrated notion of uncertainty, and no evidence trail you can audit.
• A symbolic system alone (a disease–symptom database plus a search procedure) is a strong system 2 thinker. It can compute exactly which question is most informative under its model, but it has holes in coverage, no understanding of context, and produces actions like ask: C0039070 rather than a conversation.

Dr Claude fuses the two: an AlphaGo-style Monte Carlo tree search plans what to ask over a transparent probabilistic model, and Claude decides which of the planner's suggestions fits the context and how to phrase it to the patient.

The knowledge base

MCTS needs an environment to search over. We used Columbia's Disease-Symptom Knowledge Database, a set of disease–symptom associations mined from New York Presbyterian Hospital discharge summaries (2004), with everything keyed on UMLS codes. Parsing the HTML table gives a mapping from each condition to its associated symptoms.

The database tells you that a symptom is associated with a disease, not how often. To get conditional probabilities we asked Claude 2 to assign each symptom a frequency term for its condition and mapped the terms to weights:

The result is a matrix $M \in [0,1]^{|S| \times |C|}$ with rows indexed by symptoms and columns by conditions, $M_{sc} = p(s \mid c)$. Symptoms that the database does not associate with a condition get a small noise rate rather than zero, so a single unexpected positive does not annihilate a hypothesis outright. This matrix is the entire "dynamics" of the environment.

The probabilistic model

The state of the interview is the set of pertinent positives $S^+$ (symptoms confirmed present) and pertinent negatives $S^-$ (confirmed absent). With a flat prior over conditions and symptoms assumed conditionally independent given the condition, the posterior over conditions is naive Bayes:

normalised over $c$. Given that posterior, the predictive probability that an as-yet-unasked symptom will turn out to be present is

These two equations are the whole inference engine. They are cheap enough to evaluate thousands of times a second, which is what makes tree search viable.

def compute_condition_posterior_flat_prior(matrix, pertinent_positives, pertinent_negatives):
    log_probas = np.zeros(matrix.matrix.shape[1])
    for symptom in pertinent_positives:
        log_probas += np.log(matrix[symptom, :])
    for symptom in pertinent_negatives:
        log_probas += np.log(1 - matrix[symptom, :])
    proba = np.exp(log_probas)
    return proba / proba.sum()

The game

MCTS was designed for two-player games, so we had to phrase diagnosis as a single-player game.

• State: $(S^+, S^-, R, d)$ where $R$ is the set of symptoms not yet asked about and $d$ is a committed diagnosis, or none.
• Actions: the union of two kinds. Ask about any symptom in $R$, or diagnose any condition $c$. The branching factor is therefore $|R| + |C|$.
• Transition: asking about $s$ is stochastic. We sample the answer from the model's own predictive distribution, $\text{yes} \sim \text{Bernoulli}\big(p(s \mid S^+, S^-)\big)$, and move $s$ into $S^+$ or $S^-$ accordingly. The environment is the model. Diagnosing sets $d = c$.
• Terminal: $d$ is set, or $R$ is empty.
• Reward: the posterior probability of the committed diagnosis, $p(d \mid S^+, S^-)$.

def handleSymptom(self, symptom):
    next_self = copy.deepcopy(self)
    next_self.remaining_symptoms.remove(symptom)
    proba = self.getSymptomProbabilityDict()[symptom]
    if proba > random.uniform(0, 1):
        next_self.pertinent_pos.add(symptom)
    else:
        next_self.pertinent_neg.add(symptom)
    return next_self

def getReward(self):
    conditions = compute_condition_posterior_flat_prior(
        self.dynamics, self.pertinent_pos, self.pertinent_neg
    )
    return conditions[self.dynamics.columns[self.diagnosis]]

The reward is what makes this work. Because the payoff for diagnosing is the posterior mass on that diagnosis, the search is rewarded for steering into states where one condition dominates: low-entropy states. Asking about a symptom is not rewarded directly; it is rewarded only insofar as the answer is expected to sharpen the posterior before you commit. Confirming that a patient with malaria-like symptoms has recently travelled to an endemic region makes the subsequent diagnosis "easy", and the tree learns to take that path.

The search

Selection uses standard UCT. For a child with mean return $\bar{Q}$, visit count $n$, and parent visits $N$:

The interesting design choice is the rollout. AlphaGo replaces random rollouts with a learned value network; we had 24 hours and no training data. Random rollouts to a terminal state are cheap per step but wasteful: with hundreds of symptoms, a random policy asks dozens of irrelevant questions before stumbling into a diagnosis, and the returned value is mostly noise.

Instead we used the maximum posterior probability at the leaf as the value estimate:

class ArgMaxDiagnosisRolloutPolicy(RollOutPolicy):
    def __call__(self, state):
        return max(state.getConditionProbabilityDict().values())

This is not a rollout at all; it is a static evaluation function, closer to the hand-written heuristics of early chess engines than to a learned value net. The justification is that $\max_c p(c \mid \cdot)$ is a cheap monotone proxy for negative entropy of the posterior: if one condition already carries most of the mass, the remaining interview is short and the eventual reward is high. It turns every leaf evaluation into two matrix operations, which is why 3 seconds of search per turn was enough for a live demo.

One more departure from textbook MCTS: rather than return the single best root child, we return the top-k (k = 5) children by value. The point of the LLM downstream is to inject context that the model lacks, and it cannot do that if it is handed one action.

def getBestChild(self, node, explorationValue, top_k):
    node_values = []
    for i, child in enumerate(node.children.values()):
        nodeValue = child.totalReward / child.numVisits + explorationValue * math.sqrt(
            2 * math.log(node.numVisits) / child.numVisits
        )
        heapq.heappush(node_values, (-nodeValue, i, child))
        if len(node_values) > top_k:
            heapq.heappop(node_values)
    ...

Where Claude comes in

Claude 2 plays four roles, each a separate LangChain chain with its own prompt, and none of them is allowed to do the planning.

Decision Claude receives the confirmed positives, the confirmed negatives, and the top-5 symptoms from MCTS, and must return exactly one of the five. This is the fusion point. The symbolic planner has no idea that "night sweats" is an odd thing to ask a patient who just described a broken ankle, or that two of its five suggestions are near-synonyms; Claude does. But Claude cannot wander off and ask about something the evidence does not support, because its choice set is the planner's output. The prompt is blunt about it:

The intelligent system has predicted that the following symptoms are the most valuable to confirm or reject next: {symptoms}. Your job is to determine which one of the symptoms from the intelligent system that we should inquire about next. […] Answer by quoting only the name of the symptom.

Doctor Claude turns a symptom name into a single direct question ("Are you experiencing any shortness of breath?"). Nothing else.

Patient Claude exists for the demo. It is given a patient chart and told it is playing a game: answer only the question asked, answer "no" to anything not in the chart, and never volunteer information. This let us run full interviews against synthetic charts without a human in the loop, and it turned out to be the hardest prompt to get right. Claude wants to be helpful; a good simulated patient must be unhelpful in a very specific way.

The matcher closes the loop from free text back to the model's vocabulary. Given the doctor's question and the patient's answer, Claude extracts the symptoms mentioned and whether each is present. Each extracted string is then embedded, searched against a FAISS index of every symptom name in the knowledge base, and Claude picks the best match from the retrieved candidates. Only then does it become a UMLS-coded symptom that can be added to $S^+$ or $S^-$. The same machinery, with a different prompt, extracts the initial chief complaint from the presenting note.

The loop

Putting it together, one turn of Dr Claude looks like this:

1. Extract chief complaints from the patient note; add them to $S^+$.
2. Run MCTS for 3 seconds from the current state; take the top-5 root actions.
3. If the best action is a diagnose action, stop and return it.
4. Otherwise, Decision Claude picks one symptom from the five; Doctor Claude phrases it; the patient answers; the matcher maps the answer back to symptoms and updates $S^+$, $S^-$.
5. Stream the current top-5 differential with probabilities to the UI (the "brain" panel), so the clinician can see the model's belief evolve with each answer.

The loop terminates when MCTS itself proposes a diagnosis, when any condition's posterior exceeds 0.8 ("Is it X?"), or after ten questions with no resolution ("I'm sorry, I'm not sure"). Everything runs behind a FastAPI websocket with a React chat client in front.

What is wrong with it

Twenty-four hours is not long, and there were a ton of things we could have done better.

• Flat prior. Every condition starts equally likely. Real diagnosis leans heavily on prevalence; a proper prior would change the early questions substantially.
• Naive Bayes. Symptoms are treated as conditionally independent given the condition. Fever and chills are not independent. This inflates confidence when correlated symptoms co-occur.
• Frozen chance nodes. The MCTS library keys child nodes by action, so when a symptom question is first expanded the yes/no answer is sampled once and baked into that child forever. Subsequent visits to "ask about fever" all inherit the same sampled answer. Correct treatment needs explicit chance nodes (expectimax-style) or open-loop search. This is the "improve MCTS expansion" item in the repo's to-do list, and it is the biggest algorithmic flaw.
• The noise floor is a hidden hyperparameter. A pertinent positive on a symptom the database does not list for a condition multiplies that condition's likelihood by 0.03. Given the database's coverage holes, that is an aggressive elimination rule. A patient with a real but unlisted symptom can knock out the right diagnosis in one step.
• The value function is a heuristic. $\max_c p(c \mid \cdot)$ rewards confidence, not correctness. A miscalibrated model that becomes confidently wrong is rewarded exactly like one that becomes confidently right. Proper information-gain or a learned value function would be better; both need data we did not have.
• Branching factor. With all conditions available as immediate actions at every node, the tree is wide and shallow. Three seconds of search on deep-copied Python objects does not see far.
• Discounting was switched off. The state class supports a per-question discount to penalise long interviews, but the demo ran with a discount rate of $10^{-9}$. Question count was capped by the ten-question limit instead.
• The patient is Claude. Real patients are vaguer, more contradictory and more informative than our simulated one. The matcher's extraction step is where a real deployment would break first.

Why it still matters

The result was less interesting than the pattern. A language model on its own is a fluent system 1: fast, contextual, and unable to tell you how sure it should be. A search procedure over an explicit model is a rigid system 2: calibrated within its model, blind outside it. Giving the planner authority over what is worth asking and the LLM authority over which of those to ask, and how produced a system that was more accurate than either and, importantly, auditable. Every question Dr Claude asked could be traced to a posterior and a set of database rows.

The specific mechanisms have all improved since then: better disease models, proper handling of stochastic transitions, learned value functions, tool-using models that can call the planner themselves. The division of labour has not changed. When an LLM needs to act well over many steps under uncertainty, let something that can count do the counting.

Further reading

• Dr Claude on Devpost · Source on GitHub
• Disease-Symptom Knowledge Database (Wang, Chase, Markatou, Hripcsak & Friedman, Columbia DBMI)
• Mastering the game of Go with deep neural networks and tree search (Silver et al., 2016)
• Bandit based Monte-Carlo Planning (Kocsis & Szepesvári, 2006), the UCT paper
• Thinking, Fast and Slow (Daniel Kahneman, 2011)
• Anthropic's hackathon announcement of the winners