All stories

Baseten argues attention beats cosine similarity for attribution

Baseten says attribution tools built on embedding similarity don't show what a model actually used, and argues that attention weights from open models give a more faithful record of where outputs came from.

Baseten has published a technical post arguing that many "attribution" features in generative AI products, including some clinical scribes, are poorly grounded. Their critique is summed up in the line "Cosine similarity is cosplay. Attention is attribution."

According to the company, vendors that rely on closed-model APIs only get text back. To attribute output to input, they have to split both into chunks, embed them with an encoder such as OpenAI's text-embedding-3-small, and rank pairs by cosine similarity. Baseten says this measures similarity, not influence. Each embedding encodes only its own chunk and none of the surrounding context.

To show the problem, Baseten used a synthetic task: turning a transcript into a clinical note that lists three injuries. Cosine similarity failed to link the third injury, "a bump with a small cut," to the forehead, because the transcript names the forehead in an earlier line. Baseten then averaged attention across all heads in layer 40 of Qwen2.5-32B and summed it over chunks. That method traced the injury back correctly. The company acknowledges the example is "simple and contrived."

The post also describes deployment trade-offs. One is whether to read attention from the generating model or from a cheaper proxy model, and which layers to use. Baseten says attribution must add less than 100ms of latency. It also says storing full attention tensors for a 70B, 80-layer model could take roughly 10–50GB of extra memory per request. For those reasons, Baseten runs attribution on a separate, dedicated model rather than during generation. The excerpt ends before giving further details of that setup or any quantitative evaluation.

Sources