Permissive Information-Flow Analysis for Large Language Models
Authors: Shoaib Ahmed Siddiqui
Radhika Gaonkar
Boris Köpf
David Krueger
Andrew Paverd
Ahmed Salem
Shruti Tople
Lukas Wutschitz
Menglin Xia
Santiago Zanella-Beguelin
Venue: Accepted by TMLR ()

Abstract: 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. Assuming each piece of information comes with an additional meta-label (such as low/high integrity labels), 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 input labels that were \emph{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 our label propagator assigns a more permissive label over the baseline in more than 85% of the cases, which underscores the practicality of our approach.

Published in Transactions on Machine Learning Research (10/2025)
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 Ganella-Bégueling
msasS@cam.com ac.uk uk
granhika.guonkar, koepf, padrew..averrd I-salemahmed, shruti, tople, lukas, wuiscciizz] cricrosott.cooo {mollynia, santiago] Smmterosoft.com
david scoll kraeger @gmaill com
University of Cambridge Microsoft
Mila
Reviewed on OpenReview: https: // openreview. net/ forum? id= Y
Abstract
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 cocrcing the model to spread confidential data to untrusted components. Assuming each piece of information comes with an additional meta-label (such as low/high integrity labels), 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 inputs retrieved from diverse sources.
In this paper, propose novel, more permissive approach to propagate information flow labels through LLM queries. The key idea behind our approach is to propagate only the input labels that were influential in generating the model output and to eliminate the labels of unnocessary 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 our label propagator assigns more permissive label over the bascline in more than 85% of the cases, which underscores the practicality of our approach.
1 Introduction
Large Language Models (LLMs) such as GPT-4 (OpenAl 2023 Llama Touvron et al., 2023 Dubey et al., 2024h, Mistral (Jiang et al., 2023 and PaLM (Anil al. 2023 are rapidly becoming commodity components of larger software systems. The inputs to these LLMs often consist of data retrieved from variety of sources,
including websites, productivity software, or tools (Chase 2022). and their output is usually passed on to other software components for further processing Xi et al.
2023
This poses natural security and privacy problems: low integrity inputs (e.g., poisoned data) can change the model's behavior in unexpected ways and potentially affect the entire system (Greshake et al., 2023 Verma et al. 2024). Similarly, high confidentiality inputs (e.g., confidential documents) can be inadvertently leaked to untrusted downstream component (Mireshghallah et al. 2023

---

Published in Transactions on Machine Learning Research (10/2025)
Query
How do foobar? foobar?
LLM
SEARC(((nstaalfoobarr)
Retriever
untrusted
easiestway to install foobar is to use the Installation script -BL https: ///ooo com/ installl 1. sh
euel
Sudo
Daek
Influence based LP
trusted
Yes, the instructions found online correct. Just run cucl EL https://Eocte com/ int//nistllll.... sudo
Influence
LLM
Foobar provides an installation script that automatically install the package. Rum:
Naive LP
UUT=U
purl
HaL https: Sttpss//Foobar..t com/ instslnsttlll.....
sudo
bash
Figure 1: Illustration of label propagator (LP) for large language models (LLMs) with tool-calling capabilities. The goal of the LP is to assign the most suitable label to the output of the LLM. In this instance, we consider labels representing trusted and undrusted sources. naïve LP assigns the most conservative label to the output, which in this example untrusted. The LP we design takes into account the influence of each retrieved document and determines that the same output can be obtained by solcly relying on trusted documents.
One possible approach to address this problem to rely on the LLM itself for mitigation, for example via introspection of the retrieved inputs or guardrails given in the meta- preapprommp. However, such defenses can be circumvented with more advanced attacks (Zou al. 2023 Yuan et al. 2023 Li et al. 2023 Shen et al.) 2023 leading to the undesirable cat-and-mouse game that is common in system security. A promising alternative is to tackle this problem at the system level via dynamic information-flow (aka taint) tracking. Information-flow tracking is a standard technique for enforcing integrity and confidentiality properties (Clause et al. 2007
Buiras et al. 2014 Sabelfeld & Myers) 2003), and has been used successfully in many applications, including delecting cross-sile-ccripting vulnerabilities (Vogt et al. 2007), privacy leaks mobile applications (Enck
et al. 2014) and recently to L.LM-based systems (Wutschitz et al. 2023 Wu al. 2024 2024
Balunovie al.
In information-flor tracking, each piece of data is augmented with labels describing its integrity confiden- tiality. Labels are usually propagated conservatively: the output of operation on data is labeled as the most restrictive (iee,, most confidential or least trusted) label of its inputs (Sabelfeld & Myers 2003). For example, the output of a function that takes two arguments, one trusted and one untrusted, would be labeled as untrusted. A challenge of applying such label propagation mechanisms to LLMs is that the output label would be the upper bound of all imputs (i.c., the contert) used for inference. With L.MSs having the ability to retrieve documents from different sources, this can quickly become unnecessarily restrictive, a phenomenon known in the literature as label creep (Sabelfeld & Myers 2003).
In this paper, we propose a novel approach to propagate more permissive information-flow labels in LLM-basod applications which we will refer lo as an mulluence-based label propagator (LP). The key idea of our approach is to propagate only the labels the samples that were influential in generating the model's output and drop the labels of the inputs that are not. Specifically, for a given context and fixed tolerance X, we identify all subcontexts for which the model achieves utility that is most À below the utility of the full context. Within those subcontexts, we then select the with the most permissive label. We prove that, under
2

---

Published in Transactions on Machine Learning Research (10/2025)
idealizing assumptions, algorithm identifies the most permissive label(s) possible. An example of our approach is shown in Figure]
To ground our work in relevant LLM system, we implement and evaluate two different realizations of our label propagator: (i) a prompt-based system (Lewis et al., 2020) where the retrieved documents are provided within the prompt to an autoregressive LLM, and (ii) ANN-LM architecture (Khandelwal al. 2020), where the output distribution is computed as mixture of the distribution of the model and the retrieved documents.
We demonstrate the effectiveness of our proposed label propagator by evaluating it on three datasets: (i) synthetic dataset containing personal details for evaluating the label propagator's ability to handle a large number of inpuls, (ii) a news article dataset assessing whether the LP is able to handle long free-form natural language, and (iii) a dataset consisting of LLM agent conversations as depicted Figure | that mimics the applications used with agent frameworks.
Summary of contributions
We formulate the problem of security label propagation for retrioval-augmonted LLMs and LLM agents.
We propose a permissive approach that propagates only the labels of influential inputs while maintaining safety and never propagating an overly permissive label.
We show that our permissive approach correctly identifies the exact labels in at least 75% of the cases and improves the label in at least 50% of the cases on all datasets, thus mitigating the problem of label creep.
We show that influence-basee based label propagation can lead to more permissive labels without degrading the quality of the LLM output.
2 Problem Setting and Goal
We consider a general inference scenario in which an LLM takes as input a textual prompt. x and a set of documents C, which we refer as the contezt. We refer to any subset S C a subcontent. Given prompt æ and context C, we represent the LLM as probability distribution (y/ C) over possible completions y.
This formulation is sufficiently general to represent many real-world applications of LLMs. For example, in retrieval-augmennte generation (RAG) [Lwis et al. 2020 Gun et al. 2020), the context contains documents retrieved from the knowledge-base. Alternatively, if an LLM uses plugins to retrieve external data (e.g., web search, email retrieval, calendar query, etc.) (Schick et al. 2023 the data returned by the plugins is part of the context.
2.1 Information-Flow. Labels
We assume that each document E C' in the context is assigned label from set L of labels, which we model as a label assignment function C:C. - L. Labels can be used for many purposes, including representing access control information or information about the reliability of the source of the document.
As common practice (Denning) 1976 Myers & Liskov 1997 Sabelfeld & Myers 2003), we assume L forms lattice i.c., it has a partial order which every pair labels L1, Le has least upper bound (aka join) T., L. L2, and greatest lower bound (aka meet) T., r7 L22 With this, can naturally define the label L of context C as L = Loce ((c). In all cases, labels lower in the lattice are said to be more permissive, illustrated in the following examples:
Without loss of generality, could document, webpage, email, previous output the LLM in multi-turn interaction, etc. this work, we consider the context be unordered set.
the term document to refer individual piece of textual data in the context. This
3

---

Published in Transactions on Machine Learning Research (10/2025)
(LoInt, LastMonth)
(LoInt, Last Week)
(Hilnt, Lasth)nth)
(LoInt, Today)
(Hilnt, LastWeck Week)
(Hilnt, Today)
Figure 2:
Illustration of a product lattice of labels for integrity {Hilnt, LoInt} and time (LastMonth, Last Week, Today). Each dimension is a sub-lattice with a total order S. The product lattice is the Cartesian product the two sub-lattices with a partial order E.
Confidentiality. A canonical example of a security lattice the set (Secret, General), denoting high and low confidentiality data, where General Ee Secret. The join of Secret and General is Secret (i.e., high confidentiality), but General is the more permissive label.
Integrity. Similarly for integrity, the set (Hilnt, LoInt) denotes high and low-integrity data such that Hilnt C Lolnt. The join of LoInt and Hilnt is LoInt (i.e., low integrity), but Hilnt is the more permissive label.
Separation of Duty, Another example is the lattice where labels are subsets users and where assigning label U particular action denotes that all users E U must authorize the action before it can take place. In this case, the join and meet operations correspond to set union (least permissive) and set intersection (most permissive), respectively.
Product lattices. Figure 2 shows an example of product lattice for the case of two dimensions: reliability (Hilnt, LoInt)) and timestamps (Lntt Last Week, Today). The top of the lattice (Loont, LastMonth) represents the least reliable and least recent documents, while the bottom of the lattice (Hilnt, Today) represents the most reliable and most recent documents.
Obtaining labels Example scenarios where labels are readily available are corporate email servers, which label emails from outside as "external", productivity suites such as Google Workspace or Microsoft. SharePoing which implement document classification labels and RBAC, and corporate search engines, which retrieve data from internal and external sources and label them accordingly. Tool-angmentation provides another opportunity to integrate labels by labelling inputs and outputs in tool manifests, one-off effort that can be factored into the design of each tool.
Threat model We assume that input labels are given and correct. In the most general case an adversary lives at some point in the information flow lattice and can tamper with documents or tools al below that level. For the integrity and confidentiality lattices, this means the adversary can poison data labeled as untrusted and read data labeled as public. The adversary's goal to cause illicit flows, such as data from low integrity labels (poisoned data) alfecting output that has high integrity label, or input data labelled as confidential affecting outputs labeled as public.
http:///laarn.nccoooott.. www.:etnninnooooccooooooooo http://a....toottcoeoeeaaaaaaaae http:///////////////attattttttte
4

---

Published in Transactions on Machine Learning Research (10/2025)
2.2 Goal: Permissive Label Propagation
Since the documents in the context of an LLM can influence the model's output, the label L of the output will depend on the labels of documents in the context. We refer to the process of determining L as label propagation. A naïve approach to label propagation in LLMs is to assume that all documents in the context. impact the output, and hence to propagate the label L = Leec E(c).
However, the LLM does not necessarily need all documents the context to generate the output. In the example in Figure [ the LLM did not need the web search result from the untrusted website. Labeling the output as untrusted would be overly pessimistic and possibly inhibit the system from using the output further, e.g., as the input to another tool that requires trusted data.
Our goal is to obtain more permissive output label by propagating only the labels of the inputs that are actually necessary for generating the output. However, since the lattice of labels only forms a partial order, one challenge is that different subsets of the inputs are not always comparable, As consequence, we cannot hope to find a single optimal label. Instead, a label propagator should idcally find all minimal labels and delegate the selection of one to the underlying application.
3 Permissive Label Propagation
Our core idea is to propagate only the labels of the documents in the context that were necessary for generating the output. A naïve solution to this problem would be to iterate over all subsets of documents in the context and determine whether they can be removed without significantly affecting the output. However, the computational cost of this approach grows exponentially with the number of documents in the context.
We address this issue with the observation that it is sufficient to identify the labels that improve full context's label, and consider only their corresponding subcontexts. As the label lattices that occur in practice are often small (e.g. secret vs public, trusted vs untrusted), this leads to solution that is both practical and optimal under a certain monotonicity assumption (mentioned in Eq. 3).
the
3.1 s-similar Labels
Let C' be context with label L = Llee E(c). For label L' that is at least permissive L (i.e., I' [ L), we define the I'-subcontext Cir of C as the set of all documents whose label is at or below L', i.c., CL = E C ((c) C L'}. Clearly, C1L = C' because all the documents in C' satisfy E(c) E L by definition.
Our goal is to find labels L' E L such that the output of the language model changes only negligibly when substituting with Clp. We capture 'nogligible change' by introducing a hyperparameter À and require that the utility of the model's output under the full context. drops by at most when restricting to the subcontext. Formally:
Definition 1 (A-similar labels). Let x be a prompt, y a completion, and C a context with label L. For a given utility metric U and hyperparametter À 2 O, we consider another label L' be X-similar to L if
U(PIM(yl=, C)) UpLLM(y Gp)) SA
(1)
Note that X-similarity is not equivalence relation because it is not symmetric transitive.
Definition leaves the choice of the utility function U and the language model (LM C) unspecified, as different applications require custom choices of these functions. In this paper, we focus on language modeling
where perplexity a common way to measure utility (Kaplan et al. 2020). Hence, for the remainder of this paper, we compute utility as the negative perplexity:
-1/Ml
| U UpLLMyy2e, ()),())))))=>> =- (LM (z, C.
(2)
5

---

Published in Transactions on Machine Learning Research (10/2025)
Algorithm 1: X-similar label search
parameter:
language model PLM, threshold X input :onteext C, label I,, prompt 2; completion y-
output: : Set A of labels X-similar to L Function aianimal_labeLli (C, x));))
A A 0 for L' E children(L) (L) do
3
$ - (c) E ]
4
if U(LLM(yl,, ()) - UPLMayya S)) then A AU minimal labels (C, v,,,,)
5
end end
if A A then return A
10
else
11
return (T)
end 12
3.2
Computing X-similar Labels
We describe our algorithm for identifying X-similar labels and their contexts. As we observed before, is not necessary to iterate over all subsets of the context: it suffices to iterate over all labels below the full context's label and consider their corresponding subcontexts. Technically; we ilerate over the powerset of P(lllce E(c) C' S C}) of possible output labels and we compute the similarity of the corresponding subcontexts to the full context. Note that the set of labels considered is based on documents in the context; it is finite even when the full lattice of labels C is infinite (e.g. timestamps).
Algorithm D describes this idca in pscudocode. We represent the powerset of possible labels as directed acyclic graph (DAG) where nodes represent labels and edges represent the lattice order (see Figure 2 for illustration). Starting from the root node L = =leec f(c) that corresponds to the full context, traverse the DAG depth-first to identify X-similar labels. For each label L, the function minimal labels () returns A, the set of A-similar labels at or below L (i.c., at least as permissive as L).
3.3 Correctness
The set of labels returned by Algorithm is minimal in that elements pairwise incomparable with respect to the lattice order, i.e., no label is more restrictive than the other (and hence redundant). The algorithm achieves this by recursing on each child with a more permissive but X-similar label than the parent node, and adding a node's label only there is no such child. If there is a total order over the labels (i.e., all elements are pairwise comparable) then only a single label is returned (IA) = 1). Note that the labels do not need to form a tree, so Algorithm y may visit node multiple times. This can be avoided by keeping track of visitod nodes, which we forgo for simplicity of presentation.
Algorithm improves over the naive solution by iterating only over full L-subcontexts (and not over all subsets). If a child label L' is not X-similar, we prunc the scarch assuming that also none of the children of will be X-similar. If the model's utility is monotonous in the context, as in
CCC' (ppi C)) Uppim(sle,C")
i.e, adding more documents to the context never decreases the utility, we can guarantee that Algorithm identifies all minimal labels. Proposition summarizes these guarantees.
Proposition 1. Algorithum | always terminates and returns a minimal set of X-similar labels. If the utility function is monotonous, then Algorithm 7 returns all minimal X-similar labels.

