PaperlayerView exact revision on arXiv

Source-linked AI summary

Permissive Information-Flow Analysis for Large Language Models

Shoaib Ahmed Siddiqui, Radhika Gaonkar, Boris Köpf, David Krueger, Andrew Paverd, Ahmed Salem, Shruti Tople, Lukas Wutschitz, Menglin Xia, Santiago Zanella-Béguelin

arXiv:2410.03055v3cs.LGcs.AI

TL;DR

LLMs that use diverse retrieved inputs face security and privacy risks, while propagating the most restrictive input label can be overly conservative. The paper labels outputs using influential inputs rather than every input, and in an LLM-agent setting the prompt-based propagator improved labels in more than 85% of cases without changing agent output.

  • Problem

    Conventional information-flow tracking propagates the most restrictive input label, which can be too conservative for LLMs using diverse sources, while poisoned inputs and confidential data pose system security and privacy risks.

    Source: §1 · pp. 1–2

  • Method

    The influence-based label propagator retains labels of inputs influential to the output and drops labels of unnecessary inputs.

    Source: §1 · p. 2

  • Results

    More than 85% of cases in the LLM-agent setting improved in label while the agent output remained the same.

    Source: §6 · p. 16

  • Takeaways & Limitations

    Influence-based propagation can produce more permissive labels without degrading LLM output quality.

    Source: Summary of contributions · p. 3

    Under adversarial conditions, unreliable distracting inputs can prevent label improvement and reduce downstream utility; the approach also increases LLM calls.

    Source: §6.2 · p. 17

Abstract

from arXiv · show

Large Language Models (LLMs) are rapidly becoming commodity components of larger software systems. This poses natural security and privacy problems: poisoned data retrieved from one component can change the model's behavior and compromise the entire system, including coercing the model to spread confidential data to untrusted components. One promising approach is to tackle this problem at the system level via dynamic information flow (aka taint) tracking. Unfortunately, this approach of propagating the most restrictive input label to the output is too conservative for applications where LLMs operate on inputs retrieved from diverse sources. In this paper, we propose a novel, more permissive approach to propagate information flow labels through LLM queries. The key idea behind our approach is to propagate only the labels of the samples that were influential in generating the model output and to eliminate the labels of unnecessary inputs. We implement and investigate the effectiveness of two variations of this approach, based on (i) prompt-based retrieval augmentation, and (ii) a $k$-nearest-neighbors language model. We compare these with a baseline that uses introspection to predict the output label. Our experimental results in an LLM agent setting show that the permissive label propagator improves over the baseline in more than 85% of the cases, which underscores the practicality of our approach.

1 Introduction

LLMs increasingly connect diverse data sources to downstream software, creating risks from poisoned inputs and confidential-data leaks. The paper proposes influence-based label propagation to reduce overly restrictive labels, reporting improved label precision and permissiveness while preserving output quality.

1 Introduction

  • LLM systems ingest data from diverse sources and pass outputs to other software, exposing them to poisoned-input and confidential-data leakage risks.Low-integrity inputs can change model behavior, while high-confidentiality inputs can reach untrusted downstream components.Source: §1 · p. 1
  • Conventional information-flow tracking propagates the most restrictive input label, which can make LLM outputs unnecessarily restrictive when contexts contain diverse sources.The paper identifies this problem as label creep.Source: §1 · p. 2
  • The paper’s influence-based label propagator assigns output labels using influential inputs and drops labels from inputs that were not influential.For a fixed tolerance λ, it selects the most permissive label among subcontexts whose utility is at most λ below that of the full context.Source: §1 · p. 2
  • The authors implement prompt-based retrieval augmentation and a kNN-LM as two realizations of the label propagator.The prompt-based system supplies retrieved documents in the prompt; the kNN-LM mixes model and retrieved-document output distributions.Source: §1 · p. 3

Summary of contributions

  • More than 85% of cases in the LLM-agent setting improved in label while the agent output remained the same.Output sameness was measured by reduced-context versus full-context output alignment.Source: Summary of contributions · p. 3, §1 · p. 3
  • The approach correctly identifies exact labels in at least 75% of cases and improves labels in at least 50% of cases on all datasets.The authors report that this mitigates label creep.Source: Summary of contributions · p. 3

2 Problem Setting and Goal

