Building Medical AI at Scale: TAO, Synthetic Data, and LLM Judge Architecture

Summary

Flo Health, a women's health application with 81 million monthly active users, built 60 LLM judges with priority-weighted aggregation to evaluate AI responses for both medical safety and usefulness simultaneously, where safety scores always dominate.
Using TAO (Test-Time Adaptive Optimization) on Databricks AI Runtime with privacy-first synthetic training data — no real patient data — they fine-tuned open models and doubled both medical accuracy and usefulness while reducing inference costs.
The video covers multi-objective optimization, preference-pair generation, and governance patterns applicable to any domain where partial safety is unacceptable, along with the full fine-tuning workflow on Databricks AI Runtime.

Building Medical AI at Scale: TAO, Synthetic Data, and LLM Judge Architecture

Building medical AI requires satisfying multiple safety rules simultaneously while protecting user privacy. Flo Health's challenge: achieve medical accuracy at scale without expensive expert labeling. Their solution: create 60 LLM judges with priority-weighted aggregation where safety dominates usefulness. Using TAO (Test-Time Adaptive Optimization) on Databricks AI Runtime, they fine-tuned open models without real patient data, doubling both medical accuracy and usefulness while cutting inference costs.
Learn how Flo engineered their judge hierarchy, synthesized balanced training data across health topics, and implemented iterative refinement for robust answers. Vladislav Nedosekin, Michael Shtelma, and Andras Meczner explain multi-objective optimization, preference-pair generation, and governance patterns that scale medical AI safety, directly applicable to any domain where 'mostly safe' isn't acceptable.
🤝

Chapters

FAQs

What is TAO (Test-Time Adaptive Optimization) and how did Flo Health use it?

TAO, or Test-Time Adaptive Optimization, is a fine-tuning methodology available on Databricks AI Runtime that allows models to be improved iteratively without requiring large labeled datasets of real user data. Flo Health used TAO with synthetic training data to fine-tune open models for medical safety and usefulness while ensuring no real patient data was involved.

How does Flo Health's 60-judge LLM evaluation hierarchy work?

Flo Health created 60 specialized LLM judges that each evaluate a specific dimension of response quality, such as different categories of medical risk or aspects of usefulness. The judges are aggregated using a priority-weighted system where safety rules always dominate, ensuring that a highly useful but medically unsafe answer is never surfaced to users.

Why did Flo Health use synthetic data instead of real patient data for model training?

With 81 million monthly active users and over half a billion downloads, Flo Health operates under strict privacy requirements that make using real patient health data in model training unacceptable. They synthesized balanced training data across health topics to create realistic preference pairs for fine-tuning without exposing any real user information.

What results did Flo Health achieve with their fine-tuned medical AI models?

Flo Health doubled both medical accuracy and usefulness compared to their baseline models after applying TAO fine-tuning with their 60-judge evaluation framework on Databricks AI Runtime. The fine-tuning also reduced inference costs, making the improved models economically viable to serve at the scale of 81 million monthly active users.

Full transcript

