AI2 Olmo traces everything your LLM does back to its original source

Understanding precisely how the output of a Large Language Model (LLM) Artificial Intelligence (AI) matches with its training data has long been a mystery and a challenge for enterprise IT.

Now, a new open source effort launched this week by the Allen Institute for AI (Ai2) aims to help solve that challenge by tracing LLM output to training inputs. The OLMoTrace tool allows users to trace language model outputs directly back to the original training data, addressing one of the most significant barriers to enterprise AI adoption: the lack of explainability and transparency in how AI systems make decisions.

OLMo is an acronym for Open Language Model, which is also the name of Ai2’s family of open-source LLMs. On the company’s Ai2 Playground site, users can try out OLMoTrace with the recently released OLMo 2 32B model. The open-source code is also available on GitHub and is freely available for anyone to use.

Unlike existing approaches focusing on confidence scores or retrieval-augmented generation, OLMoTrace offers a direct window into the relationship between model outputs and the multi-billion-token training datasets that shaped them.

“Our goal is to help users understand why language models generate the responses they do,” said Jiacheng Liu, researcher at Ai2.

LLMs with web search functionality, like Perplexity or ChatGPT Search, can provide source citations. However, those citations are fundamentally different from what OLMoTrace does and sometimes those citations aren’t accurate.

Liu explained that Perplexity and ChatGPT Search use Retrieval Augmented Generation (RAG). With RAG, the purpose is to improve the quality of model generation by providing more sources than what the model was trained on. OLMoTrace is different because it traces the output from the model itself without any RAG or external document sources.

The technology identifies long, unique text sequences in model outputs and matches them with specific documents from the training corpus. When a match is found, OLMoTrace highlights the relevant text and provides links to the original source material, allowing users to see exactly where and how the model learned the information it’s using.

By design, LLMs generate outputs based on model weights that help to provide a confidence score. The basic idea is that the higher the confidence score, the more accurate the output.

In Liu’s view, confidence scores are fundamentally flawed.

“Models can be overconfident of the stuff they generate and if you ask them to generate a score, it’s usually inflated,” Liu said. “That’s what academics call a calibration error - the confidence that models output does not always reflect how accurate their responses really are.”

Instead of another potentially misleading score, OLMoTrace provides direct evidence of the model’s learning source, enabling users to make their own informed judgments.

“What OLMoTrace does is show you the matches between model outputs and the training documents,” Liu explained. “Through the interface, you can directly see where the matching points are and how the model outputs coincide with the training documents.”

Ai2 is not alone in the quest to better understand how LLMs generate output. Anthropic recently released its own research into the issue. That research focused on model internal operations, rather than understanding data.

“We are taking a different approach from them,” Liu said. “We are directly tracing into the model behavior, into their training data, as opposed to tracing things into the model neurons, internal circuits, that kind of thing.”

This approach makes OLMoTrace more immediately useful for enterprise applications, as it doesn’t require deep expertise in neural network architecture to interpret the results.

For enterprises deploying AI in regulated industries like healthcare, finance, or legal services, OLMoTrace offers significant advantages over existing black-box systems.

“We think OLMoTrace will help enterprise and business users to better understand what is used in the training of models so that they can be more confident when they want to build on top of them,” Liu said. “This can help increase the transparency and trust between them of their models, and also for customers of their model behaviors.”

The technology enables several critical capabilities for enterprise AI teams:

• Fact-checking model outputs against original sources
• Understanding the origins of hallucinations
• Improving model debugging by identifying problematic patterns
• Enhancing regulatory compliance through data traceability
• Building trust with stakeholders through increased transparency

The Ai2 team has already used OLMoTrace to identify and correct their models’ issues.

“We are already using it to improve our training data,” Liu reveals. “When we built OLMo 2 and we started our training, through OLMoTrace, we found out that actually some of the post-training data was not good.”

For enterprises looking to lead the way in AI adoption, OLMoTrace represents a significant step toward more accountable enterprise AI systems. The technology is available under an Apache 2.0 open-source license, which means that any organization with access to its model’s training data can implement similar tracing capabilities.

“OLMoTrace can work on any model, as long as you have the training data of the model,” Liu notes. “For fully open models where everyone has access to the model’s training data, anyone can set up OLMoTrace for that model and for proprietary models, maybe some providers don’t want to release their data, they can also do this OLMoTrace internally.”

As AI governance frameworks continue to evolve globally, tools like OLMoTrace that enable verification and auditability will likely become essential components of enterprise AI stacks, particularly in regulated industries where algorithmic transparency is increasingly mandated.

For technical decision-makers weighing the benefits and risks of AI adoption, OLMoTrace offers a practical path to implementing more trustworthy and explainable AI systems without sacrificing the power of large language models.