ProToken: Token-Level Attribution for Federated Large Language Models

TL;DR

ProToken addresses the challenge of attributing token-level outputs in federated LLMs to contributing clients under privacy constraints. By leveraging the linearity of FL aggregation, focusing on the last Transformer blocks, and applying gradient-based relevance weighting, ProToken delivers tractable, per-token provenance signals that aggregate into sequence-level attributions. The approach achieves an average attribution accuracy of $98.62\%$ across 16 configurations and remains robust when scaling up to 55 clients, maintaining over $92\%$ accuracy and clear contributor separation. This work advances trustworthy federated LLMs by providing a practical, ground-truth-grounded provenance framework for debugging, accountability, and fair reward distribution.

Abstract

Federated Learning (FL) enables collaborative training of Large Language Models (LLMs) across distributed data sources while preserving privacy. However, when federated LLMs are deployed in critical applications, it remains unclear which client(s) contributed to specific generated responses, hindering debugging, malicious client identification, fair reward allocation, and trust verification. We present ProToken, a novel Provenance methodology for Token-level attribution in federated LLMs that addresses client attribution during autoregressive text generation while maintaining FL privacy constraints. ProToken leverages two key insights to enable provenance at each token: (1) transformer architectures concentrate task-specific signals in later blocks, enabling strategic layer selection for computational tractability, and (2) gradient-based relevance weighting filters out irrelevant neural activations, focusing attribution on neurons that directly influence token generation. We evaluate ProToken across 16 configurations spanning four LLM architectures (Gemma, Llama, Qwen, SmolLM) and four domains (medical, financial, mathematical, coding). ProToken achieves 98% average attribution accuracy in correctly localizing responsible client(s), and maintains high accuracy when the number of clients are scaled, validating its practical viability for real-world deployment settings.

Figures (7)

• Figure 1: Motivating example showing how ProToken identifies the client responsible for anomalous output.
• Figure 2: ProToken Provenance Attribution Performance.Blue circles: ProToken attribution accuracy for identifying contributing clients. Orange squares: Model accuracy on benign responses. Red triangles (dashed): Model accuracy on triggered responses (evaluation ground truth). ProToken achieves on average attribution accuracy of 98.62%.
• Figure 3: ProToken Client Contribution Probability Distributions.Red boxes: Clients 0-1 (contributors) receive high probabilities. Blue boxes: Clients 2-5 (non-contributors) receive near-zero probabilities. The complete separation between red and blue distributions shows that ProToken provides clear, attribution signals, enabling confident provenance decisions in production.
• Figure 4: Average per-layer (i.e., individual layer) attribution accuracy of ProToken across 16 configurations (4 models $\times$ 4 domains). Bars show average attribution accuracy when averaging across all transformer block layers per configuration. Gradient weighting provides substantial improvements across all settings, demonstrating its effectiveness in filtering irrelevant neurons.
• Figure 5: For each model, we vary the number of monitored layers (x-axis) and measure ProToken's average provenance computation time (left y-axis, blue) and attribution accuracy (right y-axis).