[00:08] Okay. Thank you everyone for coming in today. In today's session, we will walk you through Flo Health's journey on how we doubled medical safety in foundational model.
[00:25] My name is Vladislav Nityagov. I am director of engineering at Flo Health for data and AI platform. And I'm joined today with two of my colleagues. Hi everyone. My name is Michael Shteyman. I am a lead product specialist for GenAI at Databricks. Hi, my name is Andras Schmelczer. I'm
[00:42] medical director for accuracy and safety at Flo Health. Thank you Michael. Thank you Andras. So in today's talk, we will cover four pillars of our journey. We will start from the business challenge. Essentially, what driven us
[00:59] to actually embark on this journey. We will talk about how did we evaluated results of our work and set what the final result or what good look like. Then we obviously will share some juicy
[01:15] insights on how did we approach fine-tuning. And also cover the actual outcomes so you would know how and what we did. And also share some insights on what is next for us.
[01:32] Before we actually dive into business challenge, I'd like to double-check. Can you raise your hand if you know what Flo actually is? Okay, we have about like 30-40% of the audience who knows. For those who don't know, Flo is one of
[01:48] the biggest women health and well-being applications globally. We have about 81 million monthly active users. So any feature, any functionality that we deliver, AI or not AI, has to be thoroughly thought through
[02:04] because we'll be delivering it at scale. When we think about functionality with AI, we also have to take into account all the users that currently install in Flo to evaluate or onboarding.
[02:20] And to put this into perspective, our application currently have more than half of half a billion downloads. So, our potential audience is quite big. At the very beginning of our journey, when we only started to build an
[02:35] application 10 years ago, a lot of the things that we did was rule-based, scripted. Essentially, there was a lot of interesting things that were working, but they were not flexible. They were not as innovative as people
[02:51] demand right now. At some point, we started adoption of the AI and particularly of generative AI. And we run several experiments with a single conversational AI. We got some really good metrics. It was quite a successful endeavor, but we
[03:08] still realized that it's not enough. It can't cover for everything that we have in mind and for everything that our user demands. And we decided to move to a more compound system where we can leverage traditional ML to do the number
[03:23] crunching. We can leverage generative AI to cover different aspects. And we also realized that we can't just rely on one LLM or one chatbot to cater for variety of the topics that our customers need. So, we decided to build a compound
[03:39] system. When we thought about our compound system, we had to define for ourselves what will be our good look like. Flo as an application built a reputation on medical accuracy and medical
[03:54] correctness. We do a lot of research. We do a lot of articles to educate people. Um and that's why when we've been thinking about our AI capabilities, we set medical safety as our foundation.
[04:12] However, it was not just enough to make it medically accurate and medically safe. We also needed to ensure that it's useful for our audience, that our audience likes the functionality and they will be willing to use it. So, we also set usefulness as our goal.
[04:32] This is a slide that we actually prepared for the presentation. However, I was actually super tempted to change the slide. Can you raise your hand if you know what Fable is and what happened with it few days ago?
[04:48] Yeah, it's actually quite surprising that not everyone raised their hand, but fair enough. Again, for those who don't know, Fable was one of the most powerful LLMs that was available on the market released by Anthropic. And then suddenly people woke up to only realize that it's no longer available.
[05:04] So, if you build your pipeline, if you build your application using Fable, there was no notification, no grace period, suddenly it disappeared. And what it would mean for your business, your business will stop. We actually don't want that.
[05:22] However, there are other factors apart from actually control of the life cycle of your foundational model. As I previously mentioned, we cater for at least 80 million users. In reality, it's significantly more. So, doing inference at such scale
[05:38] cost a lot of money. And if you were to build more tailored, more lighter LLMs, your inference costs are significantly decreasing. More importantly, you can build as many agents or as many components in your system as you like, and everyone will be
[05:56] a specialist in a particular area. We also can then embed domain expertise. We have quite a big medical team uh that is working with us on building these capabilities, who are and experts in different areas. And we
[06:12] can put this expertise into our AI capabilities. And lastly, we own the life cycle. So, we don't really care if Anthropic or US government will decide to rework another model, because we can run our fine-tuned models ourselves.
[06:31] Before we actually started the fine-tuning, we actually thought about how we're going to prepare the data and how we're going to do the evaluations. So, we started from actually defining our evaluation framework. We needed the framework to understand what exactly we are going to do to,
[06:49] first of all, define our target state, define the success, and then how we're going to measure it. If you think about their fine-tuning, usually, you have to prepare a lot of data sets. And in our case, preparation of data
[07:05] sets would have cost a lot of money. If you want to hire uh a professional uh gynecologist to help you with data labeling, you will pay something about like 4 to $600 an hour for the tedious work. Instead of doing that, we decided to
[07:22] invest into AI judges. We build them to help us to essentially assess the outcomes of their artifacts produced by AI. And then we use those judges to do the actual fine-tuning, as well as to
[07:38] validate the results. Now, we will tell you how did we approach medical safety and I will pass to Andras, who will give you a bit more insights on what we did in that area.
[07:53] Thank you, Vlad. So, as Vlad mentioned, evaluation was our first task to do. Well, we had a bit of a problem because we didn't have a model to evaluate, but we needed to still build this framework and
[08:09] we need the framework to build our judges to build a model. So, a bit of a chicken and egg situation. So, what we did, we took Gemini 2.5 Flash model, which we prompt engineered, and then we had a kind of a tool to be able to build our evaluation framework.
[08:26] So, what do we need to evaluate something well? So, we need relevant test cases and the mechanism, an agent that can detect failures. We would have loved to work with an out-of-the-box solution, but it didn't really exist for our particular purpose.
[08:42] The reason is that we dealt with open-ended questions, we were consumer-facing, we didn't want to give medical advice, we were in this women's health domain with a lot of cycle data, which we knew about the users, and we had this strategic goal that in
[08:59] the end we wanted to have the judges, which we can use for fine-tuning. So, this is what we did. Firstly, we did a traditional risk assessment, which is very normal in clinical informatics. So, in a structured way, we thought through what can go wrong. So,
[09:16] to give you an example, even though this was a chatbot for women's health, we didn't want people who, let's say, tell that they want to commit suicide that we encourage them to do so. So, we came up with certain rubrics that we want to assess that if someone is is telling about that, we we deal with it
[09:33] appropriately. But, there's only a certain amount of things you can kind of uh identify this way. So, what we did, we started to test our model. So, we kept testing and testing and identifying failures
[09:48] and we clustered those failures into themes. So, an example here um that uh why am I spotting? And uh the answer is is is actually true, but it's very benign. Uh spotting can have other causes, underlying conditions, but the
[10:03] model didn't include that into the answer. So, we realized that there is this theme that a lot of time when uh questions are about the causes of a certain symptom, we always give something very benign or very serious, but we don't give a balanced answer. So, that's how this rule uh was born,
[10:20] the balanced causes rule, which um then became a rubric and subsequently a judge as well. So, we kept testing and testing. Actually, we tested all the way till 6,800 cases until we didn't find any more themes. And then we felt our annotation guides,
[10:37] our set of rubrics, are complete, at least for now. Uh and then we started to build judges. So, for us, a judge consists of a definition, a grading prompt, and training cases. And uh we decided to go for an macro F1
[10:54] score of 75% because we saw in the literature that that's kind of what uh is a good performance at the time. But then as I mentioned, we need good test cases as well. So, we had a lot of synthetic test cases, so 2,200 synthetic cases, which were mainly
[11:11] adversarial cases or typical cases, challenging cases. But to really assess whether um the model works well, you need to see whether on the real world it works well as well. So, after we felt the model is safe enough, we released to 10,000 beta users
[11:26] uh and we collected their data anonymized and started to evaluate it. So, we evaluated in two ways. So, once we had the doctors who annotated using the annotation guide, but at the same time, you also run these judges and see what the LLM judge gives
[11:41] as a verdict. If the doctors and the judges contradicted, we took those cases, gave it to three blinded human doctor human annotators, and they established the ground truth. And then after having the ground truth
[11:57] for these 12,000 cases, we could actually see that our judge is not as good as we thought. So, then we had to calibrate them. We split the data set into two. We used 1/3 of the data to calibrate these judges, and we kept the other 2/3 to assess the real performance.
[12:13] Um and we kept iterating, and then in the end, we managed to reach everywhere more than 75% F1 score, which we were happy with. And then we handed over our judges to our colleagues, Michael and Vlad.
[12:36] Thank you, Andras. Building medical safety was our foundation, but as I previously mentioned, it was not enough. So, we also paid a lot of attention to usefulness. If you think about the conversational AI experience,
[12:52] you may have a variety of answers to the same question. And to Andras' point, you might be missing balanced causes, but also, your answer may not necessarily be fulfilling.
[13:10] So, we looked at different answers that were given to us by LLMs, and similarly to our medical judges, we decided to define rubric from a product side. What good look like from a user perspective.
[13:27] This rubric was also split into three tiers or three levels. Essentially, our product, similar as a medical team, set level one, something that cannot be negotiated. And that included simple things like
[13:43] relevance of the answer, if the answer was complete or not. Then we looked at additional components as tone of voice or coherence or if we actually gave a personalized answer. So, we moved in iterations to avoid
[13:59] being stuck trying to build, you know, the spaceship. We tried to go in iterative mode, covering one rubric, calibrating those judges, and then moving on to the next wave of the judges. Once this framework was established, we
[14:15] could, similarly to our medical team, assess from the product side if answers are actually meeting our requirements or not. And we use this data to then find, essentially, a balance between usefulness and medical safety.
[14:33] If you think about conversational AI answering your question and referring you to see medical professional all the time, that will be 100% medically safe answer. There is nothing that could go wrong. You go and see a doctor.
[14:48] But would you find it useful? Probably not. You will find it tiring. Every time you go to the app, you'll already predict the answer. Why should I go to the app? It will tell me to go to a doctor. So, ultimately, using the judge reasoning, we found a way to structure
[15:04] the answers in a way that would satisfy both medical and usefulness judges, meaning we get both both useful and medically accurate answer. Once we finished with all of this, we decided to apply all of this to
[15:19] fine-tuning. And at this time I will pass to Michael who will walk you through all the technical know-how on what we did during our fine-tuning journey. Thanks, Vlad. Okay, so now let's Oh, yeah, thank you. So, now let's discuss how we actually
[15:35] run the fine-tuning. Um for the fine-tuning we have used Tau methodology developed by the Databricks research. Tau stands for test-time adaptive optimization. So, let me at first explain what that is. So,
[15:52] uh it's based on a very simple idea. So, let's say if you generate a lot of answers or a lot of responses for one prompt and after that use some mechanism to pick the best one, we can significantly improve the metric we care for. In our case, it's medical safety
[16:09] and usefulness, but you can use exactly the same for pretty much anything. It can be code quality, it can be helpfulness, it can be anything else. But, we cannot run it in production in that way because it first is going to
[16:24] like it will raise the costs and it will also eat into our latency. And here, as my colleagues have mentioned, we had uh a lot of like we have a lot of judges, so it will also take time to run all of them. So, it's not really an option
[16:40] in that setup at least. So, Tau helps us to bake that like direction, bake that knowledge into the model using fine-tuning. So, how do we do it? First, as I mentioned, we generate
[16:56] n responses per prompt. So, we have our prompt data set, we generate our n responses. After that, we apply our reward function. Uh this can be like static analysis, so it can be like a verifier. It can be LLM judges, it can
[17:12] be both things together, right? It can be really anything you want. And after that, after we have our rewards, we can depending on what optimization algorithm we're going to use, after that we can either formulate, uh, build the
[17:28] preference pairs, and then run something like DPO or similar algorithms, or we can build, a dataset which we can later use for offline RL algorithms like A star PO. So, in that case we would need, uh, prompt response and the reward.
[17:44] And important is that it's an iterative process. We run it as many times as, uh, we see that the model improves. So, basically we stop iterating when we see that there is no improvement.
[18:01] Okay. So, now, uh, let me tell you how we have, basically how we have, uh, adapted it to our current, use case. At first, we did not have any problems, so we didn't have any questions, so we needed to generate them synthetically.
[18:19] After that, uh, we have, uh, generated candidate, uh, answers. Uh, we had two approaches for that. And after that, we have run all the judges. We have run the evaluation exactly in the same in the way that, uh,
[18:34] Andre and Vlad have, uh, just told you. And then we uh, need to build two datasets, one for the supervised fine-tuning, and another one for the preference tuning. And now, Vlad is going to talk about the
[18:50] answer generation. Sorry, question generation. So, one of the things that we really, uh, paying attention to is privacy. Essentially, Flo being, uh, a health app or well-being app,
[19:07] encourage users to, uh, put a lot of data into this application. So, we can essentially help with insights, with other recommendations for the users. And we need to ensure that our audience is actually comfortable with use of the
[19:24] app, with trusting us with this data. So, we didn't want to use the data for the fine-tuning in a clear way. To not to avoid to avoid any mistrust or misperception of what we do.
[19:40] The way we approached our data generation or data synthesis is we took the interactions of the user and then asked LLM to generate messages based on the themes or
[19:56] conversations or topics that our users have engaged with. After we generated those questions and answers, we also did analysis if the dataset is balanced across different
[20:13] topics that we want our fine-tuning model to cover. So, we checked if there is a presence or enough text presence of their cycle conversations or cycle questions, different symptoms, sleeps, and activities. For all those generations, we didn't use
[20:29] any of the free text input that our users have ever provided to the app. What we wanted to avoid is model to learn anything specific about any of the individuals. We want only to ensure that the model understands the patterns, dependencies
[20:45] between different characteristics of the user, but not something that will be super specific to the user. And then once we did all of this, we also generated user-style questions.
[21:01] You can see an example at the bottom of the slide. And you You also see that both iterations have been completed by AI. Not taken by the humans. Last step was again to rebalance our question data set.
[21:17] We didn't want the model to either become an expert in a cycle and completely be inadequate for a sleep conversations or for sleep insights. So we again did a balancing on a combination of the attributes to
[21:33] ensure that when we do run the fine-tuning exercise, the model will learn correct patterns. Michael, can you also walk us through how did we work with the answers? Sure.
[21:49] Uh thanks a lot. So now let's discuss how we generated our So at first you know we are Then good. Okay, is it better? Oh yeah, it's better now. Okay, great. So let's now uh
[22:05] discuss how we have generated our answers. At first we have tried to very simple, let's say diversity sampling approach. What we did, we have just generated a lot of answers for each question. Uh to boost diversity even
[22:20] more, we have built different bands, uh different temperature bands and then sampled random temperatures from each bands. So to like get really diverse questions. But in that way we were not able to get
[22:37] valid answers according to all possible judges. Uh we have lots of them uh for all the questions. That's why we have introduced a different approach, which we have talked which we have used on top of the first one. So we have introduced iterations, so or iterative refinement. So what we did, we
[22:54] have generated a answer and after that we have run it through all the judges. And for the judges which the like which failed, we have taken the justifications and asked exactly the same model to
[23:10] improve the question. Sorry, to improve the answer. And we did it in the iteration. We I think we were using three iterations and then we just picked the best one and used it as an output. Um iterative refinement of course gives
[23:26] us significantly more valid answers, but uh those answers are of course becoming out of distribution because we are using the um input coming from the judges which happens to be a another very powerful LLM. Um but still uh if it's like not huge
[23:43] percentage, it's fine. Um Okay.
[23:58] Okay, so after we have generated all the answers, we have uh run the evaluation for all answers coming from different uh So, coming from the first and second approach all together because we are not sure that uh
[24:15] like of course the answers which are coming directly from the diversity sampling the the like lots of them actually most of them will be invalid. And uh still some answers coming from the uh iterative refinement still could be invalid. That's why we are running the
[24:30] uh evaluation for all of them. And how we did it, we have uh used our medical safety judges as the constraints and used the usefulness as the metric we want to hill climb.
[24:46] And at first, basically that uh did not work. So, actually it worked, but the result was quite interesting. So, our model became very useful, but the medical safety was not there. And it's very uh if you think of it it's simple to understand why because
[25:02] we were just giving the pairs which some of them were useful some of them were like some of them were very useful some of them were less useful basically the contrast was in usefulness not in medical safety. That's why we have introduced a different approach
[25:20] on top of the first one which is based on the on the paper you can see it at the bottom ticking all the boxes. So what we did we have used all the medical judges and we have generated uh
[25:36] one very very big prompt. We have used Opus 4.8 at first and then refined it with GPT 5.5. So this huge prompt contains the check boxes or ticks for all possible aspects of medical safety and then we asked them
[25:52] all of them to award the scores from zero to two for each of them. Right and of course we have aggregated the results. So basically in that way now we have two metrics or two scores one for usefulness
[26:08] one for the medical safety. So we have combined them and then we have used that um to find the best answers and to build the preference pairs. So how that with the rejection sampling approach for building our final
[26:25] fine-tuning data data sets looked like. So basically at first we have picked only the valid answers according to medical safety. We have then also applied certain
[26:43] flow for the usefulness and formulated a pool of valid answers and then uh answers have became chosen. And rejected were just sourced from is a completely invalid questions or
[27:00] questions which had uh lower score. Both medical and um which had uh lower score according to both metrics.
[27:18] And we had two data sets, one for the supervised fine-tuning where we just picked the chosen was used as the block response. And for the preference tuning as you see on the side, we have chosen and rejected. Um
[27:45] All right. So, uh now let me tell you how we have actually trained all that. We have used Databricks AI runtime as the uh So, Databricks AI runtime is a new service which allows you to run
[28:01] distributed large-scale deep learning training. So, you can just request on demand any number of GPUs you want. So, we have used that. We have used 32 H100 GPUs um with PyTorch. And we have used FSDP for the distribution.
[28:17] And we have used LLM Foundry and Compose RL to run the actual fine-tuning. LLM Foundry is the uh library built by MosaicML. It's available in open source and Compose RL is a plugin which brings preference
[28:33] tuning. So, DPO DPO in particular which we have used. Okay. And now let us go into talk about the results. Okay. So, let's take a look what we actually managed to achieve. The most exciting part.
[28:49] Um First of all, we did burn a lot of tokens. Um So, our evaluation pipeline was consuming 30 billion tokens per week.
[29:07] But, we didn't came here to uh brass with their number of tokens that we consumed. I just thought it'll be like really good number to show. We would like to talk about the results of the fine-tuning. If you look from the left to the right, uh those are essentially iterations of our fine-tuning.
[29:22] Or evaluations of the model before and after the fine-tuning. You may have noticed that we used quite an old model, Lama version 3. That's because it was a long-running R&D. We've been trying to figure out how to approach it, what to do, how to build
[29:37] the pipeline. It all takes time. If you'll take a look at the slide, you will see that at our initial run, Lama model scored 34% in our medical benchmark.
[29:53] Usefulness was slightly higher, 42%, and the joint score was just 16%. Do you guys know why the joint score was just 16% while the other two were closer to 40?
[30:13] Anyone? No? Yes. Correct answer given by gentleman in the front row. So, not every answer that has been ranked as medically accurate was useful, and not every useful answer was medically accurate. So, 16% is quite low. That's not enough
[30:30] even to try to run AB experiment um with uh real users. If you look at their iterations or results of the fine-tuning, you will see that we actually managed to almost double both medical and usefulness
[30:47] bumping them from 34% to about 61% and 42% to 76% respectively. And now our joint usefulness and medical medical accuracy went from 16% to 47%.
[31:11] One may argue that 47% is not good enough and I would say that we 100% agree with that. So, what we will do next is we will reuse all the artifacts that we have created on this R&D journey that we did alongside with Databricks. So, big thank
[31:28] you for the Databricks team for doing that for us. We will take our pipeline for producing synthetic data. We will take our evaluation framework with 26 judges as it was at the moment, but this number is growing. Right now we
[31:44] have probably about 56 medical judges in general, yeah. Um we will take the fine-tuning pipeline and apply this to more modern, more advanced models. Our candidates will be Gemma, Gwen,
[32:00] Kimmy and several others. Because we have introduced a pipeline, now the whole iteration of every step is much much easier. And to Michael's point, instead of configuring those GPU clusters, we can now just click a button and they will
[32:16] magically appear. So, engineers actually focusing on what is important. On this note, I would like to thank you all for attending the session. We have about 6 minutes for questions, but if you won't be able to answer your question now, we will be happy to chat
[32:33] with you after the session.

Learn more about the Databricks Data and AI platform.

The information provided herein is for general informational purposes only and may not reflect the most current product capabilities or configurations.