The paper models an LLM query as a prompt and a context of labeled documents, with labels ordered in a lattice. Its goal is to assign more permissive output labels using only necessary inputs, while recognizing that partial orders can yield multiple incomparable minimal labels.

2 Problem Setting and Goal

  • A context is a set of documents supplied with a prompt, and the LLM is modeled as a distribution over possible completions.The framework covers retrieved knowledge-base documents and data returned by plugins.Source: §2 · p. 3

2.1 Information-Flow Labels

  • Each document receives a label, and labels form a lattice whose order captures relative permissiveness; joins and meets are the least upper and greatest lower bounds.Labels can encode access-control information or source reliability.Source: §2.1 · p. 3
  • The product lattice combines integrity and time dimensions, with (LoInt, LastMonth) least reliable and least recent and (HiInt, Today) most reliable and most recent.Each dimension is totally ordered, while their product has a partial order.Source: §2.1 · p. 4, Figure 2 · p. 4
  • The threat model assumes input labels are given and correct, while an adversary may tamper with documents or tools at or below its lattice level.The adversary seeks illicit flows, such as poisoned low-integrity data affecting high-integrity output or confidential input affecting public output.Source: §2.1 · p. 4

2.2 Goal: Permissive Label Propagation

  • Naïve propagation labels an output with the join of every context document, even though the LLM may not need every document to generate it.An unnecessary untrusted document can make an output unusable by a downstream tool requiring trusted data.Source: §2.2 · p. 5
  • The goal is to propagate labels only from inputs necessary for the output, but partial-order labels may yield multiple incomparable minimal choices.The application may select among all minimal labels rather than relying on a single globally optimal label.Source: §2.2 · p. 5

3 Permissive Label Propagation

The method searches for more permissive labels by testing whether restricted contexts preserve output utility within a tolerance. It returns minimal candidate labels and wraps retrieval-augmented systems to produce a completion using the selected context.

3 Permissive Label Propagation

  • The approach searches for labels of inputs necessary for an output, avoiding a naïve search over every document subset that grows exponentially with context size.It considers labels that improve on the full context’s label and is practical and optimal under a monotonicity assumption.Source: §3 · p. 5

3.1 λ-similar Labels

  • A label L′ is λ-similar to the full-context label when restricting to its subcontext reduces utility by at most λ.For this paper, utility is negative perplexity; λ-similarity is neither symmetric nor transitive.Source: §3.1 · p. 5

3.2 Computing λ-similar Labels

  • The search tests labels represented by context documents rather than all document subsets, traversing a DAG to find λ-similar labels.The candidate label set is finite even if the full lattice is infinite.Source: §3.2 · p. 6

3.3 Correctness

  • When utility is monotonic, Algorithm 1 returns all minimal λ-similar labels; otherwise, it is a heuristic because utility need not be monotonic.It always terminates and returns a minimal set, whose labels are pairwise incomparable in the lattice.Source: §3.3 · pp. 6–7

3.4 A System for Label Propagation

  • The wrapper reruns the model using the selected restricted context and returns its completion with the selected label.Even if influence estimates are mistaken or adversarially misled, the returned completion depends only on documents at or below that label.Source: §3.4 · p. 7

3.5 Computational Cost

  • Prompt-based propagation can require exponentially many LLM calls in the worst case, while the considered kNN architecture needs one additional LLM call per query.For several fundamental lattices, the algorithm visits only a small subset of the lattice in typical queries.Source: §3.5 · pp. 7–8

4 Evaluation Setup

The evaluation tests whether influence-based label propagation can identify minimal, more permissive labels across synthetic, news, and tool-using-agent settings. It measures label-search accuracy, comparisons with introspection, and the quality of outputs regenerated from reduced contexts.

4 Evaluation Setup

  • The evaluation targets accurate identification and propagation of minimal labels, comparing the approach with an introspection baseline.The three datasets probe many inputs, long free-form language, and tool-using agent conversations.Source: §4 · p. 8, §4.5 · p. 10

4.1 Research Questions

  • The study asks how accurately the propagator identifies minimal labels, how much it improves over naïve propagation, and how aligned regenerated outputs are with full-context outputs.The first two questions evaluate Algorithm 1; the third evaluates the end-to-end system.Source: §4.1 · p. 9

4.2 Evaluation Metrics

  • Exact match, precision, and recall compare predicted label sets with ground truth, while label improvement and missed labels measure outcomes in totally ordered lattices.ROUGE-L measures output alignment; the study compares reduced-context and full-context outputs, including when label improvement is possible or not.Source: §4.2 · p. 10