---

Published in Transactions on Machine Learning Research (10/2025)
As utility functions are not necessarily monotonous (Shi et al. 2023 Algorithm 1 is heuristic in practice.
In Section 5 we evaluate how closely it matches the statement in Proposition
D
Even finding sub-optimal
labels can improve over the baseline label propagator.
3.4
A System for Label Propagation
To integrate Algorithm into existing model architectures and systems, we consider two architectures for augmenting language models with retrieved information: prompt-based retrieval augmentation (ewiss el. al., Guu et al. 2020 and kNN language models [Khandelwal et al. 2020
2020
Prompt-based Augmentation Augmenting LLM prompls with retrieved documents has become a popular
approach to incorporate paraaarmettic datastore into an LLM-based pipeline (Lewis al. 2020 Guu et al,,
2020). For example, in Retrieval Augmented Generation (RAG) setup (Lewis et al.) 2020 the documents most relevant for a query are retrieved and added to the context in the prompt to an autoregressive language model.
ANN Language Models In the ENN-LM architecture (Khandelwal et al., 2020), the language model produces an output distribution (LM x) without taking into account the retrieved documents. A separate distribution PeNN (y) xx C) is computed based on the retrieved documents only, using as the criteria. for document selection. Finally, mixture of both distributions (parameterized by a hyperparameter ?) produces the final retrieval angmented output which is given by
plyla, C) "PANN C) + +1 (1-7))Lx((1x))
(4)
While PLM requires expensive LLM inference, the expression PLNN can be computed efficiently by pre- computing key-value store with a single forward pass over the datastore that then queried at inference
time.
Label Propagation Wrapper Our approach can be integrated into both of the above architectures wrapper around the existing system. We assume that every document e that can be retrieved has a label E(c). In practice, if this is not the case, unlabeled documents can be assigned the most restrictive label (i.e. the top of the lattice). The process then proceeds as follows:
Given a prompt x, retrieve the context C and run the model usual on (x, C) to obtain the original completion y-
2. Compute the pessimistic label L of C, and run Algorithm ! on LC,L,x,y) to obtain set A of labels that are X-similar to L.
3. Choose an appropriate L' E A using application-dependaant criteria and run the model again on
(I, C to obtain a new completion y'-
Return the new completion y' and the new label L/.
Safety In practice, our algorithm is an based heuristic that can make mistakes or be misled adversarially (e.g., failing to identify an influential document, or over-estinating the influence of document).. However, these mistakes do not affect the safety of the propagated labels. By rerunning the model on the new context Ci (step 3) and returning only y/ (step 4p, our approach guarantees that the returned completion depends only on documents at or below L'. This safety property holds even in the case of adversarial input documents (c.g., prompt injection) because it enforced by the system, rather than the model.
3.5 Computational Cost
The label propagator wrapper based on prompt-based augmentation requires a number of LLM calls that, in the worst case, is exponential in the number of documents in the context. This occurs e.g., in a flat. bounded lattice where all labels between 1 and T are incomparable and when each subcontext has a different label.
However, the number of LLM queries is always bounded by the size of the lattice. When the lattice forms a totally ordered set (e.g. two-element lattices distinguishing between trusted untrusted or confidential
7

---

Published in Transactions on Machine Learning Research (10/2025)
ABCD
ABC
ABD
ACD
BCD
AB
AC
AD
BC
BD
CD
C
Figure 3: Illustration of the lattice for the synthetic key-value dataset with documents. query requires multiple documents to produce the correct response, the corresponding label is the joint label of the documents.
Sample documents from the synthetic dataset
[Doc A] The social security number and date of birth of person 1 is SSN00038242 and 26- 10-192.
[Doe BJ The social security number of person 2 is SSN00092411.
[Doc. Cj The date of birth of person 2 is 18-08-1992.
[Doc DJ The social security number and date of birth of person 2 is SSN00092411 and 18-08-1992.
Sample QA pair from the synthetic dataaset
[Question] What are the social security numbers and date of birth of person 1, and person [Answer) The social security number and date of birth of person 1 is SSN00038242 and 26 16-10-1962, and person SSN00092411 and 18-08-1992. [Label] {ABC, AD)
Figure 4: Sample documents and QA pairs from the synthetic key-value dataset. The question refers to the social security mumbers and dates of birth of person 1 and 2. This information can be obtained by accessing documents A, B, and C or A, and D, hence the resulting label of {ABC, AD).
public data), Algorithm stops as soon as it finds a subcontext whose utility drops below À difference w.r.t. the utility the full context. For richer lattices describing more fine-grained security policies, Algorithm [ visits small subset. of the lattice in typical queries, either because only few labels represented in the context because the utility drops below the tolerance X. Therefore, computational costs not major concern for several fundamental lattices. We discuss several optimizations in Section 6.3 that reduce the cost of rumning the system. Liu et al. (024aa also proposed tricks to further accelerate the computation of these context attributions.
avoid worst-case costs, allernative algorithms or model architectures may be more suitable. For example, the ANN-based architecture we consider in this paper only requires one additional LLM call per query (PLM(y)at v), in addition to the original query PLM(ylæ,CC C)).
4 Evaluation Setup
As this novel problem sctup, we first define robust evaluation scheme that we use evaluate our approach and compare against an introspection baseline. Our goal is understand the performance of our label propagator in correctly identifying and propagating minimal labels.
8

---

Published in Transactions on Machine Learning Research (10/2025)
Sample document from the news article dataset
(High Integrity/ Apple cuts prices on lower-end iPads, releases red iPhones
Apple is cutting prices two iPad models and introducing red iPhones, but the company held back updating higher-end iPad Pro tablets.
A much-speeulated 10.5-inch iPad Pro didn't materialize, did versions of existing sizes in the Pro lineup, which is aimed at businesses and creative professionals. 'he new devices are mostly refreshes of existing models. Apple unveiled them through press releases Tuesday rather than staged event, it typically does for bigger product releases.
[Low Integrity] Apple cuts prices, on lower-end iPads, adds colors to the iPhone lineup While the iPad Pro tablets didn't get an update, the two lower-end iPad models got a $100 price cut today, unveiled through quiet press release rather than large staged event. With fans clamoring for greater variety of colors for their iPhones, Apple announced in the same release five fruit- rrit-insspied colors, hearkening to the flavors of the iMac G3 in 1998. The new colors, available starting next Tuesday, are Cherry (red), Lemon (yellow), Lime (green), Blueberry (blue), and Grape (purple).
Sample QA pair from news article clttaset
[Question] What are the new fruit-inspired colors for the iPhone lineup mentioned in the article about Apple cutting prices lower-end iPads and adding colors to the iPhone lineup? [Answer) The new fruit-inspired colors for the iPhone lineup mentioned in the article are Cherry (red), Lemon (yellow), Lime (green), Blueberry (blue), and Grape (purple). [Label] [Loint}
(Question/ What is the new color introduced for the iPhone according to the article about Apple cutting prices on lower-end iPads and releasing red iPhones? (Answer] According to the article, the new color introduced for the iPhone is red.
[Label] (Hilnt}
Figure 5: Sample documents and QA pairs from the news article dataset. The first question only be answered with access to the Lolnt document whereas the second question can be also answered with the Hilnt document.
4.1 Research Questions
The goal the label propagator is to return all minimal A-similar labels. As shown Proposition | this is achieved under monotonicity assumptions that are typically not satisfied in practice. Therefore, our evaluation aims to quantify how closely the empirical results match this goal. Specifically, aim to answer the following rescarch questions:
RQI: How accurate is our label propagator in identifying the set of minimal labels?
RQ2: By how much does our label propagator improve over a naive label propagator?
RQ3: How aligned is the regenerated output for the inferred label with the full-conteat output?
Rescarch questions RQI and RQ2 focus on the performance of Algorithm T directly. Rescarch question RQ3 focuses on the quality of the final output of the end-to-end system introduced in Section 3.4
4.2 Evaluation Metrics
We now describe the metrics that we use answer the aforementioned rescarch questions. Throughout, we use A denote the set of labels returned by Algorithm | and A as ground truth, i.c., the correct set of minimal labels.
9

---

Published in Transactions on Machine Learning Research (10/2025)
Exact match The exact match metric computes the average number of completely correct predictions over all questions. That is, for each question we count. 1 if A A and 0 otherwise. Exact match gives a direct. answer to RQI, but is very sensitive in that it equally penalizes any imprecision in the label propagator.
Precision and Recall For a more fine-grained evaluation, we compute precision and recall between A and A for every question, iee, |A A A| and |A A / and average these over the entire dataset. Note that precision and recall simultancously reach their maximum of if and only if A", i.c., we have an exact match.
Label improvement When the lattice forms total order (i.e., any two labels are comparable), both A and A become singleton sets. Therefore, in this case, instead of reporting precision and recall over set outputs, we instead report label improvement and missed labels. We define label improvement as the number of cases in which our system improves the output label, as a fraction of the total number of cases where label improvement is possible. Missed labels, on the other hand, defined as the number of cases our system misses a label from the context, as fraction of the total number of cases where our system improves the label. Note that in the case of more than two labels, these metrics only consider perfect improvement (i.e., the output label exactly matches the ground truth label).
Model output alignment The system for label propagation described in Section 3.4 regencrates the output based on the inferred label. While the definition of A-similarity guarantees that under the subcontext of the inferred label the original output is almost as likely as in the full context, this does not necessarily gnarantee that the regenerated output has good utility w.r.t. to other metrics. To quantify potential deviations, we measure the alignment of the different model outputs w.r.t. the target specified in the dataset i.e., y² using ROUGE-L [Lin 2004) metric, which is based on the Longest Common Subsequence (I.CS) (Hunt & Szymanski 1977 computation and has been commonly used to estimate translation quality in the past.
In particular, we report the average alignment of the full context output ROUGE-L(y, y²), and the average alignment of the reduced context output after label propagation ROUGE-L(y's y²). Furthermore, we report the average difference of the alignment between the reduced context output and full context output i..,, ROUGE-L(y', y²) ROUGE-L(y, )) in two different cases i.e, (i) when label improvement is possible, and (ii) when label improvement is not possible. The difference in output alignment should be nearly zero when label improvement is possible, and highly negative when such improvement is not possible.
4.3 Baseline
As bascline for comparison, we use introspection in which the LLM itself askod to determine which documents in the context were influential. This is inspired by the recent successful application of language models to generate relevant citations for their own outputs (Taylor et al. 2022: Gao et al 2023). Note that the performance of the introspection technique depends significantly on the prompting technique used, and hence, should be considered as a lower-bound (Turpin et al. 2024
4.4 Models
We focus on the Llama-2 model family (Touvron et al. 2023 including its 7B and 70B variants for evaluation. We use instruction- tuned versions of Llama-2 to ensure accurate response generation unless mentioned otherwise. For the ANN-LM implementation, we follow (Khandelwal et al. 2020) and use the model's penultimate layer representation of the last token (conditioned on all precoding tokens in the document) as the context representation for &NN search. Note that while we rely on Llama-2 model family (Touvron et al.) 2023) for our experiments, our approach is also applicable to other open-source models, or even proprietary models accessible only via API (though client-side optimizations will be harder to implement).
4.5 Datasets
Since the problem of label propagation for LLMs has not been previously studied, there are no off thol datasets available for evaluation. We thus design three datasets evaluate different aspects of our label
10

---

Published in Transactions on Machine Learning Research (10/2025)
propagator. Each dataset consists of a set of documents and set of corresponding question-answer pairs. The documents and answers all have labels. The goal of the label propagator is to identify the subsets of documents in a given context that are required for answering the question and have label that is at least as permissive as the naive label propagator.
We use the target response in the dataset as the completion y for Algorithm D i.e., we do not rely on the model for generating the output unless mentioned otherwise. This is motivated by the fact that the target label computed in the dataset is only correct w.r.t. the target response. Therefore, an incorrect completion y from the model would render the ground-truth label set A incorrect. We quantify the implications of this decision in Section 5.3 where we compare the difference in performance between the dataset target and the model-generated output.
Synthetic key-value dataset We create set of key-value pairs that contain hypothetical individuals' IDs as keys, and their Social Security Numbers (SSNs) and dates of birth (DoB) as values. We randomly split, distribute, and replicate the key-value pairs across multiple documents, and we attach a unique label to each document. We design the questions such that obtaining the correct answer requires identifying the different combinations of documents that contain all the necessary values. An example of the generated documents and QA pairs is shown in Figure 4 In this example, the question can be answered with access to documents A, B, and C or with access to documents A and D. Consequently, the label scarch should yicld {ABC- = AU B U C, AD = A).
Since each document has a unique label, it corresponds to a node at the bottom of the lattice (similar to the example shown in Figure 3 Each set of documents (i.c., subset of the context) therefore corresponds to a unique element in the lattice. The documents and question-answer pairs are designed such that the question can be answered from different combinations of subsets in the context. Identifying all minimal labels corresponds to identifying all these subsets.
We generated a lotal of 128 documents (i.e., 128 labels) and 64 question-answer pairs. To ensure computational tractability, we set the context size to be 14 documents for each question. We ensure that all necessary documents are included in the context, thus emulating the case of a retriever component with perfect recall (all necessary documents are present) but lower precision (the context may contain irrelevant documents). This dataset represents challenge terms of the complexity and size of the search space because the full set of labels contains 2 clements for each question. Since there is only a partial order in this lattice, we use the precision and recall metrics for evaluation.
News article dataset Starting from an existing fake news datasct (Pérzz-Rosse Rosas et al.] 2018 1 we create pairs of high and low-integrity news articles that discuss similar topics to each other. Contrary to the previous datasct, we focus on simpler lattice but more complex language in this case. Using GPT-4 (OpenAI 2023 we generated QA pairs based on these articles, where some of the answers depended only on the LoInt document, others only on the Hilnt document, and some on both documents. An example document pair and QA pair is shown in Figure 5 The dataset in total contains 240 document pairs of corresponding LoInt and Hilnt documents, as well as 3465 QA pairs. Compared to a naïve label propagator, is possible to improve the label on 647 of these QA pairs. Since there is a total order in the lattice, we report label improvement. metrics for this dataset.
LLM agent dataset Motivated by the recent emergence of L.M agents (Xi et al. 2023) and tool use in LLMs (Schick et al., 2023), we design an LLM agent dataset that focuses on tool use. Retricval tools such web or email search typically return multiple results of varying integrity levels. For example, in the case of web search, official documentation like docs python. org might be deemed to be higher integrity than community curated sources such as stackoverflow. com. Similarly, emails from verified senders may be deemed to be higher integrity than those from unknown senders.
To demonstrate this, we create a small dataset of chat conversations between a user and an assistant. In cach conversation, the LLM calls one of three retricval tools: web scarch, email scarch, or calendar search. Each piece of retrieved data carries a label L E (Hilnt, LoInt) indicating high or low integrity respectively.
11

---

Published in Transactions on Machine Learning Research (10/2025)
An example from the Tool Calling LLM case study
UseT
How many people live in New York City?
Model
Tool: search query
Relevant context: Need to perform a search query to find the latest population data of New York City
Args: { ('query_stay str'; 'current population of New York City')
Tool
Source
Information
Security Label
population New York City is approximately million as of 2021.
Hilnt
2.
New York City has population of around million.
Lolnt
Model
Tool: respond
Relevant context: Need respond back the user's query
Args:
'text str': "The current population of New York City is approximately 8.4 million.
An example interaction from the tool calling LLM case study, where the user query triggers a web search. The items returned by the tool have different labels. The aim of our label propagator is to identify the label for the final model output. Since the high-integrity output is sufficient to answer the query in this case, the label propagator can upgrade the label of the output to high integrity instead of the naive baseline of propagating the low-integrity label.
Figure 6:
In the example shown in Figure 6 the web search retrieves two documents that both contain sufficient information to answer the question. However, one of the documents carries a low-integrity label and would therefore force the output be LoInt. The goal of the label propagator to identify the influence of the high-integrity document and assign Hilnt to the output.
To cover wider range of cases, manually create 40 distinct chat conversations that involve tool calls. Some of these tools return set of two documents C. We also generate ground truth model output y* and label L* In 20 of the cases, the final label can be improved i.e. L C A(C') due to redundant or irrelevant information within the retrieved LoInt documents. In the remaining 20 cases, the LoInt source is required to produce the output and thus the output label cannot be improved i.e. L ((C). We additionally include in-context examples that are specifically used to specify the output format for the model. Since there is total order in the lattice, we report label improvement metrics.
5 Results
5.1
Synthetic key-value dataset
In this case, we assume perfect recall retriever (all relevant documents are present), albeit lower precision (the rest of the documents out of the total limit of 14 are filled by adding irrelevant documents). For the introspection bascline, we use one-shot in-context learning (Wei et al. 2022) based on the first example in the dataset and the remaining 63 questions for evaluation. Table summarizes results.
12

---

Published in Transactions on Machine Learning Research (10/2025)
Model
Tabel Prediction Method
Exact Match
Precision
Recall
Prompt-based
85.94%
92.34 + 23.27 %
93.75 + 22.10 %
JJama-2-Ches (7B)
kNN-LM
53.12%
64.58 = 45.79 %
65.10 + 45.33 %
Introspection
1.59%
3.25 + 17.53 %
3.97 + 18.48 %
Prompt-based
85.94%
94.17 + 21.84 %
92.06 + 23.45 %
Llama-2-Chan (70B)
kNN-LM
57.81%
69.53 = 45.60 %
64.45 + 44.74 %
Introspection
12.70%
15.07 = 34.94 %
15.87 + 35.44 %
Table 1: Results on the synthetic key-value dataset assuming a perfect recall retriever, which always retrieves the relevant items while also retrieving some irrelevant ones, with person ID as the key, and person SSN and DoB as values. kNN-LM prediction uses y = 0.5. We report the mean and one standard deviation for macro-averaged precision and recall.
100
LLaMa-2 (7b)
LLaMa-2 (70b)
Exact match accuracy (%)
80
60
40
20
0.0 0.001 0.003 0.01 0.025 0.05 0.075 0.1
0.25
0.5
0.75
1.0
Perplexity tolerance
Figure 7: Hyperparameter grid search on perplexity tolerance À for the 7B and 70B models on the synthetic key-value dataset. See Figure for the full grid search over both Shapley threshold and perplexity tolerance.
The prompt-basod propagator achieves an exact match accuracy of over 85% and precision and recall of over 90%. That is, in more than 85% of the cases, the label search can identily the correct influential subcontexts out of set of 16k possibilities without error. Furthermore, the prompt based label search significantly outperforms both the KNN-LM-based label search well as the introspection baseline.
Table D also compares the exact match accuracy for two different model sizes 7B and 70B parameters. The model size does not strongly correlate with the performance of the prompt based label propagator. Consequently, this enables the use of significantly smaller models for effective label propagation on the outputs of larger models.
On the other hand, introspection performance improves significantly with the model scale. This result is consistent with the improved in-context learning ability observed in larger models (Wei al. 2022). Therefore, only the largest and most powerful models can be effectively used for introspection.
Sensitivity of hyperparameters The perplexity tolerance À mcasures the acceptable loss in model utility when removing a subset of documents. The choice of X allows to trade-off model utility and the restrictiveness of the inferred label.
We illustaate the role of the perplexity tolerance À on the resulting change the exact match performance for our prompt based label propagator in Figure 7 The figure indicates that the model performance is very sensitive to the choice of perplexity tolerance ie., either very high very low values significantly hamper
13

---

Published in Transactions on Machine Learning Research (10/2025)
Retriever Type
Model
Exact Match
Label Improvement
Missed Labels
Mamma-2-hhe (7B)
74.69%
56.72%
73%
Perfect Retriever
Llama-2- 2-Chat (70B)
73.45%
39.26%
18.59%
Llama-2 2- Chaat (7B)
63.78%
49.77%
30.30%
Realistic Retriever
Llama-2-Chat (70B) 64.50%
39.88%
26.08%
Table 2:
Results the news article dataset with prompt-based label search. The perfect retriever row assumes that the retriever always retrieves the relevant items while also retrieving some irrelevant ones randomly. The realistic retriever uses a cosine similarity-base neigeottr search and may not retrieve all relevant items, thereby roducing the label improvement metric of even a perfect label scarch.
model performance. Furthermore, the optimal hyperparameters are consistent across model scales. However, larger models, being better at language modeling, outperform their smaller counterparts at lower values of X.
Although the bafeed label propagator performs very well on its own, we have found that pruming irrelevant labels can elevate the performance further. This is particularly the case for large lattices with partial order. We introduce Shapley-value-basee heuristic pruning technique in Appendix A that can marginally improve the performance from exact match accuracy of 81% to 86%. The results reported in Table D include this saalee--vlue-baaee heuristic.
5.2 News article dataset
Moving towards a more realistic setup, we include an additional retriever component for the news article dataset that first retrieves relevant articles (either low or high integrity) from the datasel before response generation. In order to correctly understand the impact of the retriever, we compare a perfect retriever that is able to retrieve all relevant documents (similar to the synthetic key-value dataset) with a realistic retriever using cosine similarity in the embedding space computed by BGE-Larye-Bl (Xiao et al., 2023b
Table _ summarizes the results of the label propagator the news article datasct. For this dataset, the label propagator achieves slightly worse performance on all metrics in comparison to the synthetic key-value dataset due to the lower quality of the dataset (automated generation of QA pairs by GPT-4 (OpenAl 2023)). Furthermore, we see reduction in the exact match accuracy of about 10% when using a realistic retriever in contrast to a perfect retriever. This due to the retriever sometimes failing to retrieve the relevant pieces of information, ultimately leading to mismatch with the ground-truth label.
Due to the presence of total order on the labels, we report label improvement and missed labels for this dataset. Even when perfect retriever is used, naive label propagator would be overly conservative in 647 out of 3465 examples where label improvement is possible. In contrast, our influence-based label propagator is able to assign the correct label in ~ 57% of these cases. On the other hand, our label propagator suggested a more permissive label but missod a label from the ground-truth in ~ 23% of the cases. However, as explained in Section 3.4 the safety properly still holds these cases because the system regenerates the output using only the articles with the more permissive label (although this means that the regenerated output may differ from the original output, as we quantify in the next section).
Sensitivity of hyperparameters We highlight the sensitivity to changes in the perplexity threshold À the news article dataset in Figure & Similar to the synthetic key-value dataset, we see similarities across the two model scales, and a higher tolerance of the larger model to lower values of x. However, we see a significant rise in optimal perplexity tolerance in contrast to the synthetic dataset due to the higher difficulty of the responses.
5.3
LLM agent dataset
In the LLM agent dataset, we take a step further by also computing the misalignment introduced by regenerating the output conditioned on the updated context (when label improvement is possible). The
14

---

Published in Transactions on Machine Learning Research (10/2025)
100
LLaMa-2 (7b)
LLaMa-2 (70b)
80
Exact match accuracy (%)
40
20
0.001 0.003 0.01 0.025 0.05 0.075 0.1
0.25
0.5
0.75
1.0
Perplexity tolerance
Figure 8: Hyperparameter grid search perplexity tolerance X for the 7B and 70B models the news article dataset with perfect retriever. See Figure I0 for the full grid search over both Shapley threshold and perplexity tolerance.
Model
Label Prediction Method
Utility Target
Label Improvement
Missed Labels
Prompt-based
Model generated y
85.0%
5.56%
Llama-2 (7B)
Introspection
45.0%
0.0%
Prompt-based
Ground truth y"
100.0%
0.0%
Prompt-based
Model generated y
95.0%
9.53%
Llama-2 (70B)
Introspection
80.0%
11.11%
Prompt-basod
Ground truth y"
100.0%
0.0%
Table 3: Results on the LLM agent dataset. The utility target column indicates which output y we are using to compute target utility U(PLM(yla, C)) in equation 1 In practice and in the absence of ground truth, the utility target is the model-generated output y, but also compare with the ground truth y* to understand the maximum possible utility assuming perfect LLM. The difference in the numbers between the cases indicates that all the errors in the case of model-generated output are artifacts of the model generating a suboptimal response which causes incorrect propagation of labels.
results presented in Table 3
We fix threshold of À 0.2 for this experiment. We use the base model without instruction tuning in this case due to the use of custom chat format.
First, we focus on the 20 cases where label improvement is possible i.e. L C E(C). We find that the label propagator improves the label at least 17 of these cases.
In at most 2 of the cases, the LP overly optimistic and returns more permissive label L' than the ground truth label L*. Again, the safety property still holds in these cases because the system regenerates the output using the reduced context Cjz. This regeneration step introduces the possibility that the new output y differs significantly compared to the output y from the full context.
To quantify this, we measure the aligmment between y and y² as well as y' and y" using the ROUGE-L F-score (see Table 4p. We obtain a ROUGE-L F-score of at least 0.85, suggesting that conditioning on the reduced context leads to a similar model response. We see negligible difference in alignment between the full context output and the reduced context output when label improvement possible, but observe significant drop when such upgrade is not possible.
15

---

Published in Transactions on Machine Learning Research (10/2025)
Model
ROUGE-I(W.V")
ROUGE-L.(y'''")
ROUGF-L(y'."" ROUGE-L(U..")
Improvement possible Not possible
Llama-2 (7B)
0.77
0.85
0.0051
-0.34
Llama-2 (70B)
0.82
0.90
0.025
-0.31
Table 4:
Model output alignment with the dataset ground truth computed using ROUGE-L F-score on the LLM agent datasct. Utility computation is only applicable to the prompt-based label propagator. ROUGE- (") refers to the alignment of the full context output (no regeneration required), while ROUGE-L(y/ yy*) represents the alignment of the regenerated response after label improvement. Interestingly, we see that the regenerated response is more aligned with the ground truth response suggesting that in some cases label propagation can improve utility. ROUGE-L(y/ ,y"= ROUGE-L(y, y*) compares the difference in alignment between the full context and the subcontext of the inferred label. When label improvement is possible, the difference in alignment is close to nil indicating that the regenerated response y' is highly aligned with the ground truth response y" Otherwisc, the difference in alignment is highly negative, indicating a significant drop in model aligument with y"
In a robust retrieval augmentation setup, adding additional documents to the context should not negatively influence the model's ability to answer a query since the model is free to ignore irrelevant inputs (Shi et al., 2023). Thus, we expect that the ground truth completion y* always similar to the completion given the full context y. However, in almost all label propagation errors we observe, we find that the model not able to answer the query with the full context. When controlling the model generation by using the ground truth target in the dataset instead of using the model-generated output, we see that our label propagator achieves perfect precision and recall.
Since the lattice is particularly simple in this case, we observe strong performance from the introspection baseline in contrast to the more complex lattice the synthetic key-value dataset. Furthermore, similar to the previous datasets, we see significant improvement in the introspection performance with increasing model scale, highlighting that introspection might be particularly well-suited for larger models.
6 Analysis and Discussion
In this section, we highlight the main findings of our work and discuss the implications of our results when used in a real-world sctting.
RQL. We find that our prompt-based label propagator can find the exact set of minimal labels in 86% of the cases, for large lattice of 16k possible labels.
RQ2. We evaluate the label propagator smaller lattice with total order, allowing us to compare labels directly and quantify the label improvement. In this case, our label propegator improves the label in 56% of the cases for the news article dataset and 85% of the cases for the LLM agent dataset.
RQ3. We showcase the label propagator in real-world use case a tool calling LLM agent setup where the propagated label is used to determine whether a sensitive tool call allowed. We find that the prompt-based label propagator able to improve the label in more than 85% of the cases while the output of the LLM agent remains the same as measured by the difference between the reduced context output alignment and the full context output alignment.
Comparison to baselincs The iitrospection-baset label propagator achioves a noticcable improvement in performance when using a 70B parameter model compared to a 7B model. Furthermore, despite its simplicity and computational convenience, the introspection based label propagator constitutes a strong baseline when considering simple lattices. We find a similar performance of the prompt-based label propagator for both model sizes suggesting that our approach could be implemented on smaller models, saving computational resources.
16

---

Published in Transactions on Machine Learning Research (10/2025)
6.1 Use-cases
Algorithm iddentfies sub-contexts with labels that are more permissive than that of the full context. However, whenever the desired output label L can be determined up-front, it is possible to skip the search over A-similar sub-contexts and restrict the retrieval component to documents at or below L.
Our approach reveals its true benefits when the use of the generated content is not yet determined after the initial retrieval step. In such cases, having a too-restrictive label comes at the cost of limiting future uses of the generated output. For example, semantic caches (Bang 2023) can be extended to store answers along with sensitivity label, where more permissive labels facilitate broader reuse of the generated content. Likewise, Emails are often forwarded beyond their initial set of recipients, which is facilitated by using permissive labels.
Our approach naturally allows users to endorse information. Let's consider a setup as illustrated in Figure ! Initially, there might only be an untrusted source answering the initial query. The LP system would correctly output an untrusted label to the potentially dangerous command. However, after a trusted authority (e.g., a university department IT admin) confirms the suggestion in the untrusted source, the LP recognizes that both sources similarly influential and assigns a trusted label to the output. Therefore, by quoting repeating untrusted information trusted source can endorse information.
6.2 Limitations
Our label propagator is able to improve on the baseline label by more than 50% and 85% in two realistic datasets. However, this improvement comes with the cost of an increased number of LLM calls (discussed in Scction 3.5
In the presence of an adversary, our label propagator assigns the lowest integrity label by design in order to avoid any harmful side effects such cross-prompt injection (Liu et al.] 2023). However, this can potentially lead to a degradation of service attack as the adversary can add unreliable distracting information that makes it impossible for the system to improve the label, reducing the downstream utility of the model output.
6.3 Extensions
Other applications Our main focus has been on predicting the least conservative label while ensuring that utility is not compromised beyond a certain threshold X (Delinition 1j. However, our proposal is flexible enough to accommodate various use cases. For instance, it can be reversed to determine the utility for a specific target label. Alternatively, a more complex use case would be to make the system dynamic with respect to X, which means the system able to compromise more utility to achieve a better or less conservative label.
discusses the computational cost incurred to execute our label propagator in Algorithm [ Despite an increase in the number of LLM calls for label propagation, the cost incurred per call can be drastically reduced ulilizing common inference optimizations implemented in model serving backends (Zheng al. 2023 Kwon al. 2023). For instance, compute-bound prompt processing in Transformers KV cache. For a totally ordered lattice, appending retrieved documents at the end of the prompt sorted by their labels will result. in no additional prompt processing costs in subsequent. LLM calls because their prompts are obtained by peeling off documents at the end of the prompt.
Efficiency improvements to Algorithm 1 Section 3.5
be amortized across calls sharing common prompt prefix by reusing
We can also optimize the decoding phase of an LLM call by keeping a running calculation of the cumulative likelihood of the tokens decoded and stop decoding when the utility drops below difference À of the utility of the full context. Liu ct al. (2024a) also explored tricks to improve the efficiency of context attribution in LLMs.
Advanced influence estimation The label propagation mechanism developed in Section 3 based estimating the influence of context elements on the model's predictions. There several alternatives
17

---

Published in Transactions on Machine Learning Research (10/2025)
techniques, such Shapley value estimation, that can be used as drop-in replacements. See Brophy et al.
(2023] for an overview of influence ion techniques and (Nguyen & Wong 2023 Cohen- Wang et al. 2024 for examples of their application to LLM contexts.
Another approach is to leverage the attention weights for infiuence estimation. Phenomena such as attention sinks (iio et al., 2023a) or lost-in-the-middle (Tiu et al. (024b) makes identification based on simple attention scores difficult.
7 Related Work
Information-foow analysis Information-fio analysis and the use of flow labels has a long-standing history, see c.g., Denning (1976), with an early survey of language-based approaches (Sabelfeld & Myers) 2003). The operations considered in many program analysis frameworks (e.g. [Volpano et al. (1996); Sabelfeld & Myers (2003)) have well-defined semantics, and the rules for propagating the flow labels of each operator can be hard-coded. The situation is different when using ML models such as Transformers: due to their attention mechanism and autoregrossive docoding, their output tainted by all inputs. Morcover, their bchavior is only implicitly defined the solution of loss-mininization problem. In our approach, we propagate labels through MI. components under the constraint that the model's loss does not increase.
The literature distinguishes between static and dynamic approaches to information flow control (Sabelfeld & Myers. 2003 Buiras et al., 2014 Mardziel et al. 2011). The key difference is that dynamic approaches track the flow during program execution, whereas static approaches reason collectively about all possible executions without actually executing the program. Our approach relies on analyzing one specific context and is hence fundamentally dynamic.
Secure multi-cxocution (SME) (Devriese & Picssens 2010: Rafnsson & Sabelfeld 2016) is an approach to dynamic information-flow analysis where programs are executed once per security level. The idea is that an execution that produces output for level I. receives data only from levels I,' C L; and dummy values from levels L" C L. Programs that are multi-executed are secure by design, and correct if the original program was secure. Our approach is related to SME in that we also call the LLM once per security level, and we propagate the most permissive label for which the original output is still "correct", the sense that the drop in utility is bounded by X.
Information-flow analysis in AI applications Several emerging approaches rely on information-flow tracking to obtain deterministic security guarantees for LLM-based applications, even if the underlying models are vulnerable to attacks: WWutschitz al. (2023) study how to leverage existing metadata in datastores used for retrieval augmentation to enforce confidentiality guarantees at inference time. Wu et al. (2024) propose system that uses information-fion tracking to prevent LLM-based applications from being compromised by untrusted information. Balunovic et al. (2024) develop monitor and policy language that can check execution traces of LLM applications for violations of different safety and information flow properties. An emergint line of work uses data and information-flow tracking to enforce deterministic security guarantecs for the tools called by AI agents (Costa et al. 2025 Debenedetti et al. 2025 Zhong et al. 2025 In contrast to these approaches, our work aims to dynamically and soundly propagate labels through LLM calls, which can be used to increase the permissiveness of system-level information- tracking.
Related to the introspection baseline, Mireshghallah et al. (0024] and Ghalebikesabi et al. 2024 use Contextual Integrity theory to explore how well LLMS can be relied upon to discern which information is appropriate to use in given context. Finally, Wallace et al. (2024) train LLMs to distinguish between instructions in their context totally ordered in a hierarchy, teaching them to prioritize higher-privilege instructions over instructions that appear lower in the hierarchy.
8 Conclusion
We presented permissive approach to propagating information-foor labels of documents retrieved in RAG systems. The key idea is to propagate only the labels of those documents that are actually used for generating
18

---

Published in Transactions on Machine Learning Research (10/2025)
the model's output. We show that approach is practical in terms of performance and infers more permissive labels than an introspection baseline. Unlike commonly used introspection--useed methods, our approach can satisfy a hard safety guarantee.
References
Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent El Shafey, Yanping Huang, Kathy Meier-Hellstern, Gaurav Mishra, Erica Morcira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, Yuanzhong Xu, Yujing Zhang, Gustavo Hernandez Abrego, Junwhan Ahn, Jacob Austin, Paul Barham, Jan Botha, James Bradbury, Siddhartha Brahma, Kevin Brooks, Michele Catasta, Yong Cheng, Colin Cherry, Christopher A. Choquette-Choo, Aakanksha Chowdhery, Clément Crepy, Shachi Dave, Mostafa Dehghani, Sunipa Dev, Jacob Devlin, Mark Diaz, Nan Du, Ethan Dyer, Vlad Feinberg, Fangxiaoyu Feng, Vlad Fienber, Markus Freilag, Xavier Garcia, Sebastian Gehrmann, Lucas Gonzalez, Guy Gur-Ari, Steven Hand, Hadi Hashemi, Le Hou, Joshua Howland, Andrea Hu, Jeffrey Hui, Jeremy Hurwitz, Michael Isard, Abe Ittycheriah, Matthew Jagiclski, Wenhao Jia, Kathleen Konealy, Maxim Krikun, Sneha Kudugunta, Chang Lan, Katherine Lee, Benjamin Lee, Eric Li, Music Li, Wei Li, YaGuang Li, Jian Li, Hyeontaek Lim, Hanzhao Lin, Zhongtao Liu, Frederick Liu, Marcello Maggioni, Aroma Mahendru, Joshua Maynez, Vedant Misra, Maysam Moussalem, Zachary Nado, John Nham, Eric Ni, Andrew Nystrom, Alicia Parrish, Marie Pellat, Martin Polacek, Alex Polozov, Reiner Pope, Siyuan Qiao, Emily Reif, Bryan Richter, Parker Riley, Alex Castro Ros, Aurko Roy, Brennan Saeta, Rajkumar Samuel, Rence Shelby, Ambrose Slone, Daniel Smilkov, David R. So, Daniel Sohn, Simon Tokumine, Dasha Valter, Vijay Vasudevan, Kiran Vodrahalli, Xuezhi Wang, Pidong Wang, Zirui Wang, Tao Wang, John Wieting, Yuhuai Wu, Kelvin Xu, Yunhan Xu, Linting Xue, Pengcheng Yin, Jiahui Yu, Qiao Zhang, Steven Zheng, Ce Zheng, Weikang Zhou, Denny Zhou, Slav Petrov, and Yonghui Wu. PaLM 2 technical report, 2023.
Mislav Balunovic, Luca Beurer Kellner, Mare Fischer, and Martin Vechev. AI agents with formal security guarantees. In TCML 2024 Next Generation of AT Safety Workshop, 2024. URI. htttps://openreeeeee aet/foorumidd=oofMWPtttsss
Fu Bang. GPTCache: An open-source semantic cache for LLM applications enabling faster answers and cost savings. In Proceedings of the 3rd Workshop for Natural Language Processing Open Source Software (NLP-OSS 2023), pp. 212 218, 2023.
Jonathan Brophy, Zayd Hammoudeh, and Daniel Lowd. Adapting and evaluating influence-sstimation methods for mondient-boosted decision trees. J. Mach. Learn. Res., 24:154-1, 2023.
Pablo Buiras, Deian Stefan, and Alejandro Russo. On dynamic flow-sensitive floating-labe systems. In 27th IEEE Computer Security Foundations Symposium, CSF 2014, pp. 65 79. IEEE Computer Society, 2014.
Harrison Chase. Langchain, 2022. URL http:p///ggthh...cooo/ceeae////p
James Clause, Wanchum Li, and Alessandro Orso. Dytan: generic dynamic taint analysis framework. In 2007 international symposium on Software testing and analysis, pp. 196-206, 2007.
Benjamin Cohen- Wang, Harshay Shah, Kristian Georgiev, and Aleksander Madry. ContextCite: Attributing model generation to context. arXio preprint 2rXiu:2409.00729, 2024.
Manuel Costa, Boris Köpf, Aashish Kolluri, Andrew Paverd, Mark Russinovich, Ahmed Salem, Shruti Tople, Lukas Wutschitz, and Santiago Zanella- Béguulln. Securing ai agents with information flow control, 2025.
URL https::/arxiv..... wwtp://ggggabbbaaeeeeoo aas//bb////555 23643
Edoardo Debenedetti, Πla Shumailov, Tianqi Fan, Jamic Hayes, Nicholas Carlini, Daniel Fabian, Christoph Kern, Chongyang Shi, Andreas Terzis, and Florian Tramèr. Defeating prompt injections by design, 2025.
URL http:::/aaax...oo////aaaaa//
19

---

Published in Transactions on Machine Learning Research (10/2025)
Dorothy E Denning. A lattice model of secure information flow. Communications of the ACM, 19(5):236-243, 1976.
Dominique Devriese and Frank Piessens. Noninterference through secure multi-execution. In 31sl. IEEE Symposium on Security and Privacy, SP 2010, 16-19 May 2010, Berleley /Oakland, California, USA, pp. 109-124. IEEE Computer Society, 2010.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama 3 herd of models. arXiv preprint WxXi:2447.2178.3 2024.
William Enck, Peter Gilbert, Seungyeop Han, Vasant Tendulkar, Byung-Gon Chun, Landon P Cox, Jaeyeon Jung, Patrick McDaniel, and Anmol N Sheth. Taintdroid: an information-foow tracking system for realtime privacy monitoring on smartphones. ACM Transactions on Computer Systems (aninnssT (TOCS), 32(2):1 29, 2014.
Tianyu Gao, Howard Yen, Jiatong Yu, and Danqi Chen. Enabling large language models to generate text with citations, 2023.
Sahra Ghalebikesabi, Eugene Bagdasaryan, Ren Yi, Itay Yona, Ilia Shumailov, Aneesh Pappu, Chongyang Shi, Laura Weidinger, Robert Stanforth, Leonard Berrada, Pushmeet Kohli, Po-Sen Huang, and Borja Balle. Operationalizing contextual integrity in privacy-conscious assistants. arXiv preprint arXiv:2408.028) 023773 Als,AlJ, 2024. doi: (0.48550/arXivv222 210.Xi..2000.0373..
Kai Greshake, Sahar Abdelnabi, Shailesh Mishra, Christoph Endres, Thorsten Holz, and Mario Fritz. Not what you've 've signed up for: Compromising real-world LLM-integrated applications with indirect prompt injection, 2023.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938 PMLR, 2020.
James W Hunt and Thomas Szymanski. fast algorithm for computing longest common subsequences. Communications of the ACM, 20(5):350 353, 1977.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillamme Lample, Lucile Saulnier, et al. Mistral 7B. arXiv preprint arXit:2310.06825, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint orXiv:2001.0836) 2020.
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference Leurning Representations, 2020. URL https:///penrreview new. hett//oorumidd-h111
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gon- zalez, Hao Zhang, and Ion Stoica. Efficiont memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Patrick Lewis, Fthan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmentee generation for knowledge- intensive NLP tasks. Advances in Neural Information Processing Systems, 33:9450-94h 9474, 2020.
Haoran Li, Dadi Guo, Wei Fan, Mingshi Xu, Jie Huang, Fanpu Meng, and Yangqiu Song. Multi-step jailbreaking privacy attacks on ChatGPT, 2023.
Chin-Yer Lin. ROUGE: A package for automatic evaluation of summaries. In Text summarization branches oul, pp. 74 81, 2004.
20

---

Published in Transactions on Machine Learning Research (10/2025)
Fengyuan Liu, Nikhil Kandpal, and Colin Raffel. Attribot: A bag of tricks for efficiently approximating leave-one-out context attribution. arXiv preprint. ar arXiv:2411.51100 2024a.
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models long contexts. Transactions of the Associalion for Computational Linguistics, 12:157-173, 2024b.
Yi Liu, Gelei Deng, Yuekang Li, Kailong Wang, Zihao Wang, Xiaofeng Wang, Tianwei Zhang, Yepang Liu, Haoyu Wang, Yan Zheng, et al. Prompt injection attack against LLM-integrated applications. arXiv preprinl ar Xiv:23006 05499, 2023.
Piotr Mardziel, Stephen Magill, Michael Hicks, and Mudhakar Srivatsa. Dynamic enforcement of knowledge- based security policies. In Proceedings IEEE CSF, pp. 114-128. IEEE Computer Society, 2011.
Niloofar Mireshghallah, Hyunwoo Kim, Xuhui Zhou, Yulia Tsvetkov, Maarten Sap, Reza Shokri, and Yejin Choi. Can LLMs keep secret? testing privacy implications of language models via contextual integrity theory, 2023.
Niloofar Mireshghallah, Hyunwoo Kim, Xuhui Zhou, Yulia Tsvetkov, Maarten Sap, Reza Shokri, and Yojin Choi. Can LLMs keep secret? testing privacy implications of language models via contextual integrity theory. In 12th International Conference on Learning Representations, ICLR 2024. OpenReview.net, 2024.
Andrew C. Myers and Barbara Liskov. A decentralized model for information flow control. In SOSP, pp. 129-142. ACM, 1997.
Tai Nguyen and Eric Wong. In-context example selection with influences. arXiv preprint 1104, 2023.
OpenAL. GPT-4 technical report, 2023.
Verónica Pérez-Rosas, Bennett Kleinberg, Alexandra Lefevre, and Rada Mihalcea. Automatic detection of fake news. In 27th International Conference on Computational Linguistics (COLING), pp. 3391 3401. ACL, 2018.
Willard Rafnsson and Andrei Sabelfeld. Secure multi-exccution: Fine-grained, declasification-awarr, and transparent. J. Compul. Secur., 24(1):39 90, 2016.
Andrei Sabelfeld and Andrew C Myers. Language-bard based information-foom security. IEEE Journal on selected areas in communications,, 21(1):5-11 19, 2003.
Timo Schick, Jane Dwivedi-Yu, Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Luke Zeltlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools, 2023.
Lloyd Shapley. A value for N-person games. Contribution to the Theory of Games, 2, 1953.
Xinyue Shen, Zeyuan Chen, Michael Backes, Yun Shen, and Yang Zhang. 'Do Anything Now": Characterizing and evaluating in-the-wild jailbreak prompts on large language models, 2023.
Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed Chi, Nathanael Schärli, and Denny Zhou. Large language models can be casily distracted by irrelevant context. In International Conference on Machine Learning, pp. 31210 31227. PMLR, 2023.
Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, Andrew Poulton, Viktor Kerkez, and Robert Stojnic. Galactica: A large language model for science. arXiv preprint arX0o:2211.09005, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blocher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan
21

---

Published in Transactions on Machine Learning Research (10/2025)
Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Tlyaan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Miles Turpin, Julian Michael, Ethan Perez, and Samuel Bowman. Language models don't don't always say what they think: unfaithful explanations in chain-of-thonght prompting. Advances in Neural Information Processing Systems, 36, 2024.
Apury Verma, Satyapriya Krishna, Sebastian Gehrmann, Madhavan Seshadri, Anu Pradhan, Tom Ault, Leslie Barrett, David Rabinowitz, John Doucette, and NhatHai Phan. Operationalizing threat model for red-tcaming large language models (Ilms). arXiv preprint arXiv:2407.14997 2024.
Philipp Vogt, Florian Nentwich, Nenad Jovanovic, Engin Kirda, Christopher Kruegel, and Giovanni Vigna. Cross site scripting prevention with dynamic data tainting and static analysis. In NDSS, volume 2007, pp. 12, 2007.
Dennis Volpano, Cynthia Irvine, and Geoffrey Smith. A sound type system for secure flow analysis. Journal of computer security, ((2-3))167-187, 1996.
Eric Wallace, Kai Xiao, Reimar Leike, Lilian Weng, Johannes Heidecke, and Alex Beutel. The instruction hierarchy: Training Iims to prioritize privileged instructions. arXiv preprint arXiv:2404. 18208 fes. CR], 2024. doi: 10.48550/ Xiv.2404.1322
Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilitics of large language models, 2022.
Fangzhou Wu, Ethan Cecchetti, and Chaowei Xiao. System-level defense against indirect prompt injection 19091
attacks: An information flow control perspective, 2024. URL httpp:://axxvv.rrggpppp
Lukas Wutschitz, Boris Köpf, Andrew Paverd, Saravan Rajmohan, Ahmed Salem, Shruti Tople, Santiago Zanella-Béguelin, Menglin Xia, and Victor Rühle. Rethinking privacy in machine learning pipelines from an information flow control perspective. arXiv preprint Xrxi:22311. 15792 [cs.LG], 2023. doi:
00.48550/arXiv.3311 2311.1579..
Zhiheng Xi, Wenxiang Chen, Xin Guo, Wei He, Yiwen Ding, Boyang Hong, Ming Zhang, Junzhe Wang, Senjie Jin, Fnyu Zhou, et al. The rise and potential of large language model based agents: A survey. arXiv preprint arXio:2309.072: 078604 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint 2xxc:2309.11455 2023a.
Shitao Xiao, Zheng Liu, Peitian Zhang, and Niklas Muennighoff. C-Pack: Packaged resources to advance general chinese embodding, 2023b.
Youliang Yuan, Wenxiang Jiao, Wenxuan Wang, Jen tse Huang, Pinjia He, Shuming Shi, and Zhaopeng Tu. GPT-4 is too smart to be safe: Stealthy chat with LLMs via cipher, 2023.
Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Josoph E. Gonzalez, al. SGLang: Efficient exccution of structured language model programs, 2023.
Poter Yong Zhong, Siyuan Chen, Ruiqi Wang, McKenna McCall, Ben L. Titzer, Heather Miller, and Phillip B. Gibbons. RTBAS: Defending LLM Agents Against Prompt Injection and Privacy Leakage, 2025. URL https://arxiv.orr//bb
22

---

Published in Transactions on Machine Learning Research (10/2025)
Andy Zou, Zifan Wang, J. Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks on aligned language models, 2023.
23

---

Published in Transactions on Machine Learning Research (10/2025)
best exact match accuracy 85.94%
best exact match
85.94%
0.00
1.96
1.56
4.69
29.69
$1.99
76.56
79.69
46.88
23.44
17.19
4.69
0.0
0.80
0.00
3.12
20:11
51.12
67.14
76.12
70.12
46.03
21.08
18.75
4.69
0.001
3.12
4.69
25
71.88
Shapler thrested 0,003
0.003
Exoct matthh accuracy (96)
thresheld
0.00
4.89
4.69
10.94
32.01
62.50
76.06
81.25
48.44
23.44
17.19
4.69
0.00 0.01
0.00
17.19
$4.38
57.81
73.44
79.69
81.25
51.56
21.88
15.75
0.00
21.88
23:44
31.25
66.08
68
01.25
85:94
50.00
25.00
38.75
4.05
0.00
29.69
37.50
48.44
$4.06
75.00
85:94
84.33
56.25
26.12
23.44
7.81
0.
57.81
0.1
57.83
71.88
78:12
68.44
48
0.3
43.75
56
23.44
21.89
20
0.00
12.50
12.50
12.50
12.50
12.50
12.50
12:50
12.50
12.50
12.50
12.50
1.0
0.00
32.50
12.50
12.50
12.50
12.50
12.50
12.50
12.50
12.50
12.50
12.50
0,001
0.003
0.01
0.025
0.05
0.25
0.001
0.003
0.05 0.075
0.25
0.5
Perplexity tolerance
Perplexity telerance
Figure Hyperparameter grid search on Shapley threshold and perplexity tolerance for the 7B (left) and 70B (right) models the synthetie key-value dataset.
accuracy
best match accuracy
0,0
45.25
49.15
71.40
70.53
0.0
0.01 0,003 0,001
0.03 0.01 0.003 0.001
0.00
44.85
45.25
46.41
49.15
53.71
57.78
60.52
71.40
74.69
72.73
70.53
0.00
54.11
$4.45
58.16
29
63.06
66.28
67.91
73.45
72.64
53.09
64.36
80
Beact matth accuraccy (96)
0.00
44.85
45.25
46.41
49.15
53.71
57.78
60.52
71.40
74.69
72.73
70.53
0.00
54.11
54.45
58.16
59.26
63.06
66.38
67.91
73.45
72.64
69.09
64.35
Shaplo threshold
Shapley threshold
60
54.45
58:16
67.91
73.45
72.64
72.64
0.00
44.85
45.25
46.41
45.15
53
57.78
60:52
71.40
74.69
72.73
70.53
0.00
54.11
54.06
56.16
59:28
63.00
66.38
67.91
73.45
72.64
68.09
64.36
0.3
45.25
16.41
69.15
60.52
71.40
74.69
72.73
70.51
0.3
0.00
54.11
54.46
58.16
59.28
63.06
66.38
67.91
73.45
72.64
68.09
64.36
0.001
0.003
0.025
0.05
0.075
0.5
0.01
0.65
0.75
Perplexity tolerance
Perplexity tolerance
accuracy
accuracy
40.14
10.41
41.65
44.18
51.08
53.3 14 41.15
61.78
62.66
59.57
0.0
0.001
0.03 0.1 0.003 0 001
0.00
40.14
40.43
41.65
44.18
47.59
51.08
53.35
61.15
63.78
62.66
59.17
0.00
45.54
45.97
67.16
50.10
54.63
57
40
59.42
64.50
62.91
59.19
54.00
80
Bxodt matthh sccouralcy/ (56)
0.00
60.14
40.61
41.65
44.18
47.
51
63.78
59.57
0.00
45.54
45.97
47.16
50.10
54.63
57
$7.40
59.42
64.50
62.91
58.19
$4.60
Shaples threshold
Shapley thresholdd
60
60.14
40.43
64.18
45
59.57
45.54
45.97
40
0.00
A0.14
80.43
41.85
44.15
47.39
51.00
$3.3
41.15
63.78
62.65
59.37
0.00
45.54
45.97
67.16
50.10
56.61
57.40
54.42
64.50
62.91
59.19
54.60
0.3
40.43
41.65
47.59
51.08
41.15
63.78
5937
0.3
0.00
45.54
45.97
67.16
50.10
54.63
57
40
59.02
64.50
62.91
59.19
54.00
1.0
0.05
0.075
0.075
Perplesity
Perplexity
Figure 10: Hyperparameter grid search on Shapley threshold and perplexity tolerance for the 7B (left) and 70B (right) model on the news article dataset with a perfect retriever (first row) and realistic retriever (second row).
A Improving performance on fine- line-erained lattices
Algorithm D relies only on À produce the minimal of labels. However, in the cases where documents share similarities and provide important information regarding the model output such the format. (which
24

---

Published in Transactions on Machine Learning Research (10/2025)
is the case for synthetic key-value dataset 4.5 the overall contribution of irrelevant documents goes up. Therefore, additional irrelevant documents are introduced in the predicted label set due to being X-similar.
In order to cover these cases when dealing with a complex lattice and large number of relations, we propose a simple heuristic based on Shapley values that have been commonly used in the past [Shalley (1953): Nguyen & Wong )002)) In particular, we compute the Shapley value for each of the different labels the lattice and filter out any label combinations where labels below a particular threshold are present.
The Shapley value defines the marginal contribution of any label L by computing the average difference in outcomes when a particular label is present and absent. We compute the label Shapley value via perplexity, which is computed as the average drop in perplexity when a particular label is included. This provides a notion of the importance of each label. This now rntrues an additional hyperparameter i.e., the Shapley value threshold.
We visualize the results for the 2D grid search by considering both perplexity tolerance À as well as Shapley threshold the synthetic key-value dataset. in Figure g and the news article dataset. in Figure T0
The results indicate that simpler lattices such as in the case of news articles, we see almost no impact due to the Shapley value threshold. On the other hand, for complex lattices, such as in the case of our synthetic key-value dataset, we observe 6% absolute improvement in the exact-match accuracy, highlighting the utility of this heuristic in such cases.
25


Reviews:
[{"review_id": "Pylkh2Ux1N", "reviewer": "Reviewer_SC3C", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths:\n- The approach is simple and the paper discusses both correctness and the computational cost.\n- The technical part of the paper is well written and formalized to a reasonable extent.\n- The paper attempts to solve a new problem and thus curates various synthetic and manually generated datasets for the problem. \n\n## Weakness:\n\n**Writing** (Mainly in abstract and intro):\n\nOverall the abstract and introduction are unclear and it is hard to infer the motivation and contribution. Since this is a novel problem, I encourage the authors to significantly improve the writing of the paper to convey the primary research contribution.\n\nIn intro:\n>Labels are usually propagated conservatively: the output of an operation on data is labeled as the most restrictive (i.e., most confidential or least trusted) label of its inputs.\n\nWhat work does this? There is no citation to this approach.\n\n> 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.  The line in abstract does not read well. “Labels” are not defined at this point. \n\n\n> 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.\n\n“Samples” are not defined at this point\n\n> improves the label in at least 50% of the cases on all datasets, thus mitigating the problem of label creep.\n\nUnclear what improvement means here?\n\n> We address this issue with the observation that it is sufficient to identify the labels that improve over the full context’s label, and consider only their corresponding subcontexts.\n\nThis line is not clear\n\n - The citations everywhere are not formatted correctly. Use \\citep and \\citet appropriately.\n\n**Evaluation:**\n\n- Why is only Llama-2 used for the evaluation? The authors should consider using more recent Llama-3 and more recent open-weight models. \n\n- Dataset complexity: News article dataset only uses pairs of datasets which is High integrity and Low integrity. The QAs for the dataset were generated using GPT-4 and were not manually validated.  LLM agent dataset is manually curated but it also only has two labels. The proposed technique should be evaluated on a lattice that is more complex than just 2 classes.\n\n- Figure-8 should also show a line for introspection baseline", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a technique for improved label attribution for LLM generation given multiple sources in the context. The proposed algorithm propagates the labels of the samples that were influential in generating the model output and to eliminate the labels of unnecessary inputs.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1) Rewrite the abstract and the introduction of the paper for giving better motivation and overview of the problem. \n2) Perform experiments with 3-4 more open-source models such as various Qwen and Llama-3 models. \n3) Increase the complexity of the dataset and present individual error analysis for subset of the problems for both the baseline and the proposed technique.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JytlgufbQ9", "reviewer": "Reviewer_pCjM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. The motivation is clear: The problem of label creep in naive label propagation for LLMs is convincingly articulated.\n\n2. The influence-based label propagator is formalized with clear mathematical definitions.\n\n3. The paper evaluates the method on three datasets to test label propagation in practical scenarios (synthetic, news, agent), and in most cases the method is effective.\n\n**Weaknesses**\n\n1. The method assumes all context documents are pre-labeled with correct integrity/confidentiality. In many real-world LLM workflows (such as web search, multi-agent collaboration, tool outputs), this assumption can break. \n\n2. The introspection baseline is a bit weak for smaller models, which may make the improvement margins appear large.\n\n3. The datasets, while thoughtfully constructed, are mostly synthetic or manipulated, thus, generalizability to more diverse, noisy, or large-scale real-world deployments is not verified.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper addresses the problem of information-flow label propagation in LLM-based systems, where LLMs ingest and generate outputs conditioned on data of varying integrity. The key contribution is the proposal of a more permissive, influence-based label propagation algorithm that propagates only the labels of context documents actually required to produce the output, mitigating label creep from overly conservative standard approaches.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. It would be better to discuss how to tune $\\lambda$ given its significant sensitivity.\n\n2. It would be better to include experiments and analysis on mis-labeled situations.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "ybI10MAAFp", "reviewer": "Reviewer_mg3P", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- Clear conceptual framework:\n\t- Clearly defines λ-similarity and minimal influence labels, ensuring practical interpretability.\n\t- Formal correctness and safety guarantees enhance trustworthiness.\n- Practical applicability:\n\t- Directly addresses \"label creep,\" an important real-world security and privacy challenge.\n\t- Enables practical deployment scenarios, particularly where data sources have varying trust or confidentiality levels.\n- Comprehensive empirical evaluation:\n\t- Novel setting, robust experiments covering multiple scenarios and lattice complexities.\n\t- Demonstrates consistent performance improvements over baselines, particularly introspection methods.\n- Algorithmic efficiency:\n\t- Provides efficient heuristic methods (e.g., Shapley pruning) for large search spaces, maintaining performance without compromising safety.\n- Security focus:\n\t- Effectively addresses indirect prompt-injection attacks by limiting label propagation to necessary contexts only.\n\n\n**Weaknesses**\n\n- Deployment complexity and cost:\n\t- Reliance on repeated model execution (especially in commercial APIs) might increase costs and latency significantly.\n- Dependence on accurate labeling:\n\t- The approach assumes accurate and complete labeling of all retrieved documents, which might not always be realistic in practice.\n- Limited insights on kNN-LM performance:\n\t- Notably weaker results in kNN-LM implementation lack detailed analysis and explanation.\n- Hyperparameter sensitivity:\n\t- Narrow optimal range for the hyperparameter λ; insufficient guidance on practical setting strategies.\n- Potential denial-of-service vulnerabilities:\n\t- Acknowledged but not thoroughly explored possibility of denial-of-service via injection of low-utility, low-integrity context.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces influence-based label propagation (LP) for retrieval-augmented LLM systems. Instead of conservatively assigning the join of all input labels to a model’s output, the proposed LP only propagates the labels of influential context documents, determined by bounding the drop in utility (perplexity) within λ. The authors study two concrete instantiations: (i) prompt-based retrieval augmentation in which retrieved passages are concatenated into the prompt (typical RAG setup), and (ii) kNN-LM, where an external datastore contributes a mixture distribution to the next-token probabilities. Across three datasets (synthetic key-value, news, and an LLM-agent tool-use corpus) the prompt-based LP achieves exact-match accuracies up to 86% and improves or preserves the output label in > 85% of agent scenarios 2025. The paper provides a formal problem statement, correctness guarantees under monotonicity, and an algorithm whose search complexity is bounded by the lattice size.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Connect to existing threat-model frameworks:\n\t- Please discuss and explicitly cite Operationalizing a Threat Model for Red-Teaming Large Language Models (LLMs) by Verma et al., 2024 especially in the context of Infusion Attacks presented in their work, exploring integration possibilities with structured adversarial capabilities and evaluation frameworks presented therein. \n\t\t- https://arxiv.org/pdf/2407.14937\n- Provide practical guidance on setting λ:\n\t- Recommend an analytic approach or adaptive heuristic for setting λ in real-world deployments to reduce the tuning burden.\n- Clarify and analyze kNN-LM shortcomings:\n\t- Conduct and include an ablation study on the γ hyperparameter and provide qualitative insights on factors contributing to poor kNN-LM performance.\n- Include cost and latency evaluations:\n\t- Add concrete runtime metrics (wall-clock time, token count) comparisons between LP, introspection, and baseline methods to quantify practical trade-offs.\n- Evaluate robustness to incorrect labeling:\n\t- Include experiments quantifying the impact on accuracy and safety guarantees when document labels are partially incorrect or missing.\n- Clarify fallback behavior:\n\t- Explicitly describe the system's response strategy when no sub-context meets the λ-similarity condition to clearly communicate the safety fallback mechanisms.\nOverall, the paper presents valuable insights and methods addressing key security and usability concerns for LLM deployments. Addressing the suggested revisions will strengthen its practical applicability, robustness, and connection to broader security research.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)