Technical Summary: ProToken – Token-Level Attribution for Federated Large Language Models

1. Problem Statement

Federated Learning (FL) enables the collaborative training of Large Language Models (LLMs) across distributed, private data sources without centralizing raw data. While this preserves privacy, it creates a "black box" scenario regarding model provenance. When a global federated LLM generates a specific response, it is unclear which client(s) contributed to that output. This lack of attribution hinders critical system applications, including:

• Debugging: Identifying which client introduced problematic behaviors or poisoned data.
• Trust Verification: Validating that the global model reflects the intended contributions.
• Fair Reward Allocation: Allocating incentives based on actual contribution.
• Malicious Client Identification: Detecting participants injecting backdoors or adversarial patterns.

Existing interpretability techniques for centralized models (e.g., gradient integration, perturbation analysis) are inapplicable because they assume access to raw training data and single-model architectures. Furthermore, standard FL debugging methods are designed for classification tasks with discrete outputs, failing to address the variable-length, autoregressive nature of LLM generation. A core challenge is distinguishing whether a generated token stems from a specific client's federated training data or the model's pre-existing pre-trained knowledge.

2. Methodology: ProToken

ProToken is a provenance methodology designed to attribute token-level generation in federated LLMs to specific clients while strictly adhering to FL privacy constraints (i.e., operating only on model updates, activations, and gradients, not raw data).

Core Insights

ProToken leverages three key structural insights to make attribution tractable and accurate:

1. Linear Aggregation: FL aggregation strategies (e.g., FedAvg, FedProx) are linear at the parameter level. This allows the global model's forward computation to be mathematically decomposed into a weighted sum of per-client layer-wise computations.
2. Layer Specialization: Transformer architectures concentrate task-specific signals in later blocks (specifically self-attention output projections and final feed-forward layers). Early layers capture generic linguistic features, while later layers encode the specific knowledge influencing the final output.
3. Gradient-Based Relevance: Not all neurons are relevant for generating a specific token. Gradient-based weighting filters out irrelevant neural activations, focusing attribution only on neurons that directly influence the token prediction.

Technical Workflow

The attribution process operates during the autoregressive generation of a response:

1. Layer Selection: Instead of tracking all parameters (which is computationally prohibitive), ProToken restricts analysis to a strategic subset of layers: the Output Projection layer of the self-attention mechanism and the final layer of the MLP within the last $N$ transformer blocks.
2. Per-Token Attribution: For each generated token $x_j$, the system computes a provenance score for each client $i$.
   • The global model processes the input to generate hidden states ($h^\ell_G$) and gradients ($g^\ell_{x_j}$) for the selected layers.
   • For each client $i$, the system simulates what the layer output would be ($h^\ell_i$) if the client's specific weights were used, keeping the input from the global forward pass.
   • The client's contribution score for the token is calculated as the inner product of the client's hypothetical activation and the global gradient:
     $$P^\ell_{i, x_j} = \langle h^\ell_i, g^\ell_{x_j} \rangle$$
   • This score is summed across selected layers to get a per-token score, and then accumulated across the entire generated sequence to produce a final sequence-level attribution score.
3. Normalization: Scores are normalized via softmax to produce a probability distribution over clients, identifying the most likely source.

3. Evaluation Framework

To address the lack of ground truth in LLM provenance (distinguishing between pre-trained knowledge and federated contributions), the authors constructed a verifiable evaluation framework using backdoor injection:

• Ground Truth Generation: Distinct trigger-response pairs (e.g., a specific trigger phrase leading to a fixed "sentinel" refusal response) were injected into the local training data of designated "malicious" clients.
• Evaluation Metric: Attribution accuracy is measured by the system's ability to correctly identify the trigger-bearing clients when the sentinel response is generated. This provides an unambiguous ground truth that the model behavior originated from specific client contributions.

4. Experimental Results

The authors evaluated ProToken across 16 configurations involving four LLM architectures (Gemma, SmolLM, Llama, Qwen) and four domains (medical, financial, mathematical, coding).

• Attribution Accuracy: ProToken achieved an average attribution accuracy of 98.62% across all configurations. It successfully identified responsible clients with high confidence, showing clear binary separation between contributing and non-contributing clients.
• Scalability: In a scalability test increasing the client count from 6 to 55 (a 9.2× increase), ProToken maintained high accuracy (>92%), demonstrating robustness in large-scale federated deployments.
• Gradient Weighting Impact: An ablation study showed that gradient-based relevance weighting is critical. Disabling gradient weighting reduced average accuracy from 66.34% to 35.71%, confirming that filtering irrelevant neurons is essential for reliable attribution.
• Computational Tractability: By monitoring only the last subset of transformer blocks, ProToken reduced computational overhead significantly compared to tracking all parameters. For a 100-token response, the method remained efficient (e.g., ~1.1s to 1.8s depending on layer depth), making real-time attribution feasible.

5. Significance and Claims

The paper claims ProToken as the first token-level provenance attribution method specifically for federated LLMs. Its significance lies in:

• Solving the Attribution Gap: It addresses the fundamental challenge of tracing outputs in federated generative models where data is never centralized.
• Privacy Preservation: It operates entirely on model updates and gradients, maintaining the privacy guarantees of FL.
• Practical Viability: The method is computationally tractable and scalable, validated across diverse architectures and domains.
• Enabling Trust: It provides the necessary tools for debugging, identifying malicious participants, and ensuring fair reward allocation in collaborative LLM training, thereby laying a foundation for interpretable and accountable federated LLM systems.

The authors emphasize that while backdoor injection was used to create ground truth for evaluation, the method itself is designed for general attribution and is not an attack or defense strategy. The results validate ProToken as a foundational step toward trustworthy federated LLM deployment.