4.3 Baseline

  • The introspection baseline asks the LLM which documents influenced its output, and its results are treated as a lower bound because prompting choices affect performance.The baseline is inspired by work on models generating citations for their own outputs (Taylor et al., 2022).Source: §4.3 · p. 10

4.4 Models

  • Experiments use instruction-tuned Llama-2 7B and 70B models, with the kNN-LM implementation using the penultimate-layer representation of the last token for retrieval.The authors state the approach also applies to other open-source models and API-accessible proprietary models, though client-side optimizations may be harder.Source: §4.4 · p. 10

4.5 Datasets

  • For evaluation, the target response supplies the completion unless noted otherwise, because the dataset’s ground-truth label is correct only for that response.The authors compare against model-generated outputs to quantify the implications of this choice.Source: §4.5 · p. 11
  • The synthetic dataset distributes personal identifiers, SSNs, and birth dates across documents, with questions answerable from different document subsets; its 14-document contexts contain all necessary documents plus possible irrelevant ones.It has 128 documents and 64 question-answer pairs, and its partial-order lattice has 2^14 possible labels.Source: §4.5 · p. 11
  • The news dataset pairs similar high- and low-integrity articles and contains 3,465 QA pairs; the agent dataset has 40 tool-using conversations, half permitting label improvement and half requiring low-integrity sources.The news dataset has 240 article pairs; agent conversations use web, email, or calendar search, with two in-context examples specifying output format.Source: §4.5 · pp. 11–12

5 Results

The influence-based propagator performs well across synthetic, news, and agent evaluations, with results depending on the dataset, retriever, and model output. Regenerating responses from reduced contexts generally preserves alignment when labels improve, while model-generated errors can affect label-search results.

5.1 Synthetic key-value dataset

  • Over 85% exact match accuracy and over 90% precision and recall were achieved by prompt-based label search, significantly outperforming kNN-LM and introspection on the synthetic dataset.The search identifies influential subcontexts among 16k possibilities.Source: §5.1 · p. 13
  • Very high or very low perplexity tolerance λ significantly hampers exact-match performance; optimal settings are consistent across model scales, while larger models do better at lower λ.The larger models’ advantage at lower λ is attributed to stronger language modeling.Source: §5.1 · pp. 13–14
  • 81% to 86% exact match accuracy was achieved after adding Shapley-value-based pruning, which can marginally improve performance, particularly for large partial-order lattices.The results in Table 1 include this heuristic.Source: §5.1 · p. 14

5.2 News article dataset

  • With a perfect retriever, the propagator assigns the correct more permissive label in approximately 57% of improvable news cases and misses a ground-truth label in approximately 23%; a realistic retriever lowers exact-match accuracy by about 10%.The realistic retriever sometimes fails to retrieve relevant information, causing a mismatch with the ground-truth label.Source: §5.2 · p. 14
  • The optimal perplexity tolerance is significantly higher for news responses than for synthetic questions, and larger models tolerate lower λ better.The authors attribute the higher news tolerance to greater response difficulty.Source: §5.2 · p. 14

5.3 LLM agent dataset

  • In the agent dataset, label improvement occurred in at least 17 of 20 eligible cases, and regenerated outputs had ROUGE-L F-scores of at least 0.85 with negligible alignment change when improvement was possible.Alignment dropped significantly when label improvement was not possible; the experiment used λ = 0.2.Source: §5.3 · p. 15
  • Using ground-truth targets instead of model-generated outputs yields perfect precision and recall; label-propagation errors with generated outputs are associated with the model’s suboptimal responses.Introspection also performs strongly on the agent dataset’s simpler lattice and improves with model scale.Source: §5.3 · p. 16

6 Analysis and Discussion

The prompt-based label propagator finds minimal labels across large and ordered lattices and improves labels in a tool-calling agent setting. The discussion connects these results to reuse and endorsement use cases, while noting added LLM-call costs and adversarial risks.

