Using roBERTa models + LLMs to improve NER results in healthcare data

The healthcare industry treats data privacy very seriously, and rightfully so. Besides the obvious ethical reasons for doing so, there is, of course, the Health Insurance Portability and Accountability Act of 1996 better known as HIPAA. 

HIPAA dictates how healthcare data can be used and what constitutes a successful de-identification of the data so it can be used more freely. Common reasons for de-identifying healthcare data are:

• Development and testing use cases, which are improved by using more realistic data.
• Data commercialization: selling data to data vendors or through marketplaces or directly to LLM labs.
• Model training. fine-tuning healthcare models with additional data to improve them.

For most use cases, companies typically go for an ‘expert determination’ of their data which engages a statistician to review the de-identified data to ensure patient privacy is being maintained (read here to learn more).

De-identifying unstructured healthcare data

To properly de-identify unstructured data like lab reports, discharge notes, medical summaries, doctor’s notes, etc. you can either manually review each document and manually remove/replace PHI, or you can use various machine learning models to do the same but more quickly and over far more data.

With Tonic Textual, we primarily use fine-tuned roBERTa models to identify sensitive health information. Our roBERTa models are fine-tuned on a variety of real or synthesized datasets commonly found in the healthcare space, as well as other industries. We regularly assist customers in obtaining an Expert Determination on their unstructured data. And thanks to Textual’s automations, we’ve seen customers get an ED using our tooling in as little as 5 days from start to finish, on a combination of medical notes and PDF lab reports, which is typically quite a challenge.

Using LLMs in de-identification

Despite not being ideal for this specific task in many ways, LLMs can be useful. We’ve found that feeding all text through an LLM does not work for use cases where latency is paramount or there is simply a very large volume of words to de-identify (think on the order of billions of words). The issue comes down to the tok/s throughput of frontier LLM models. One way around this is to selectively feed data to the LLM to use it as another signal to improve the results from the faster, cheaper roBERTa models. 

We specifically focused on the task of lowering the false positive rate for entities tagged by our roBERTa model. In this way, we send much less data to the LLM (just the entities, not the original text) and are able to keep LLM cost and latency both low. 

As a motivating example, consider a piece of text such as:

A roBERTa model will often flag Parkinsons as a last name, especially if the model is not trained on disease names or medical conditions. However, feeding this result into an LLM can quickly discard such a result as its clearly a mislabel. 

By using LLMs to reduce false positives we saw a net increase in f1 scores across all entities types with the exception of our DATE_TIME entity which saw a negligible decrease in f1 of 0.0022. As expected, precision was most affected. We noted that 7 entity types saw no improvements (the roBERTa model was already operating very well), while the remaining 15 entity types tracked saw improvements . For those entities which did improve we saw a P50 improvement of ~1% and a P75 improvement of 2%. Of special note, precision for last name and state both saw improvements of ~4%.

LLM-assisted NER in Tonic Textual

To take advantage of this feature in Textual, you can simply enable it in our UI under your ‘Dataset Settings’. Or alternatively, you can take advantage of our Python SDK (its open source) and simply pass it as an argument to the redact function. Below we show two examples on the same input with the LLM classification disabled and then enabled:

With it disabled:

response = ner.redact(‘My name is Paris. I live in Paris. I am an artist that works a lot with Plaster of Paris.’, enable_llm_classification = False)

{
    "original_text": "My name is Paris. I live in Paris. I am an artist that works a lot with Plaster of Paris.",
    "redacted_text": "My name is [NAME_GIVEN_R6QJj3]. I live in [LOCATION_CITY_R6QJj3]. I am an [OCCUPATION_XdRCfjF] that works a lot with [ORGANIZATION_A5bU9oouNlHEdg8iQM].",
    "usage": 20,
    "de_identify_results": [
        {
            "start": 11,
            "end": 16,
            "new_start": 11,
            "new_end": 30,
            "label": "NAME_GIVEN",
            "text": "Paris",
            "score": 0.9138046503067017,
            "language": "xx",
            "new_text": "[NAME_GIVEN_R6QJj3]"
        },
        {
            "start": 28,
            "end": 33,
            "new_start": 42,
            "new_end": 64,
            "label": "LOCATION_CITY",
            "text": "Paris",
            "score": 0.9367831349372864,
            "language": "xx",
            "new_text": "[LOCATION_CITY_R6QJj3]"
        },
        {
            "start": 43,
            "end": 49,
            "new_start": 74,
            "new_end": 94,
            "label": "OCCUPATION",
            "text": "artist",
            "score": 0.85813307762146,
            "language": "xx",
            "new_text": "[OCCUPATION_XdRCfjF]"
        },
        {
            "start": 72,
            "end": 88,
            "new_start": 117,
            "new_end": 150,
            "label": "ORGANIZATION",
            "text": "Plaster of Paris",
            "score": 0.888020396232605,
            "language": "xx",
            "new_text": "[ORGANIZATION_A5bU9oouNlHEdg8iQM]"
        }
    ]
}

We can see that ‘Plaster of Paris’ was incorrectly flagged as an organization.  On a positive note though the roBERTa model alone was able to tease about the different meanings of the word ‘Paris’.

With LLM classification enabled we remove the false positive:

{
    "original_text": "My name is Paris. I live in Paris. I am an artist that works a lot with Plaster of Paris.",
    "redacted_text": "My name is [NAME_GIVEN_R6QJj3]. I live in [LOCATION_CITY_R6QJj3]. I am an [OCCUPATION_XdRCfjF] that works a lot with Plaster of Paris.",
    "usage": 20,
    "de_identify_results": [
        {
            "start": 11,
            "end": 16,
            "new_start": 11,
            "new_end": 30,
            "label": "NAME_GIVEN",
            "text": "Paris",
            "score": 0.9138046503067017,
            "language": "xx",
            "new_text": "[NAME_GIVEN_R6QJj3]"
        },
        {
            "start": 28,
            "end": 33,
            "new_start": 42,
            "new_end": 64,
            "label": "LOCATION_CITY",
            "text": "Paris",
            "score": 0.9367831349372864,
            "language": "xx",
            "new_text": "[LOCATION_CITY_R6QJj3]"
        },
        {
            "start": 43,
            "end": 49,
            "new_start": 74,
            "new_end": 94,
            "label": "OCCUPATION",
            "text": "artist",
            "score": 0.85813307762146,
            "language": "xx",
            "new_text": "[OCCUPATION_XdRCfjF]"
        }
    ]
}

With this new support for LLM-assisted NER redaction we are seeing our healthcare customers de-identify their data faster and obtain expert determination attestations more quickly, saving them precious time and money. If you are interested in trying it out yourself, connect with our team today.