A transformer that can show which words shaped its answer

A language model can give a confident answer without making clear which parts of its input mattered. Attention weights are often shown as a window into a transformer’s reasoning, but a high attention weight alone does not establish that a word caused or determined the output. A new preprint by Teerawat Tongloy, Jiraporn Pongsopa and Rangsan Jomtarak asks whether a model can instead be designed so that its answer has a calculable, word-by-word accounting.

The authors built a transformer for natural-language inference: deciding whether a hypothesis follows from a premise, contradicts it or could be either true or false. Its architecture replaces standard operations such as softmax, LayerNorm and ReLU with quadratic, measurement-like calculations. Attention is based on squared overlaps between word representations, normalized across each row. The model’s final class scores are probabilities derived from squared amplitudes, following the mathematical form of the Born rule in quantum physics.

This design has a specific benefit for attribution. If the attention normalizers and other denominators are held fixed at their values for the original input, the class probability becomes a polynomial in the word amplitudes. The team uses integrated gradients—a method that assigns contributions by tracking how a model’s output changes as its input is scaled from a baseline—to calculate each word’s evidence. Because the function is polynomial, a finite numerical integration rule can compute the attribution without integration error. The contributions sum to the predicted class probability above a defined “no-information” share.

That precision has an important qualification: the denominators are held fixed during the attribution calculation. The result is exact for that defined calculation, not a uniquely true account of why the model answered as it did. The authors also note that removing a word changes the model’s normalizers, so deletion tests provide a separate check rather than a direct restatement of the attribution.

On the e-SNLI benchmark, the model reached 77.45% test accuracy, averaged across three training seeds. A similarly sized softmax transformer scored 76.62%; the reported seed-to-seed spreads overlap. An unconstrained polynomial model was somewhat more accurate, at 78.98%. The comparison therefore suggests that the measurement-inspired design can approach the standard transformer’s performance in this experiment, not that it is more accurate overall.

The researchers tested whether the word rankings were useful in two ways. First, they removed the word ranked most important and checked whether the predicted answer changed. This happened in 44% of test pairs for the new model’s evidence ranking, compared with 31% for the standard transformer’s integrated-gradient ranking and 31% for the new model’s attention ranking. Some single deletion changed the answer in 53% of pairs, so the top-word result fell short of a perfect test. Second, the authors compared top-ranked words with human highlights in e-SNLI. The new model’s evidence agreed with highlights more than the standard model’s evidence did, but attention rankings could also align well with highlights. Human annotations are useful reference points, not definitive proof of a model’s internal reasoning.

The paper also explores physical implementations. The model’s weights were compiled to an optical setup in simulation, not tested on a photonic chip. On IBM’s ibm_kingston quantum processor, the team measured attention overlaps for two of the shortest test examples, while performing the rest of the computation classically. The device reproduced the model’s answers for those examples, including one after deleting a decisive word. Evidence calculated using measured attention row sums closely tracked the exact evidence for one example and less closely for the other, where calibration left a residual. This is a small demonstration, not evidence of a quantum speed advantage; the authors explicitly make no such claim.

The scope remains limited: one benchmark, models much smaller and less accurate than state-of-the-art language systems, and only two short examples on quantum hardware. The optical results are simulated, and a separate fully coherent quantum circuit was also simulated rather than run on hardware. The work is an arXiv preprint, so its claims should be treated as provisional pending independent scrutiny and replication. Its central contribution is a way to make a particular kind of word attribution mathematically exact—while showing why exactness, usefulness and a complete explanation are different things.