6 Analysis and Discussion

  • 86% of cases yield the exact set of minimal labels on a lattice with 16k possible labels.This is the prompt-based propagator’s result for RQ1.Source: §6 · p. 16
  • 85% of LLM-agent cases and 56% of news-article cases show label improvement on a smaller, totally ordered lattice.The ordered lattice allows direct label comparison and quantification of improvement.Source: §6 · p. 16
  • More than 85% of tool-calling agent cases show improved labels while output alignment with the full context remains the same.Alignment is measured by the difference between reduced-context and full-context output alignment.Source: §6 · p. 16
  • The introspection baseline improves with a 70B model over a 7B model, while prompt-based performance is similar at both sizes.Introspection remains a strong baseline on simple lattices; similar prompt-based performance suggests the approach could use smaller models.Source: §6 · p. 16

6.1 Use-cases

  • More permissive labels can broaden answer reuse in semantic caches and facilitate forwarding emails beyond their initial recipients.The approach is most beneficial when the generated content’s eventual use is not yet known; if the target label is known up front, retrieval can instead be restricted to documents at or below that label.Source: §6.1 · p. 17
  • A trusted authority confirming an untrusted source can lead the propagator to assign a trusted output label, enabling endorsement by repeating that information.The sources must be similarly influential for the propagator to assign the trusted label.Source: §6.1 · p. 17

6.2 Limitations

  • Label improvement on two realistic datasets costs additional LLM calls, and adversarial distracting information can prevent improvement and reduce downstream utility.In the presence of an adversary, the propagator assigns the lowest integrity label by design to avoid harmful side effects such as cross-prompt injection.Source: §6.2 · p. 17

6.3 Extensions

  • The approach can be adapted to target a specified label or dynamically trade utility for a less conservative label; KV-cache reuse and alternative influence estimates are proposed extensions.Simple attention scores are difficult for influence estimation in light of attention sinks and lost-in-the-middle effects.Source: §6.3 · pp. 17–18

7 Related Work

The related work situates the approach within information-flow analysis, secure multi-execution, and emerging LLM security systems. It distinguishes this work by dynamically propagating labels through LLM calls under a utility constraint.

7 Related Work

  • Unlike conventional program operators with hard-coded label rules, Transformer outputs are tainted by all inputs, motivating label propagation constrained by model loss.The paper describes Transformer behavior as implicitly defined by loss minimization.Source: §7 · p. 18
  • This approach is fundamentally dynamic because it analyzes one specific context during execution rather than reasoning over all possible executions.The distinction between dynamic and static information-flow control follows Sabelfeld & Myers (2003); Buiras et al. (2014); Mardziel et al. (2011).Source: §7 · p. 18
  • Like secure multi-execution, the approach calls the LLM once per security level, but selects the most permissive label whose output remains correct within utility tolerance λ.Here, correctness means the drop in utility is bounded by λ.Source: §7 · p. 18
  • Compared with other information-flow systems for LLM applications, this work dynamically and soundly propagates labels through LLM calls to increase system-level permissiveness.The cited systems address datastore metadata, untrusted information, execution-trace monitoring, or tool-call security.Source: §7 · p. 18

8 Conclusion

The paper concludes that it propagates labels for documents actually used to generate an output, yielding more permissive labels than an introspection baseline while retaining a hard safety guarantee.

8 Conclusion

  • The approach propagates labels only for retrieved documents actually used to generate the model’s output.The conclusion presents this as a permissive alternative for information-flow label propagation in RAG systems.Source: §8 · p. 18
  • Unlike commonly used introspection-based methods, the approach can satisfy a hard safety guarantee while inferring more permissive labels than the baseline.The paper reports the approach as practical in terms of performance.Source: §8 · p. 19

A Improving performance on fine-grained lattices

A Shapley-value threshold heuristic filters label combinations when λ-similar subcontexts include irrelevant documents. Its measured benefit is concentrated in the complex synthetic key-value lattice.

A Improving performance on fine-grained lattices

  • λ-similar subcontexts can include irrelevant documents when those documents share similarities and contribute useful information such as output format.The appendix identifies this issue for the synthetic key-value dataset.Source: A Improving performance on… · pp. 24–25
  • The heuristic computes a Shapley value for each lattice label using average perplexity drop when that label is included, then filters combinations below a threshold.This adds a Shapley-value threshold hyperparameter.Source: A Improving performance on… · p. 25
  • 6% absolute improvement in exact-match accuracy is observed on the synthetic key-value dataset, while the threshold has almost no impact on the simpler news-article lattice.The results compare a two-dimensional grid search over perplexity tolerance λ and Shapley threshold.Source: A Improving performance on… · p. 25

Related papers

Loading 2410.03055v3…