# Cascade: Token-Sharded Private LLM Inference
> Project lead and corresponding author.
Abstract
As LLMs continue to increase in parameter size, the computational resources required to run them are available to fewer parties. Therefore, third-party inference services – where LLMs are hosted by third parties with significant computational resources – are becoming increasingly popular. However, third party inference raises critical concerns about user data privacy. To mitigate these risks, privacy researchers have developed provably secure schemes for third-party inference, such as Secure Multi-Party Computation (SMPC). However, SMPC protocols have significant computational and communication overhead, and do not scale to large models. In this work, we propose a new multi-party inference protocol, Cascade, that avoids these punitive costs by leveraging sharding in the sequence dimension to maintain privacy, trading off cryptographic privacy guarantees for increased performance and scalability. We demonstrate that Cascade is resistant to a generalization of a recent attack that is highly effective against other statistical privacy schemes, and that it is further resistant to learning-based attacks. As Cascade is orders of magnitude faster than existing schemes, our findings offer practical solutions for secure deployment of modern state-of-the-art LLMs.
1 Introduction
Due to recent advances in their abilities, Large Language Models (LLMs) have been increasingly used for a wide range of tasks, including but not limited to code generation, language processing, and complex reasoning (Zhu et al., 2024; Kasneci et al., 2023; Thirunavukarasu et al., 2023). However, modern LLMs often contain hundreds of billions of parameters and require significant hardware resources to deploy. In particular, recent open-weights models achieve cutting-edge performance (DeepSeek-AI et al., 2025; Qwen et al., 2025) but remain too costly for individuals or organizations to run on their own. Thus, the demand for third-party LLM inference providers has grown significantly. This raises critical privacy risks when data confidentiality is imperative, such as in areas with strict data privacy laws (e.g. GDPR in Europe), or in domains where sensitive individual information is revealed (e.g. healthcare).
As a result, privacy researchers have focused on creating performant protocols for LLM inference, which do not reveal underlying information about the input prompt to the third-party provider. Many of these schemes rely on an approach called Secure Multi-Party Computation (SMPC), which involves multiple parties jointly performing inference, in such a way that no party can reconstruct the input from the information it receives (Yao, 1982; Goldreich et al., 1987). Recently, various SMPC schemes have been formulated for for LLMs (Huang et al., 2022; Hao et al., 2022; Pang et al., 2023; Akimoto et al., 2023; Dong et al., 2023; Li et al., 2024), which rely on protocols based on additive-secret-sharing. These protocols are efficient for the bulk of LLM operations, such as matrix multiplication, but remain bottlenecked at non-linearities. Many works attempt to bypass this limitation by using SMPC-friendly approximations for non-linearities, such as piecewise polynomials. However, such approximations can impact downstream performance, and still remain far more expensive than direct computation of non-linearities.
Therefore, recent works have sought to mitigate non-linearity costs by using statistical obfuscation approaches rather than cryptographic primitives. These schemes are significantly faster than SMPC schemes, but they can only guarantee statistical security, not theoretical security. For example, recent schemes (Zheng et al., 2024; Yuan et al., 2024; Luo et al., 2024a) leveraged the permutation-equivariance properties of transformers (Xu et al., 2024) to obtain permutation-based schemes for private inference. Here, hidden states are revealed as permuted plaintext to the party performing the inference. These works justify security by referring to the extremely large permutation space, and concluding that the reversal of these permuted states to the original user prompts is infeasible.
However, this argument has recently been shown to be flawed. In recent work Thomas et al. (2025) devise a reconstruction attack, the vocab-matching attack, that they show is able to near-perfectly decode input tokens from LLM hidden states. The authors extend the attack to further decode input tokens from various permutations of hidden states, breaking the security of three statistical permutation-based schemes (Zheng et al., 2024; Yuan et al., 2024; Luo et al., 2024a).
In this paper, we introduce a new multi-party scheme, Cascade, that is resistant to the vocab-matching attack (see Figure 1). Although this is a statistical scheme, unlike previous works, it leverages sharding at the token level for obfuscation. While Cascade does not have the rigorous privacy guarantees of cryptographic MPC schemes, it is much more efficient than them – up to 100x faster overall, and with over 150x lower communication costs – and presents a new paradigm in the trade-off between scalability and security.
The main contributions of our paper are:
1. We generalize the vocab-matching attack introduced in Thomas et al. (2025) to sharded hidden state reversal. While the vanilla vocab-matching attack decodes full hidden states, our generalization is able to decode arbitrary slices of $N× d$ hidden states along the token dimension $N$ .
1. We introduce Cascade, a multi-party statistical scheme for inference in the open-weights setting that relies on sharding at the token level. We demonstrate that Cascade is secure against our generalized vocab-matching attack for certain choices of sharding along the token dimension. We also show that Cascade is resistant to existing learning-based reversal approaches in the literature (Wan et al., 2024; Morris et al., 2023).
1. We demonstrate the efficiency and scalability of Cascade. As our scheme does not employ cryptographic primitives like additive-secret-sharing or homomorphic encryption, it is significantly more lightweight than existing schemes, and is able to scale to modern state-of-the-art LLMs.
<details>
<summary>extracted/6603058/images/cascade-final.png Details</summary>
### Visual Description
## Diagram: Attention Mechanism with Sharding
### Overview
The image is a diagram illustrating an attention mechanism, likely within a neural network architecture. It shows the flow of data and operations between different components, including token indices, computation nodes (CompNodes), and attention nodes (AttnNodes). The diagram also indicates a sharding operation.
### Components/Axes
* **Title:** None explicitly given, but the diagram depicts an attention mechanism.
* **Labels:**
* "β x β" and "AttnNodes" at the top, indicating the attention nodes and their dimensionality.
* "α" and "CompNodes" in the middle-left, indicating the computation nodes.
* "Token indices: [1, 2, 3, ..., N]" at the bottom, representing the input tokens.
* "User text: 'Computer viruses are parasitic programs which are...'" below the token indices.
* "R-sharding" at the bottom-right.
* Within each CompNode block: "h → q, k, v", "Layernorm", "O-proj", "MLP".
* Within the AttnNodes block: "softmax(qkᵀ + m) . v", "expsum(qkᵀ + m)".
* **Nodes:**
* Three CompNodes, each containing "h → q, k, v", "Layernorm", "O-proj", and "MLP".
* One AttnNodes block containing "softmax(qkᵀ + m) . v" and "expsum(qkᵀ + m)".
* **Connections:** Arrows indicate the flow of data between the components.
* **Index Sets:**
* CompNodes: "[1, 2, 8, 9]", "[3, 4, 10, 11]"
* AttnNodes: "[1, 3, 8, 10]", "[2, 4, 9, 11]"
### Detailed Analysis
1. **Token Indices:** The diagram starts with "Token indices: [1, 2, 3, ..., N]" and a user text snippet: "Computer viruses are parasitic programs which are...". This suggests the input is a sequence of tokens.
2. **CompNodes (Computation Nodes):**
* There are three CompNodes, each performing the same operations:
* "h → q, k, v": This likely represents the transformation of input embeddings (h) into query (q), key (k), and value (v) vectors.
* "Layernorm": Layer normalization.
* "O-proj": Output projection.
* "MLP": Multi-layer perceptron.
* The index sets associated with the first two CompNodes are "[1, 2, 8, 9]" and "[3, 4, 10, 11]". The third CompNode has no explicit index set listed.
3. **AttnNodes (Attention Nodes):**
* The AttnNodes block performs the core attention calculations:
* "softmax(qkᵀ + m) . v": This calculates the attention weights using a softmax function applied to the dot product of query (q) and key (k) plus a mask (m), and then multiplies the result by the value (v).
* "expsum(qkᵀ + m)": This calculates the exponential sum of the dot product of query (q) and key (k) plus a mask (m).
* The index sets associated with the AttnNodes are "[1, 3, 8, 10]" and "[2, 4, 9, 11]".
4. **Data Flow:**
* Arrows indicate the flow of data from the token indices to the CompNodes.
* Arrows also show data flowing from the CompNodes to the AttnNodes.
* There are connections between the AttnNodes and the CompNodes.
5. **Sharding:** "R-sharding" at the bottom-right suggests that the data or computation is being sharded (split) across multiple devices or processors.
### Key Observations
* The diagram illustrates a standard attention mechanism, with query, key, and value vectors being computed and used to calculate attention weights.
* The presence of "Layernorm" and "MLP" indicates that this is likely part of a larger neural network architecture.
* The index sets associated with the nodes suggest that the tokens are being processed in specific groups or shards.
* The "R-sharding" label indicates that the computation is distributed.
### Interpretation
The diagram depicts a sharded attention mechanism within a neural network. The input tokens are processed by computation nodes (CompNodes) to generate query, key, and value vectors. These vectors are then used in the attention nodes (AttnNodes) to calculate attention weights and produce a weighted sum of the value vectors. The sharding operation suggests that the computation is distributed across multiple devices, which is common in large-scale neural network training. The index sets likely represent the specific tokens assigned to each shard. The "softmax(qkᵀ + m) . v" and "expsum(qkᵀ + m)" operations are central to the attention mechanism, allowing the model to focus on the most relevant parts of the input sequence.
</details>
Figure 1: Schematic representation of our proposed privacy-preserving multi-party inference scheme, Cascade.
2 Setup & Threat Model
We consider the setting where a user $U$ wants to perform inference with an LLM model $M$ on some input prompt $\bm{x}$ , which can be considered as an ordered sequence of tokens $[x_{1},x_{2},...,x_{N}]$ . We denote the size of the hidden state of the LLM by $d$ , and the sequence length by $N$ .
As the user $U$ does not have the resources to perform the inference themselves, they rely on a set of third-parties $P_{1},P_{2},...,P_{K}$ . The model weights of $M$ are known to all parties. We consider the setting where each of the parties behaves semi-honestly, a common assumption of past works (Zheng et al., 2024; Luo et al., 2024a; Dong et al., 2023; Yuan et al., 2024). Semi-honest parties will follow the defined protocol faithfully, but may exploit any information that they receive during the execution of the protocol to attempt to recover the user’s data.
Although it is not necessary for the setup and description of Cascade, we further assume that the LLM model $M$ employs unidirectional self-attention. This can be considered a worst-case assumption in security analysis, as the vocab-matching attack relies on this property to succeed.
3 Related Work
Recent cryptographically-based SMPC schemes aim to alleviate the computational overhead of computing non-linearities in transformer architectures by devising new protocols for SMPC-unfriendly operations. In MPCFormer, Li et al. (2023) replace GeLU by a quadratic approximation and softmax with a normalized quadratic approximation. They employ knowledge distillation to fine-tune the model for downstream tasks, but they still see significant degradation in performance. To avoid the expensive operation of fine-tuning, Puma (Dong et al., 2023) formulates the first SMPC protocol for GeLU using piecewise polynomial fitting. However, this protocol still requires many multiplication and inequality protocol calls, and it is expensive. To improve the efficiency of GeLU, SecFormer (Luo et al., 2024b) replaces GeLU with a Fourier series approximation using segmented polynomials, and incorporates knowledge distillation as in MPCFormer. They are able to achieve similar runtimes as MPCFormer, with better downstream performance.
Most existing statistical schemes that are related to our work are permutation-based. Zheng et al. (2024) introduce the two-party scheme PermLLM, where hidden states are permuted at non-linear components to ‘safely‘ reveal elements to parties. Yuan et al. (2024) also develop STIP, a three-party scheme where the model server carries out inference on the user’s input, without learning the proprietary model weights belonging to the model provider. This is accomplished with random permutations of both hidden states and model weights. In Luo et al. (2024a), the Cenatur scheme follows the same three-party model. However, they now avoid exposing the embedding lookup table and certain unpermuted intermediate results (e.g. attention logits), by incorporating additive-secret-sharing from SMPC at most stages of self-attention. In all of these schemes, permutations of hidden states are exposed to a party performing inference. Thomas et al. (2025) use this observation to break the security of these schemes with their vocab-matching attack.
Later, to analyze the security of our statistical scheme, we will test reconstruction quality of learning-based attacks on exposed shards. The most relevant attacks in our setting are Wan et al. (2024) and Morris et al. (2023). The former work focuses on reversal of hidden states, while the latter emphasizes logit distribution reversal. In both papers, the authors fine-tune a transformer-based architecture to reverse hidden states into their original input tokens. The latter paper does not assume any access by the adversary to model weights, while the former explicitly denotes the case of a model provider performing inference on user provided embeddings, and so is more analogous to our setting.
We will also test the security of our scheme against a generalization of the attack from Thomas et al. (2025), which is able to reverse input tokens from sharded hidden states. We describe this in greater detail in the next section.
4 Generalized Vocab-Matching Attack
Here, we describe the vanilla attack from Thomas et al. (2025), where the adversary receives full layer $L$ hidden states $h_{1},...,h_{N}∈\mathbb{R}^{d}$ and attempts to recover the $N$ -token input prompt. Then, we introduce a variant where the adversary only receives some of these $N$ hidden states but can perform a search over multiple unknown tokens.
4.1 Vanilla Attack
Suppose an LLM inference provider hosts a unidirectional LLM like Gemma-2-2B-IT and is given the following prompt by a user: "What is currently the most populated city in Spain?". After tokenization, this might be represented as the list [BOS, "What", "is", "currently", "the", "most", "populated", "city", "in", "Spain", "?", EOS], with BOS and EOS denoting the special tokens of begin-of-sequence and end-of-sequence respectively. These $12$ tokens have layer $0 0$ embeddings $e_{1},...,e_{12}∈\mathbb{R}^{d}$ , and these give layer $L$ embeddings $h_{1},...,h_{12}∈\mathbb{R}^{d}$ . Denoting the stacked decoder layers in Gemma-2-2B-IT as $\phi_{1},...,\phi_{f}$ , which are all unidirectional, one sees for $\phi_{≤ L}=\phi_{L}\circ...\circ\phi_{2}\circ\phi_{1}$ that
$$
\phi_{\leq L}(e_{1})=[h_{1}]\in\mathbb{R}^{1\times d},\;\phi_{\leq L}(e_{1},e_%
{2})=[h_{1},h_{2}]\in\mathbb{R}^{2\times d},\;\ldots,\;\phi_{\leq L}(e_{1},%
\ldots,e_{12})=[h_{1},\ldots,h_{12}]\in\mathbb{R}^{12\times d}.
$$
First, the provider takes the Gemma-2-2B-IT vocabulary $\mathcal{V}$ with size $V=256000$ . For each token $v∈\mathcal{V}$ with embedding $e_{v}∈\mathbb{R}^{d}$ , they compute $\phi_{≤ L}(e_{v})∈\mathbb{R}^{1× d}$ , and then set $\widehat{e}_{1}$ as the $e_{v}$ which minimizes the L1 error $\|\phi_{≤ L}(e_{v})_{1}-h_{1}\|_{1}$ . Because $\phi_{≤ L}(e_{1})=h_{1}$ , then assuming there are no other collisions $\phi_{≤ L}(e_{v})=h_{1}$ , we will have $\widehat{e}_{1}=e_{1}$ . Next, they compute $\phi_{≤ L}(\widehat{e}_{1},e_{v})∈\mathbb{R}^{2× d}$ for each $v∈\mathcal{V}$ , and set $\widehat{e}_{2}$ as the $e_{v}$ which minimizes $\|\phi_{≤ L}(\widehat{e}_{1},e_{v})_{2}-h_{2}\|_{1}$ . Again, without collisions, we have $\widehat{e}_{2}=e_{2}$ . This continues until finally, they compute $\phi_{≤ L}(\widehat{e}_{1},...,\widehat{e}_{11},e_{v})∈\mathbb{R}^{11%
× d}$ , and set $\widehat{e}_{12}$ as the $e_{v}$ which minimizes $\|\phi_{≤ L}(\widehat{e}_{1},...,\widehat{e}_{11},e_{v})_{12}-h_{12}\|_{1}$ , which will be $e_{12}$ . Now, the provider has true input embeddings $[\widehat{e}_{1},...,\widehat{e}_{12}]=[e_{1},...,e_{12}]$ , which can be reversed into the $12$ input tokens via the lookup table. Here, non-collision is necessary for provable success of the attack, as otherwise we may not have each $\widehat{e}_{i}=e_{i}$ .
4.2 Sharding Generalization
Now, we discuss a generalization of the vanilla attack above, with our psuedocode shown in Algorithm 1. This pertains to settings where the provider does not receive all $N$ hidden states. For instance, suppose they only receive $h_{3},h_{5},h_{10}$ . Can they still recover all input tokens?
The answer is negative. For instance, none of these $3$ hidden states are functions of $e_{11},e_{12}$ , so one cannot always One might be able to infer these from previous tokens for certain prompts. For instance, in the example prompt, ”?” and EOS could be inferred from the rest of the prompt. guarantee recovery of tokens $11$ and $12$ . Nevertheless, our generalized attack, with sufficiently large computation cost, allows recovery of all tokens that the last hidden state $h_{10}$ depends on, meaning the first $10$ tokens.
First, since the provider does not know $h_{1}$ , they cannot carry out the first step of the vanilla attack to decipher $e_{1}$ . However, if they generalize the first step, they can obtain something stronger from $h_{3}$ : all of $e_{1},e_{2},e_{3}$ . This is because $[h_{1},h_{2},h_{3}]=\phi_{≤ L}(e_{1},e_{2},e_{3})$ , so assuming non-collision Actually, since we are matching only the last row to $h_{3}$ , all that is needed is non-collision of the last row. of the $V^{3}$ possible $3$ -token forward passes $\{\phi_{≤ L}(e_{v_{1}},e_{v_{2}},e_{v_{3}})\}_{v_{1},v_{2},v_{3}∈\mathcal%
{V}}$ , the provider can recover $e_{1},e_{2},e_{3}$ in at most $V^{3}$ forward passes. Next, from $[h_{1},...,h_{5}]=\phi_{≤ L}(e_{1},...,e_{5})$ and non-collision of the $V^{2}$ possible $2$ -token forward passes $\{\phi_{≤ L}(e_{1},e_{2},e_{3},e_{v_{4}},e_{v_{5}})\}_{v_{4},v_{5}∈%
\mathcal{V}}$ , the provider can recover $e_{4},e_{5}$ in at most $V^{2}$ forward passes. Finally, from $[h_{1},...,h_{10}]=\phi_{≤ L}(e_{1},...,e_{10})$ and non-collision of the $V^{5}$ possible $5$ -token forward passes $\{\phi_{≤ L}(e_{1},...,e_{5},e_{v_{6}},...,e_{v_{10}})\}_{v_{6},%
...,v_{10}∈\mathcal{V}}$ , the provider can recover $e_{6},...,e_{10}$ in at most $V^{5}$ forward passes.
Now, the provider has recovered the first $10$ tokens in at most $V^{3}+V^{2}+V^{5}$ forward passes, where exponents $3,2,5$ correspond to gaps between indices of known hidden states. In general, if the provider gets $h_{i_{1}},...,h_{i_{k}}$ with $1≤ i_{1}<...<i_{k}≤ N$ , then they can carry out this attack to reverse the first $i_{k}$ tokens in at most $V^{i_{1}}+V^{i_{2}-i_{1}}+...+V^{i_{k}-i_{k-1}}$ forward passes. Although this attack appears powerful, the upper bound on forward passes can be quite large, so it is not always feasible. For instance, the dominating term in $V^{3}+V^{2}+V^{5}$ is $V^{5}≈ 10^{27}$ for Gemma-2-2B-IT, which is entirely infeasible to search over for any adversary. In general, as $k≤ N=o(V)$ in practice, the cost bound $V^{i_{1}}+V^{i_{2}-i_{1}}+...+V^{i_{k}-i_{k-1}}$ is dominated by the maximum-gap term $V^{g}$ with $g=\max_{j}(i_{j+1}-i_{j})$ . Thus, if $g$ is large enough so that the adversary cannot perform $V^{g}$ forward passes, we would expect the attack to be infeasible. We define the minimum value of $g$ satisfying this as the vocab-matching threshold $\rho$ .
**Definition 4.1**
*Let $t_{\text{max}}$ be the maximum value of $t$ for which an adversary has resources to perform $V^{t}$ forward passes. Then we define the vocab-matching threshold $\rho$ as $\rho:=t_{\text{max}}+1$ .*
Note that $\rho$ can be set on a case-by-case basis to suit the security demands of the user. Because $V\sim O(100000)$ in typical modern LLMs, the vocab-matching threshold of any adversary is likely no more than $\rho=3$ in practice.
Finally, we remark that the generalized attack requires stronger assumptions than the vanilla attack to succeed. For instance, in the vanilla attack, we only require non-collision to hold across $V$ forward passes at a time. In our example, however, we require it across up to $V^{5}$ forward passes, and in general, it is required across up to $V^{g}$ forward passes. However, as our generalization of the attack is considered primarily as a theoretical threat, we make the worst-case assumption that non-collision almost always holds in our security analysis.
Algorithm 1 Generalized Vocabulary-Matching Attack
0: Model $M$ , indices $1≤ i_{1}<i_{2}<...<i_{k}≤ N$ , layer $l$ hidden states $[h_{i_{1}},h_{i_{2}},...,h_{i_{k}}]$ , vocabulary $\mathcal{V}$
0: Nearly decoded token sequence $\widehat{\bm{x}}=[\widehat{x}_{1},\widehat{x}_{2},...,\widehat{x}_{i_{k}-1}%
,\widehat{x}_{i_{k}}]$
1: Initialize empty sequence $\bm{\widehat{x}}←[]$
2: $i_{0}← 0$
3: for $j=0$ to $k-1$ do
4: $\text{gap}← i_{j+1}-i_{j}$
5: $\text{min\_dist}←∞$
6: $\text{best\_match}←\text{None}$
7: for $v_{1},v_{2},...,v_{\text{gap}}∈\mathcal{V}$ do
8: $g← M_{≤ l}([\widehat{\bm{x}},v_{1},v_{2},...,v_{\text{gap}}])$ {Forward pass up to layer $l$ }
9: $\text{dist}←\|g-h_{i_{j+1}}\|_{1}$ {Calculate L1 distance}
10: if $\text{dist}<\text{min\_dist}$ then
11: $\text{min\_dist}←\text{dist}$
12: $\text{best\_match}← v$
13: end if
14: end for
15: end for
16: return $\widehat{\bm{x}}$
5 Cascade: Token-Sharded Multi-Party Inference
In Section 4, we reviewed the vocab-matching attack from Thomas et al. (2025) and showed how it could be extended to reverse tokens from token-sharded hidden states. However, we also saw that the cost of the generalized attack scaled exponentially with gaps between the revealed token indices, and so it is likely infeasible for sharded hidden states with sufficiently large token gaps. This motivates our defense to the generalized attack based on token-sharding of hidden states, which leads to a new statistical multi-party inference scheme: Cascade.
Notably, Cascade does not use cryptographic primitives; the computations are nearly unchanged from a standard forward pass, so almost no additional computational overhead is incurred. There is also no degradation of performance, as unlike SMPC schemes, no approximations need to be made to efficiently compute non-linearities. The scheme also does not require any trusted party interaction during inference.
At a high level, Cascade exploits the fact that only the self-attention mechanism in transformers has interaction among tokens in a sequence; for all other parts of the architecture, the tokens are treated like batch dimension elements. Therefore, the CompNodes each receive only a subset of the tokens, and perform the bulk of batch operations in the LLM. Then, self-attention is performed by separate AttnNodes, who each receive $Q,K,V$ -projections of sharded hidden states from the CompNodes. The AttnNodes send partial attention outputs at each layer back to the CompNodes, who can reconstruct shards of the full attention output.
In this section, we formally define the notation and protocols for Cascade, without initializing a specific sharding setup. In Section 6, we will motivate the exact sharding scheme through a careful security analysis.
5.1 Sharding and Notation
In multi-headed attention, we let $H$ and $H_{KV}$ be the attention head and key-value head (for grouped-query attention) counts, $N$ be the token count, $d_{emb}$ be the hidden dimension, and $d$ be the attention hidden dimension. There are three axes of sharding used along the token dimension: (1) the sharding of hidden states $\bm{h}∈\mathbb{R}^{N× d_{emb}}$ , (2) the sharding of query states $\bm{q}∈\mathbb{R}^{H× N× d}$ , and (3) the sharding of key and value states $\bm{k},\bm{v}∈\mathbb{R}^{H_{KV}× N× d}$ . These involve splitting token indices $[N]=\{1,2,...,N\}$ into a union of disjoint subsets $\{R_{i}\}_{i=1}^{\alpha}$ for hiddens, $\{S_{j}\}_{j=1}^{\beta}$ for queries, and $\{T_{k}\}_{k=1}^{\gamma}$ for keys and values, where $\alpha,\beta,\gamma$ are shard counts. We also shard the positional embeddings $\bm{p}∈\mathbb{R}^{N× d_{emb}}$ and the attention mask $\bm{s}∈\mathbb{R}^{H× N× N}$ , which are initialized pre-inference, as well as the masked logits $\bm{a}∈\mathbb{R}^{H× N× N}$ and attention output $\bm{o}∈\mathbb{R}^{N× d_{emb}}$ .
In practice, we can symmetrize $S$ and $T$ sharding. This arises from a pairwise coverage requirement: as $\{S_{j}\}_{j=1}^{\beta}$ and $\{T_{k}\}_{k=1}^{\gamma}$ cover $[N]$ , each index pair $(x,y)∈[N]×[N]$ lies in some $S_{j}× T_{k}$ . Because $S$ and $T$ sharding correspond to query and key/value sharding, this means for all such index pairs $(x,y)$ , there is some AttnNode jk that holds both the $x$ th query row and the $y$ th key row. Observe that a node holding these rows has – in the worst case – the same information as if they hold the $y$ th query and $x$ th key rows. This is because the query and key matrices are linear projections of the same hidden states. Therefore, we can set $S$ and $T$ sharding equal at no loss of security. From here on, we replace $T$ with $S$ and $\gamma$ with $\beta$ .
Our symmetrized notation for sharded intermediate LLM states at any layer is in Table 1. Note that the last six rows rely on tensors that are not directly shards of LLM states, but are derived from them, as in Table 2. Here, max, softmax, and expsum are performed row-wise, with $\text{expsum}(\bm{x}):=\sum_{i}\text{exp}(x_{i})$ . Also, $\bm{e}^{S_{k}}$ only performs an expsum after subtracting row-wise maximums, as in the numerically stable subtract-max form of softmax.
| $\bm{h},\bm{p},\bm{o}$ | $(N,d_{emb})$ | $(R_{i},*)$ | $\bm{h}^{R_{i}},\bm{p}^{R_{i}},\bm{o}^{R_{i}}$ |
| --- | --- | --- | --- |
| $\bm{q},\bm{k},\bm{v}$ | $(H,N,d)$ | $(*,R_{i},*)$ | $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}$ |
| $\bm{q},\bm{k},\bm{v}$ | $(H,N,d)$ | $(*,S_{j},*)$ | $\bm{q}^{S_{j}},\bm{k}^{S_{j}},\bm{v}^{S_{j}}$ |
| $\bm{q},\bm{k},\bm{v}$ | $(H,N,d)$ | $(*,R_{i}\cap S_{j},*)$ | $\bm{q}^{R_{i}\cap S_{j}},\bm{k}^{R_{i}\cap S_{j}},\bm{v}^{R_{i}\cap S_{j}}$ |
| $\bm{a},\bm{s}$ | $(H,N,N)$ | $(*,*,S_{k})$ | $\bm{a}^{S_{k}},\bm{s}^{S_{k}}$ |
| $\bm{a},\bm{s}$ | $(H,N,N)$ | $(*,R_{i},S_{k})$ | $\bm{a}^{R_{i}S_{k}},\bm{s}^{R_{i}S_{k}}$ |
| $\bm{a},\bm{s}$ | $(H,N,N)$ | $(*,S_{j},S_{k})$ | $\bm{a}^{S_{j}S_{k}},\bm{s}^{S_{j}S_{k}}$ |
| $\bm{a},\bm{s}$ | $(H,N,N)$ | $(*,R_{i}\cap S_{j},S_{k})$ | $\bm{a}^{(R_{i}\cap S_{j})S_{k}},\bm{s}^{(R_{i}\cap S_{j})S_{k}}$ |
| $\bm{m}^{S_{k}},\bm{e}^{S_{k}}$ | $(H,N,1)$ | $(*,R_{i},*)$ | $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}}$ |
| $\bm{m}^{S_{k}},\bm{e}^{S_{k}}$ | $(H,N,1)$ | $(*,S_{j},*)$ | $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}}$ |
| $\bm{m}^{S_{k}},\bm{e}^{S_{k}}$ | $(H,N,1)$ | $(*,R_{i}\cap S_{j},*)$ | $\bm{m}^{(R_{i}\cap S_{j})S_{k}},\bm{e}^{(R_{i}\cap S_{j})S_{k}}$ |
| $\bm{u}^{S_{k}}$ | $(H,N,|S_{k}|)$ | $(*,R_{i},*)$ | $\bm{u}^{R_{i}S_{k}}$ |
| $\bm{u}^{S_{k}}$ | $(H,N,|S_{k}|)$ | $(*,S_{j},*)$ | $\bm{u}^{S_{j}S_{k}}$ |
| $\bm{u}^{S_{k}}$ | $(H,N,|S_{k}|)$ | $(*,R_{i}\cap S_{j},*)$ | $\bm{u}^{(R_{i}\cap S_{j})S_{k}}$ |
Table 1: Sharding for intermediate LLM states, where $*$ is a full slice. In slicing notation, $\bm{a}^{(R_{i}\cap S_{j})S_{k}}=\bm{a}[:,R_{i}\cap S_{j},S_{k}]$ and $\bm{m}^{(R_{i}\cap S_{j})S_{k}}=\bm{m}^{S_{k}}[:,R_{i}\cap S_{j},:]$ , with $\bm{m}^{S_{k}}$ as in Table 2.
| $\bm{m}^{S_{k}}$ | $(H,N,1)$ | $\max(\bm{a}^{S_{k}})$ |
| --- | --- | --- |
| $\bm{e}^{S_{k}}$ | $(H,N,1)$ | $\text{expsum}(\bm{a}^{S_{k}}-\bm{m}^{S_{k}})$ |
| $\bm{u}^{S_{k}}$ | $(H,N,|S_{k}|)$ | $\text{softmax}(\bm{a}^{S_{k}})\bm{v}^{S_{k}}$ |
Table 2: Tensors derived from sharded states $\bm{a}^{S_{k}},\bm{v}^{S_{k}}$ in Table 1.
5.2 Single Layer Pipeline
We now describe how Cascade operates for a single LLM layer. There are two types of nodes which hold our defined sharded states: CompNodes and AttnNodes. We initialize $\alpha$ CompNodes and $\beta^{2}$ AttnNodes, indexed as CompNode i and AttnNode jk for all $i∈[\alpha]$ and $j,k∈[\beta]$ . Cascade breaks single layer inference into the pre-pass by CompNodes, attention-pass by AttnNodes, and post-pass by CompNodes. We describe these in Algorithm 3, Algorithm 4, and Algorithm 5 in subsequent subsections. These form a single layer pass in Algorithm 2.
Algorithm 2 Cascade Single Layer Forward Pass
0: $\bm{h}^{R_{i}}$ at layer $l$ from CompNode i for $1≤ i≤\alpha$
0: $\bm{h}^{R_{i}}$ at layer $l+1$ to CompNode i for $1≤ i≤\alpha$
1: for $i=1$ to $\alpha$ : CompNode i do
2: $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}←\textbf{pre\_pass}(\bm{h%
}^{R_{i}})$
3: for $j,k=1$ to $\beta$ : AttnNode jk gets
4: $\bm{q}^{R_{i}S_{j}},\bm{k}^{R_{i}S_{k}},\bm{v}^{R_{i}S_{k}}$
5: end for
6: end for
7: for $j,k=1$ to $\beta$ : AttnNode jk do
8: $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}},\bm{u}^{S_{j}S_{k}}←\textbf{%
attn\_pass}(\bm{q}^{S_{j}},\bm{k}^{S_{k}},\bm{v}^{S_{k}})$
9: for $i=1$ to $\alpha$ : CompNode i gets
10: $\bm{m}^{(R_{i}\cap S_{j})S_{k}},\bm{e}^{(R_{i}\cap S_{j})S_{k}},\bm{u}^{(R_{i}%
\cap S_{j})S_{k}}$
11: end for
12: end for
13: for $i=1$ to $\alpha$ : CompNode i do
14: $\bm{o}^{R_{i}}←\textbf{post\_pass}\{\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}%
S_{k}},\bm{u}^{R_{i}S_{k}}\}_{k=1}^{\beta}$
15: $\bm{h}^{R_{i}}←\textbf{mlp}(\bm{h}^{R_{i}}+\bm{o}^{R_{i}})$ {Residual and MLP}
16: end for
During Cascade layer $l$ inference, each CompNode i begins with $\bm{h}^{R_{i}}$ at layer $l$ , and ends with $\bm{h}^{R_{i}}$ at layer $l+1$ . This output is the input to CompNode i for layer $l+1$ inference, so inference for all layers can follow Algorithm 2.
After the last layer, CompNodes apply the LM head to get $R_{i}$ -sharded logits, and the CompNode with the last token will use this to generate the next token’s embedding The logits can also be returned to the user if desired.. This is sent back to the CompNode with the next token index in $R_{i}$ . Then, Cascade inference repeats to generate the next token. This can be sped up with KV-caching: see Appendix F.
Now, we describe the pre-pass, attention-pass, and post-pass components of Cascade. All shard-specific slicing and concatenation operations are initialized in the node setup. Furthermore, we assume that all AttnNodes wait until all CompNodes finish the pre-pass to do the attention-pass, and all CompNodes wait until all AttnNodes finish the attention-pass to do the post-pass.
5.2.1 Pre-Pass
At layer $l$ , each CompNode i starts with the $R_{i}-$ sharded hidden states $\bm{h}^{R_{i}}$ , and applies layer normalization if necessary. Then, it $Q,K,V$ -projects these to get the $R_{i}-$ sharded query, key and value states $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}$ . CompNode i then applies rotary or positional embedding to $\bm{q}^{R_{i}},\bm{k}^{R_{i}}$ , using sharded positional embeddings $\bm{p}^{R_{i}}$ (the node can generate these upon setup to avoid any communication overhead, since it only depends on its index set $R_{i}$ ), and returns all of $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}$ , as described in Algorithm 3.
Algorithm 3 CompNode i Single Layer Pre-Pass
0: $\bm{h}^{R_{i}},\bm{p}^{R_{i}}$
0: $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}$
1: $\bm{q}^{R_{i}}←\text{q\_proj}(\bm{h}^{R_{i}})$
2: $\bm{k}^{R_{i}}←\text{k\_proj}(\bm{h}^{R_{i}})$
3: $\bm{v}^{R_{i}}←\text{v\_proj}(\bm{h}^{R_{i}})$
4: $\bm{q}^{R_{i}}←\text{rotary\_pos\_emb}(\bm{q}^{R_{i}},\bm{p}^{R_{i}})$
5: $\bm{k}^{R_{i}}←\text{rotary\_pos\_emb}(\bm{k}^{R_{i}},\bm{p}^{R_{i}})$
6: return $\bm{q}^{R_{i}},\bm{k}^{R_{i}},\bm{v}^{R_{i}}$
5.2.2 Attention-Pass
After the pre-pass, each AttnNode jk receives shards $\bm{q}^{R_{i}S_{j}},\bm{k}^{R_{i}S_{k}},\bm{v}^{R_{i}S_{k}}$ from CompNode i. By concatenating their rows over $1≤ i≤\alpha$ , AttnNode jk obtains $\bm{q}^{S_{j}},\bm{k}^{S_{k}},\bm{v}^{S_{k}}$ . For example, for $\bm{q}^{S_{j}}$ , concatenation is in the order in which one concatenates elements of sorted $R_{i}\cap S_{j}$ over $1≤ i≤\alpha$ to get sorted $S_{j}$ . Then, AttnNode jk computes $\bm{a}^{S_{j}S_{k}}=\bm{q}^{S_{j}}(\bm{k}^{S_{k}})^{T}+\bm{s}^{S_{j}S_{k}}$ , where matrix multiplication is per-head and we broadcast $H_{KV}$ to $H$ . For the post-pass, AttnNode jk also stores row-wise maximums and subtract-max expsums $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}}$ . Finally, AttnNode jk takes the row-wise softmax and performs value multiplication to get $\bm{u}^{S_{j}S_{k}}=\text{softmax}(\bm{a}^{S_{j}S_{k}})\bm{v}^{S_{k}}$ . All of $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}},\bm{u}^{S_{j}S_{k}}$ are returned, as in Algorithm 4.
Algorithm 4 AttnNode jk Single Layer Attention-Pass
0: $\bm{q}^{S_{j}},\bm{k}^{S_{k}},\bm{v}^{S_{k}},\bm{s}^{S_{j}S_{k}}$
0: $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}},\bm{u}^{S_{j}S_{k}}$
1: $\bm{k}^{S_{k}}←\text{repeat\_kv}(\bm{k}^{S_{k}})$
2: $\bm{v}^{S_{k}}←\text{repeat\_kv}(\bm{v}^{S_{k}})$
3: $\bm{a}^{S_{j}S_{k}}←\bm{q}^{S_{j}}(\bm{k}^{S_{k}})^{T}+\bm{s}^{S_{j}S%
_{k}}$
4: $\bm{m}^{S_{j}S_{k}}←\text{row\_max}(\bm{a}^{S_{j}S_{k}})$
5: $\bm{a}^{S_{j}S_{k}}←\text{exp}(\bm{a}^{S_{j}S_{k}}-\bm{m}^{S_{j}S_{k}})$
6: $\bm{e}^{S_{j}S_{k}}←\text{row\_sum}(\bm{a}^{S_{j}S_{k}})$
7: $\bm{a}^{S_{j}S_{k}}←\bm{a}^{S_{j}S_{k}}/\bm{e}^{S_{j}S_{k}}$
8: $\bm{u}^{S_{j}S_{k}}←\bm{a}^{S_{j}S_{k}}\bm{v}^{S_{k}}$
9: return $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}},\bm{u}^{S_{j}S_{k}}$
5.2.3 Post-Pass
Finally, after the attention-pass, each CompNode i receives $\bm{m}^{(R_{i}\cap S_{j})S_{k}},\bm{e}^{(R_{i}\cap S_{j})S_{k}},\bm{u}^{(R_{i}%
\cap S_{j})S_{k}}$ from each AttnNode jk, which are slices of its attention-pass outputs $\bm{m}^{S_{j}S_{k}},\bm{e}^{S_{j}S_{k}},\bm{u}^{S_{j}S_{k}}$ along the second-to-last dimensions. Then, for each fixed $1≤ k≤\beta$ , CompNode i concatenates the rows of $\bm{m}^{(R_{i}\cap S_{j})S_{k}},\bm{e}^{(R_{i}\cap S_{j})S_{k}},\bm{u}^{(R_{i}%
\cap S_{j})S_{k}}$ over all $1≤ j≤\beta$ to obtain $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u}^{R_{i}S_{k}}$ . Next, CompNode i aims to combine these results $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u}^{R_{i}S_{k}}$ over $1≤ k≤\beta$ into the $R_{i}-$ sharded (pre $O$ -proj) output of attention. Using slicing notation, treating matrix multiplication as per-head, and broadcasting $H_{KV}$ to $H$ , this is
$$
\text{softmax}(\bm{a})[:,R_{i},:]\bm{v}=\sum_{k=1}^{\beta}\text{softmax}(\bm{a%
})[:,R_{i},S_{k}]\bm{v}[:,:,S_{k}]\in\mathbb{R}^{H\times|R_{i}|\times d}
$$
by blocked matrix multiplication. The terms in the summation are not known to CompNode i, since slicing here is performed post-softmax. To correct for this, observe that for a row vector $\bm{x}∈\mathbb{R}^{N}$ and any $1≤ k≤\beta$ , again with slicing notation,
$$
\text{softmax}(\bm{x})[S_{k}]=\frac{\text{expsum}(\bm{x}[S_{k}])}{\sum_{l}%
\text{expsum}(\bm{x}[S_{l}])}\cdot\text{softmax}(\bm{x}[S_{k}])\in\mathbb{R}^{%
|S_{k}|}.
$$
Thus, the above summation can be simplified as follows, with $\odot$ denoting elementwise multiplication:
| | $\displaystyle\text{softmax}(\bm{a})[:,R_{i},:]\bm{v}$ | $\displaystyle=\sum_{k}\text{softmax}(\bm{a})[:,R_{i},S_{k}]\bm{v}[:,:,S_{k}]$ | |
| --- | --- | --- | --- |
with $\bm{n}^{R_{i}}∈\mathbb{R}^{H×|R_{i}|× 1}$ being the elementwise maximum of $\bm{m}^{R_{i}S_{k}}∈\mathbb{R}^{H×|R_{i}|× 1}$ over all $1≤ k≤\beta$ . The fraction is elementwise divison, and expsum is performed row-wise along the last dimension. This expression is numerically stable because each $\bm{m}^{R_{i}S_{k}}-\bm{n}^{R_{i}S_{k}}≤ 0$ and each $\bm{e}^{R_{i}S_{k}}≤ 1$ . Now, each aggregate term in the numerator and denominator summations is known to the CompNode. In essence, the CompNode is performing a weighted average of concatenated AttnNode $\bm{u}$ results, with the weights also coming from AttnNodes $\bm{m},\bm{e}$ results. To get the final output of attention corresponding to row indices in $R_{i}$ , the CompNode finally performs $O$ -projection. Algorithm 5 implements this.
Algorithm 5 CompNode i Single Layer Post-Pass
0: ${\bm{m}}^{R_{i}S_{k}},{\bm{e}}^{R_{i}S_{k}},{\bm{u}}^{R_{i}S_{k}}$ for $1≤ k≤\beta$
0: $\bm{o}^{R_{i}}$
1: Initialize $\bm{o}^{R_{i}}$ with zeroes and shape like $\bm{u}^{R_{i}S_{1}}$
2: Initialize $\bm{w}^{R_{i}}$ with zeroes and shape like $\bm{e}^{R_{i}S_{1}}$
3: $\bm{n}^{R_{i}}←\text{elementwise\_max}\{\bm{m}^{R_{i}S_{k}}\}_{k=1}^{\beta}$
4: for $k=1$ to $\beta$ do
5: $\bm{w}^{R_{i}}←\bm{w}^{R_{i}}+\text{exp}(\bm{m}^{R_{i}S_{k}}-\bm{n}^{%
R_{i}S_{k}})\odot\bm{e}^{R_{i}S_{k}}$
6: $\bm{o}^{R_{i}}←\bm{o}^{R_{i}}+\text{exp}(\bm{m}^{R_{i}S_{k}}-\bm{n}^{%
R_{i}S_{k}})\odot\bm{e}^{R_{i}S_{k}}\odot\bm{u}^{R_{i}S_{k}}$
7: end for
8: $\bm{o}^{R_{i}}←\bm{o}^{R_{i}}/\bm{w}^{R_{i}}$
9: $\bm{o}^{R_{i}}←\text{o\_proj}(\bm{o}^{R_{i}})$
10: return $\bm{o}^{R_{i}}$
6 Security Analysis
We now examine the security properties of Cascade. Cascade does not employ cryptographic techniques, so we can only elucidate on statistical security. Nevertheless, we examine a wide range of security considerations below. The security of Cascade is a function of its implementation parameters: the number of nodes participating, as well as the sharding strategy used, i.e. $\{R_{i}\}_{i=1}^{\alpha}$ and $\{S_{j}\}_{j=1}^{\beta}$ .
To analyze information leakage, we examine isolated computational stages of nodes. In all stages (pre-pass, attention-pass, or post-pass) of Cascade, no information is received during the computation. Therefore, all input leakage comes from the stage initialization. Thus, any leakage must occur from the following sharded tensors:
| | $\displaystyle\text{CompNode}_{i}$ | $\displaystyle→\bm{h}^{R_{i}},\{\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u%
}^{R_{i}S_{k}}\}_{k=1}^{\beta},$ | |
| --- | --- | --- | --- |
When more nodes participate, each has access to fewer tokens and hidden states, which makes reversing the full input more difficult. Even holding the node count fixed, the choice of sharding may also affect the difficulty of reversal. For example, is it more secure for a node to have access to $[h_{1},h_{5},h_{9},h_{13}]$ or $[h_{1},h_{2},h_{12},h_{13}]$ ?
As the space of sharding strategies is vast, we focus on $\bm{(c,\delta)}$ -sharding strategies, where each sharded index set takes the form of a ‘clustered-arithmetic’ or $(c,\delta)$ -sequence.
**Definition 6.1**
*We say a subset of indices $[N]$ is a $(c,\delta)$ -sequence if it takes the following form for some $i$ :
| | $\displaystyle\{$ | $\displaystyle i,i+1,...,i+c-1,\delta+i,\delta+i+1,...,\delta+i+c-1,2%
\delta+i,2\delta+i+1,...,2\delta+i+c-1,...\}.$ | |
| --- | --- | --- | --- |*
We focus on these strategies as they fulfill several security desiderata. We emphasize however that other strategies may be preferable to $(c,\delta)$ -sharding depending on the use-case. Note that under $(c,\delta)$ -sharding, the minimum number of CompNodes $\alpha$ needed to ensure all indices in $[N]$ are held by some node is given by $\alpha=\lceil\delta/c\rceil$ .
We will examine leakage to CompNodes from $\bm{h}^{R_{i}}$ in Section 6.1 and Section 6.2. We will analyze leakage from all shards to CompNodes in Section 6.3. Finally, we examine leakage from $\bm{q}^{S_{j}},\bm{k}^{S_{k}},\bm{v}^{S_{k}}$ to AttnNodes in Section 6.4.
6.1 Layer 0 CompNode Hidden Security
We first consider the security of CompNodes at layer 0 of the scheme, who each hold some token embeddings of the input prompt $\bm{x}$ . As embeddings are immediately reversible to their tokens through the lookup table, we conclude that Cascade should not be used when the security of every token is paramount. If individual token security is imperative, Cascade can be integrated with SMPC as in Appendix A.
Given that a CompNode has access to some of the $N$ token embeddings, say $[e_{n_{1}},e_{n_{2}},...,e_{n_{t}}]$ , the possibility of reconstruction of the full prompt $x$ is theoretically lower bounded by the entropy of the distribution
| | $\displaystyle p(\{x_{j}:j∈[N]\setminus\{n_{1},...,n_{t}\}\}\mid x_{n_{1}}%
,x_{n_{2}},...,x_{n_{t}})$ | |
| --- | --- | --- |
The true distribution cannot easily be computed, but we may approximate it using masked token infilling, as in the training of models like BERT (Devlin et al., 2019), RoBERTa (Liu et al., 2019), and ModernBERT (Warner et al., 2024). We use the recently released state-of-the-art ModernBERT-large to probe properties of this distribution under $(c,\delta)$ -sharding. We use 200 samples from FineWeb-Edu to compute the ROUGE-L (Lin, 2004) score over argmax-token generation. Our results are shown in Figure 2. We see that ROUGE-L score diminishes as both $c$ and $\alpha$ increase; good security seems to be achieved for $c,\alpha≥ 8$ and thus $\delta≥ 64$ .
<details>
<summary>extracted/6603058/images/rouge_c_alpha.png Details</summary>
### Visual Description
## Heatmap: ROUGE Score of token-infilling, c vs α
### Overview
The image is a heatmap visualizing the ROUGE score of token-infilling as a function of two parameters, 'c' and 'α'. The x-axis represents 'c', and the y-axis represents 'α'. The color intensity indicates the average ROUGE score, with lighter colors (yellow) representing higher scores and darker colors (black/purple) representing lower scores. The heatmap shows a triangular region of data, with the upper-right triangle being empty (white).
### Components/Axes
* **Title:** ROUGE Score of token-infilling, c vs α
* **X-axis:**
* Label: c
* Scale: Logarithmic
* Markers: 1, 2, 4, 8, 16, 32, 64, 128, 256
* **Y-axis:**
* Label: α
* Scale: Logarithmic
* Markers: 2, 4, 8, 16, 32, 64, 128, 256, 512
* **Colorbar (Legend):**
* Label: Average ROUGE Score
* Scale: Linear
* Range: 0.1 to 0.5
* Colors: Progresses from dark purple (0.1) to yellow (0.5) through shades of orange and red.
### Detailed Analysis
The heatmap displays the average ROUGE score for different combinations of 'c' and 'α'. The ROUGE score is highest (yellow) when both 'c' and 'α' are low (close to 1 and 2 respectively). As 'c' and 'α' increase, the ROUGE score decreases, transitioning to darker colors (purple/black). The upper triangle of the heatmap is white, indicating missing or invalid data points for those combinations of 'c' and 'α'.
* **Region of High ROUGE Score:** Located in the bottom-left corner, where both 'c' and 'α' are small. The color is a bright yellow, corresponding to a ROUGE score of approximately 0.5.
* c = 1, α = 2: ROUGE Score ≈ 0.5
* c = 2, α = 2: ROUGE Score ≈ 0.45
* c = 1, α = 4: ROUGE Score ≈ 0.45
* **Transition Region:** As 'c' and 'α' increase, the color transitions to orange, then red, then purple, and finally black. This indicates a decreasing ROUGE score.
* c = 4, α = 4: ROUGE Score ≈ 0.35
* c = 8, α = 8: ROUGE Score ≈ 0.2
* **Region of Low ROUGE Score:** Located towards the top and right of the data region. The color is dark purple/black, corresponding to a ROUGE score of approximately 0.1 or lower.
* c = 32, α = 32: ROUGE Score ≈ 0.1
* c = 64, α = 64: ROUGE Score < 0.1
* **Empty Region:** The upper-right triangle of the heatmap is white, indicating no data.
### Key Observations
* The highest ROUGE scores are achieved when both 'c' and 'α' are low.
* The ROUGE score decreases as 'c' and 'α' increase.
* There is a clear negative correlation between 'c' and 'α' and the ROUGE score.
* The data is only available for combinations of 'c' and 'α' where α >= c.
### Interpretation
The heatmap suggests that the token-infilling performance, as measured by the ROUGE score, is highly sensitive to the parameters 'c' and 'α'. Lower values of both parameters lead to better performance. The absence of data in the upper triangle implies a constraint or condition where α must be greater than or equal to c for the token-infilling process to be valid or meaningful. The rapid decrease in ROUGE score as 'c' and 'α' increase indicates that these parameters might need to be carefully tuned to achieve optimal performance. The relationship between 'c', 'α', and the ROUGE score could be further explored to understand the underlying mechanisms of the token-infilling process and identify the optimal parameter settings.
</details>
Figure 2: ROUGE-L scores for Layer $0 0$ token prediction using ModernBERT-Large, as a function of $c$ , the number of ‘clusters’ in the sharding scheme, and $\alpha$ , the number of CompNodes. Higher $\alpha$ and higher $c$ tend towards lower ROUGE-L, increasing security.
6.2 Layer $>0$ CompNode Hidden Security
We now turn our attention to CompNode security of $\bm{h}^{R_{i}}$ at deeper layers of the LLM. Can we select a sharding strategy that defends against our attack outlined in Section 4? Also, do we remain secure to learning-based attacks as in prior works Wan et al. (2024) and Morris et al. (2023)?
Vocab-Matching Attack Our defense against vocab-matching is to ensure large token gaps in nodes. For a $(c,\delta)$ -sharding scheme, in each shard, the distance from one ‘cluster’ of indices to the next is $\delta-c+1$ . Therefore, when one carries out the generalized vocab-matching attack, the cost scales exponentially with the maximum gap, so it is on the order of $V^{\delta-c+1}$ . Note $\delta-c+1=(\alpha-1)c+1$ , so regardless of the value of the vocab-matching threshold $\rho$ , we can increase $\alpha$ or $c$ until $\delta-c+1≥\rho$ to prevent vocab-matching. For $\rho=3$ , all $\alpha,c≥ 2$ satisfy this.
One objection to this line of argument is that if the adversary has a strong prior on the input text distribution, they can eliminate unlikely tokens from forward passes in the generalized attack, essentially only iterating over tokens in $\mathcal{V}_{0}⊂eq\mathcal{V}$ . However, the cost of this optimized attack is on the order of $|\mathcal{V}_{0}|^{\delta-c+1}$ , so it still scales exponentially. Thus, even if $|\mathcal{V}_{0}|$ is hundreds of times smaller than $|\mathcal{V}|\sim 100000$ , large enough $\rho$ ensures that $\delta-c+1≥\rho$ remains secure.
Learning-Based Attacks We now consider learning-based reversal attacks. Although hidden states may not be as easily reversible as token embeddings, it may be the case that as attention propagates information among hidden states, the text-infilling task is easier at deeper layers of the LLM than at the purely textual level. We examine this hypothesis by fine-tuning Gemma-2-2B-IT and Llama-3.1-8B-Instruct on $(c,\delta)$ -masked input sequences from FineWeb-Edu, with the target labels being the full input sequence. Note this is a ‘worst-case’ scenario where the adversary knows the $c$ and $\delta$ parameters. We update both models to use a bidirectional mask, in line with the token-infilling nature of this task. We train until the evaluation loss on a held-out set converges over layer 1 representations from each model, and evaluate on layer 1 hidden states. Our approach is similar to Wan et al. (2024); Morris et al. (2023).
Our ROUGE-L scores for reconstructed Gemma-2-2B-IT prompts are shown in Table 3. When $c,\alpha≥ 8$ , we have a ROUGE-L score less than 0.25, indicating significant reconstruction difficulty. We also tested on $c=4,\alpha=16$ and $c=8,\alpha=24$ and obtained ROUGE-L scores of 0.173 and 0.144, supporting that security continues to improve by scaling $c$ and $\alpha$ .
Table 3: ROUGE-L scores of text reconstruction from the hiddens of layer 1 of Gemma-2-2B-IT for various values of $c$ and $\alpha$ under $(c,\delta)$ -sharding. Increasing $c$ or $\alpha$ results in worse reconstruction.
| $c=1$ $c=4$ $c=8$ | 0.701 0.427 0.355 | 0.467 0.290 0.222 | 0.349 0.230 0.191 |
| --- | --- | --- | --- |
We similarly analyze BLEU scores Papineni et al. (2002) in Table 4. We find a similar trend as in Table 3, observing that security improves as $c$ or $\alpha$ increases. Here, for $c≥ 4,\alpha=8$ , we have a BLEU score between $0.1$ and $0.2$ , indicating only marginal reconstructionability; and for $c≥ 4,\alpha=12$ , the BLEU score lies below $0.1$ , indicating nearly no reconstructionability.
Table 4: BLEU scores of text reconstruction from layer $1$ hiddens of Gemma-2-2B-IT for different values of $c,\alpha$ under $(c,\delta)$ -sharding. Increasing $c$ or $\alpha$ results in worse reconstruction.
| $c=1$ $c=4$ $c=8$ | 0.537 0.352 0.246 | 0.229 0.133 0.123 | 0.130 0.086 0.083 |
| --- | --- | --- | --- |
We further examine the choice of hidden layer, and if Llama representations are decoded better, in Appendix B. We find similar or better security across these other parameter choices.
6.3 CompNode Full Security
In Section 6.1 and Section 6.2, we considered leakage to CompNodes from $\bm{h}^{R_{i}}$ . Now, we consider how the shards $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u}^{R_{i}S_{k}}$ impact security. Our main theorem shows that with this leakage, a scheme like $(c,\delta)$ -sharding is, in some sense, necessary to maintain security against vocab-matching: without sufficiently large gaps between consecutive clusters of indices, a variant of the attack can be carried out at Layer 0 to reveal additional tokens to a node.
**Theorem 6.2**
*Suppose the attention mask $\bm{s}$ is unidirectional. Furthermore, denote the vocab-matching threshold as $\rho$ , as defined in Definition 4.1. Then to prevent the generalized vocab-matching attack, for all $i∈[\alpha]$ , each gap between clusters of consecutive indices in $R_{i}$ must have size $≥\rho$ .*
* Proof*
Fix $i$ , and denote $R_{i}=\{r_{1},r_{2},...,r_{m}\}$ in ascending order. Consider any $k∈[\beta]$ . Now, since $\bm{a}^{R_{i}S_{k}}=\bm{q}^{R_{i}}(\bm{k}^{S_{k}})^{T}+\bm{s}^{R_{i}S_{k}}$ , and $\bm{s}$ is unidirectional, we see that the $l$ th row of $\text{exp}(\bm{a}^{R_{i}S_{k}})\bm{v}^{S_{k}}$ is exactly $f(r_{l},\{s∈ S_{k}:s<r_{l}\})$ , where we denote
$$
f(r,S):=\sum_{s\in S}\text{exp}\left(\bm{q}[:,r,:](\bm{k}[:,s,:])^{T}\right)%
\bm{v}[:,s,:].
$$
This $l$ th row is known to CompNode i because we can write it in terms of known shards:
$$
\text{exp}(\bm{a}^{R_{i}S_{k}})\bm{v}^{S_{k}}=(\text{softmax}(\bm{a}^{R_{i}S_{%
k}})\bm{v}^{S_{k}})\odot\text{expsum}(\bm{a}^{R_{i}S_{k}})=\bm{u}^{R_{i}S_{k}}%
\odot\bm{e}^{R_{i}S_{k}}\odot\text{exp}(\bm{m}^{R_{i}S_{k}}).
$$
Thus, for each $r_{l}∈ R_{i}$ and $k∈[\beta]$ , we see that CompNode i knows $f(r_{l},\{s∈ S_{k}:s<r_{l}\})$ . At Layer $0 0$ , since $\bm{q},\bm{k},\bm{v}$ are linear projections of token embeddings, we see that $f(r,S)$ only depends on tokens with indices in the set $\{r\}\cup S$ . Since CompNode i knows all tokens with indices $R_{i}$ from $\bm{h}^{R_{i}}$ at layer 0, then each $f(r_{l},\{s∈ S_{k}:s<r_{l}\})$ only depends on unknown tokens at indices $\{s∈ S_{k}:s<r_{l}\}\setminus R_{i}$ . Thus, $f(r_{l+1},\{s∈ S_{k}:s<r_{l+1}\})$ depends on the same unknown tokens as $f(r_{l},\{s∈ S_{k}:s<r_{l}\})$ , and extra unknown tokens at $\{s∈ S_{k}:r_{l}<s<r_{l+1}\}$ , which we call the $(k,l)$ -gap. If there are $<\rho$ tokens in a $(k,l)$ -gap for some $l$ , then following the generalized vocab-matching attack in Section 4, CompNode i can perform a forward pass over all $V^{<\rho}$ candidate sequences of such unknown tokens in the $(k,l)$ -gap and compute candidate values for $f(r_{l},\{s∈ S_{k}:s<r_{l}\})$ and $f(r_{l+1},\{s∈ S_{k}:s<r_{l+1}\})$ , and then select the candidate sequence which allows these quantities to match their known values. This would allow them to recover all tokens in the $(k,l)$ -gap. Thus, to prevent this vocab-matching attack, we need each $(k,l)$ -gap to have size $≥\rho$ or zero. This forces clusters of consecutive indices in $R_{i}$ to have size $≥\rho$ gaps between them. ∎
Note that $(c,\delta)$ -sharding satisfies the requirement of the theorem if and only if $\delta-c+1≥\rho+1$ . Thus, this theorem induces a slightly stronger requirement than $\delta-c+1≥\rho$ from Section 6.2, which only considered leakage from $\bm{h}^{R}_{i}$ . In Appendix E, we argue further that $(c,\delta)$ -sharding with sufficiently large $\delta-c+1$ is sufficient for Layer 0 security, by demonstrating a reduction to an intractable linear program.
6.4 AttnNode Full Security
We now turn our attention to the analysis of AttnNode security. Recall that $S$ and $T$ sharding are set to be equal by the argument in Section 5.1, and so in the following we discuss only $S$ sharding with the implication that $T$ receives the same treatment.
Given a particular $R$ sharding, there are many possibilities for $S$ sharding. One option is to set $S$ and $R$ sharding equal, and have each AttnNode jk receive the union of query shards from CompNode j and key/value shards from CompNode k. However, this results in each AttnNode having access to twice as many $Q,K,V$ -rows as each CompNode. To reduce such leakage to AttnNodes, we propose a further $m$ -split of AttnNodes as follows.
To form $S$ sharding from $R$ sharding, we can let $\beta=m\alpha$ , and partition each $R_{i}$ into $m$ shards $R_{i,1},...,R_{i,m}$ . Then the $S$ shards are $S_{m(i-1)+x}=R_{i,x}$ for all $1≤ i≤\alpha$ and $1≤ x≤ m$ . This ensures that pairwise coverage is maintained, but reduces the number of tokens (technically $Q,K,V$ -projections of hidden state rows, but these are reversible to tokens in the worst case at layer $0 0$ ) that each AttnNode has access to by a factor of $m$ . Using this split increases the value of $\beta^{2}$ , the total number of AttnNodes, by a factor of $m^{2}$ .
There still remains a degree of freedom in deciding the exact choice of subdividing $R_{i}$ into the subsets $R_{i,x}$ . Under the assumption that $R_{i}$ follows a $(c,\delta)$ -sharding scheme, we propose that $R_{i,x}$ contains the elements of sorted $R_{i}$ at indices $x,x+\delta,...,x+{(t-1)\delta}$ , where $t=\frac{N}{c\alpha}$ and the split factor is $m=c$ .
For example, suppose that $\alpha=3$ , $c=2$ , $\delta=6$ , and $N=18$ . With a split of $m=2$ , we have $\beta=m\alpha=6$ . Then,
| | $\displaystyle R_{1}=\{1,2,7,8,13,14\}\quad$ | $\displaystyle R_{2}=\{3,4,9,10,15,16\}\quad$ | $\displaystyle R_{3}=\{5,6,11,12,17,18\}$ | |
| --- | --- | --- | --- | --- |
where we denote $S_{ixjy}=R_{i,x}\cup R_{j,y}$ , which is the set of tokens that AttnNode a,b has access to for $a=m(i-1)+x,b=m(j-1)+y$ . In other words, $R_{i}$ above are the sets of token indices that the CompNodes receive, and the $S_{ixjy}$ are the sets of token indices that the AttnNodes receive. Note that some $S_{ixjy}$ entries are not included above (there are $21$ listed, but $\beta^{2}=36$ AttnNodes) because they exactly match a listed entry by $S_{ixjy}=S_{jyix}$ .
Vocab-Matching Attack Sharding $S$ in this way prevents vocab-matching, like in Section 6.2. Indeed, as each $R_{i,x}$ has elements that are separated by $\delta$ , and each $S_{ixjy}$ combines elements from two different $R_{i,x}$ ’s, then there cannot be $3$ consecutive elements in $S_{ixjy}$ if $\delta>2$ . Also, the largest number of missing tokens between two elements of $S_{ixjy}$ (i.e. the largest ‘gap’) is lower bounded by $\frac{\delta}{2}$ . Therefore, for sufficiently large $\delta$ , the vocab-matching attack is infeasible.
Learning-Based Attacks To test security against learning-based attacks, we conducted experiments with the above scheme for $m=\{2,3,4\}$ , with $c=8$ and $\alpha=8$ , and the same dataset and setup as in Section 6.2. Due to computational constraints, we focus our experiments on Gemma-2-2B-IT on layer 1; we expect similar trends for Llama-3.1-8B-Instruct and other layers. We train a single model for all shard possibilities that arise from $S_{ixjy}$ . Experiments are conducted with the same dataset and model setup as described . Our results are shown in Table 5. We see that although $m=2$ results in a relatively higher ROUGE-L than that for the CompNodes in Table 3, the score for $m=4$ is very similar; therefore, we recommend using $m≥ 4$ for security.
Table 5: ROUGE-L scores for different values of splitting parameter $m$ on layer 1 of Gemma-2-2B-IT with $c=8,\alpha=8$ . We see that the score for $m=4$ is similar to that of CompNodes in Table 3 for the same $c$ and $\alpha$ .
| 2 3 4 | 0.3057 0.2643 0.2376 |
| --- | --- |
6.5 Collusion Resistance
Most SMPC schemes break when nodes share information. For example, in two-party additive secret sharing, a private value $x$ is split into shares $\langle x\rangle_{0},\langle x\rangle_{1}$ held by parties $0,1$ , and they can reconstruct $x=\langle x\rangle_{0}+\langle x\rangle_{1}$ if they collude. In Appendix G, we similarly analyze the collusion resistance properties of Cascade. We find that colluding CompNodes and AttnNodes gain access to the union of their token shards. Because unions of shards have smaller gaps between consecutive clusters, collusion directly impacts the feasibility of generalized vocab-matching, such as in Theorem 6.2. Thus, although the exact collusion limit depends on security parameters like the vocab-matching threshold $\rho$ and the particular sharding scheme, Cascade generally requires a large number of non-colluding nodes for security.
7 Communication and Computational Costs
We now enumerate communication and computational costs associated with Cascade. As in previous SMPC works (Li et al., 2023; Dong et al., 2023; Li et al., 2024), we assume perfect parallel transport in communication, a homogeneous node-wise bandwidth of $B$ , and an inter-node latency of $\tau$ . We denote $F$ as the number of bytes per element.
7.1 Theoretical Costs
We first provide theoretical estimates of costs in Cascade.
Computation There is almost no floating point overhead in Cascade relative to vanilla inference: see Appendix C. We provide empirical evidence for this in Section 7.2.
Communication We present single layer communication byte and time overhead formulae, derived in Appendix D:
$$
\displaystyle=\beta F(2dH+2dH_{KV}+2H)\cdot N \tag{1}
$$
$$
\displaystyle=2\tau+\frac{\beta Fd(H+2H_{KV})}{B}\cdot\max_{i}|R_{i}| \displaystyle+\frac{F(d+2)H}{B}\cdot\max_{j}|S_{j}|. \tag{2}
$$
Note the only sharding parameter Equation 1 depends on is $\beta$ , and in fact it scales linearly with $\beta$ . Thus, byte overhead is minimized when $\beta$ is minimized, which coincides with minimizing the AttnNode count. Furthermore, communication time is minimized when $\beta\max_{i}|R_{i}|$ and $\max_{j}|S_{j}|$ are minimal. For fixed $\alpha,\beta$ , since $\{R_{i}\}_{i=1}^{\alpha}$ and $\{S_{j}\}_{j=1}^{\beta}$ partition $[N]$ , then $\max_{i}|R_{i}|≥\lceil N/\alpha\rceil$ and $\max_{j}|S_{j}|≥\lceil N/\beta\rceil$ . Here, equality holds when all $R_{i}$ , and all $S_{j}$ , are around the same size. Thus, $(c,\delta)$ -sharding, which approximately divides tokens between CompNodes in equal quantity, has optimal communication time.
7.2 Performance Experiments
We now evaluate the real-world performance of Cascade through the distributed computing framework Ray (Moritz et al., 2018). We run our experiments on Paperspace machines with 16 vCPU and 64GB RAM, with the CPU model being Intel Xeon Gold 6226R CPU @ 2.90GHz. All machines are colocated in the same region with an average bandwidth of 2 Gbps and latency of 0.38 ms.
We benchmark against two recent SMPC schemes for LLM inference, MPCFormer (Li et al., 2023) and Puma (Dong et al., 2023). For MPCFormer, we modify the Crypten implementation to use public rather than private weights, to match our open-weights setting. Puma data is taken from Dong et al. (2023), as it is built on SPU with its own set of optimizations. We perform inference (a single forward pass) on Bert-Base and Bert-Large with an input prompt of 128 tokens – we repeat this measurement 100 times as there is run-to-run variability in the timing. Our results are shown in Table 6.
We first compare Cascade for $\alpha=1$ without Ray against vanilla inference, to estimate protocol overhead. The mean runtime is 109ms vs. 91ms for vanilla inference for Bert-Base, and 320ms vs. 273ms for Bert-Large. Profiling shows this minor increase is due to the attention-score compilation step discussed in Appendix C. Nevertheless, the mean runtime is within the 95% confidence interval of vanilla inference in both cases.
Next, we compare the performance of Cascade with $\alpha=1$ and using Ray. As seen above, this is slower than not using Ray by around a factor of $3×$ . This slowdown can be attributed to framework-specific overhead, such as serialization. In other words, Cascade is so efficient that the distributed-compute framework overhead now constitutes a significant proportion of its slowdown from vanilla inference, rather than protocol-specific overhead.
Table 6: Total runtime means and 95% confidence intervals in seconds over 100 trials, for a single $128$ -token prompt forward pass on Bert-Base and Bert-Large for MPCFormer, Puma, and various settings of Cascade. Higher $\alpha$ corresponds to increased node counts and security.
| Scheme MPCFormer Puma | Bert-Base (s) 55.32 33.91 | Bert-Large (s) 141.22 73.72 |
| --- | --- | --- |
| Cascade ${}_{\alpha=1,\text{no Ray}}$ | 0.11 [0.10, 0.13] | 0.32 [0.23,1.07] |
| Cascade α=1 | 0.32 [0.31, 0.36] | 1.01 [0.97, 1.09] |
| Cascade α=4 | 0.59 [0.51,0.69] | 1.57 [1.44, 1.73] |
| Cascade α=8 | 0.74 [0.62, 0.96] | 1.58 [1.27, 1.97] |
| Vanilla | 0.09 [0.08,0.12] | 0.27 [0.20, 0.99] |
We further benchmark Cascade with $\alpha=4,8$ , using a $(c,\delta)$ -sharding scheme with $c=1$ and no $m$ -splits. We use clusters of $6$ and $18$ machines for $\alpha=4$ and $8$ . By testing the saturation of parallelization gains, we take an optimal $4$ cores per node. Performance in these cases is slower than for $\alpha=1$ , due to communication overhead. Still, Cascade is $90×$ faster than MPCFormer and $45×$ faster than Puma for Bert-Large, even in its slowest setting.
Furthermore, although it is difficult to make strong conclusions from 3 data points, we observe that the increase in runtime for both models from $\alpha=1$ to $\alpha=8$ appears sublinear, even as the number of distinct machines increases superlinearly. This is expected due to parallel computation and communication across the nodes. Indeed, for Bert-Large, there is no statistically significant increase in runtime from $\alpha=4$ to $\alpha=8$ , suggesting that it may even be saturating. In future work, we hope to run on even larger values of $\alpha$ , with potential $m$ -splitting, to confirm this hypothesis.
To support our theoretical estimates in Section 7.1, we also examined the total communication in the $\alpha=1$ case using tshark. We found that the total communicated bytes are within $2\%$ of Equation 1. The extra bytes can be attributed to Ray-specific metadata and a slight increase from the raw bytes size due to serialization. We present a comparative table of total communicated bytes for Cascade versus MPCFormer and Puma in Table 7. We see that even in the most expensive $\alpha=8$ setting, Cascade is $\sim 160×$ more efficient in total bytes transferred than MPCFormer, and $\sim 140×$ more efficient than Puma.
Table 7: Total gigabytes (GB) communicated for a single forward pass on Bert-Base and Bert-Large for MPCFormer, Puma, and Cascade with various settings. A prompt length of 128 is used.
| MPCFormer Puma Cascade α=1 | 12.089 10.773 0.009 | 32.577 27.246 0.025 |
| --- | --- | --- |
| Cascade α=4 | 0.038 | 0.101 |
| Cascade α=8 | 0.076 | 0.203 |
7.3 Scalability
We also tested the scaling behaviour of Cascade as a function of the model size. We again perform a forward pass with an input prompt of size 128 tokens, setting $\alpha=2$ , and without $m$ -splitting. We repeat this across 100 runs. Our results are shown in Table 8. We observe that the runtime grows sublinearly with the model size, suggesting Cascade will also scale well to the largest modern LLMs. We also note that Puma reported around $300$ seconds for a full forward pass on Llama-2-7B; Cascade is over $20×$ faster.
Table 8: Total runtime means and 95% confidence intervals in seconds over 100 trials, for a single $128$ -token prompt forward pass on various sizes of LLM models with Cascade, with $\alpha=2$ and no m-splitting. We observe sublinear increase in mean runtime as the model size grows, indicating that Cascade is well suited to scale to the largest modern LLMs.
| Bert-Base Bert-Large Llama-3.2-1B-Instruct | 110M 335M 1B | 0.70 1.33 2.67 | [0.62, 0.74] [1.24, 1.46] [2.33, 2.96] |
| --- | --- | --- | --- |
| Llama-2-7B | 7B | 12.71 | [11.07, 14.07] |
| Llama-2-13B | 13B | 22.72 | [20.58, 25.99] |
8 Conclusion & Future Work
We generalized a recent attack for decoding LLM hidden states into their original user text in the setting of sharding in the token dimension. To defend against this, we introduced a novel multi-party scheme, Cascade, that employs token sharding of hidden states. We showed that, under appropriate settings, Cascade is resistant to our attack, as well as existing learning-based attacks in the literature. Finally, we benchmarked Cascade on modern state-of-the-art LLMs, demonstrating significant improvements compared to existing multi-party inference schemes in terms of runtime, total communicated bytes, and scalability.
Future directions of work could examine the security and cost of other sharding strategies, particularly in unreliable network settings. This could include considerations of the asychronous setting, alongside other optimizations in Appendix F. Future work could also precisely optimize the security of Cascade for specific use cases, based on the form and frequency of private user data in prompts. We also encourage investigation into mixing Cascade with other SMPC schemes, such as our proposal in Appendix A, to offer a flexible tradeoff between cryptographic security and efficiency. Such work could aim to prevent direct exposure of tokens to nodes, which is one of the main limitations of Cascade.
References
- Akimoto et al. (2023) Yoshimasa Akimoto, Kazuto Fukuchi, Youhei Akimoto, and Jun Sakuma. Privformer: Privacy-preserving transformer with mpc. In 2023 IEEE 8th European Symposium on Security and Privacy (EuroS&P), pages 392–410, 2023. doi: 10.1109/EuroSP57164.2023.00031.
- DeepSeek-AI et al. (2025) DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, J. L. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Liang Zhao, Litong Wang, Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li, Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang, Qinyu Chen, Qiushi Du, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. S. Li, Shuang Zhou, Shaoqing Wu, Shengfeng Ye, Tao Yun, Tian Pei, Tianyu Sun, T. Wang, Wangding Zeng, Wanjia Zhao, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang, Xiaokang Chen, Xiaotao Nie, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, X. Q. Li, Xiangyue Jin, Xiaojin Shen, Xiaosha Chen, Xiaowen Sun, Xiaoxiang Wang, Xinnan Song, Xinyi Zhou, Xianzu Wang, Xinxia Shan, Y. K. Li, Y. Q. Wang, Y. X. Wei, Yang Zhang, Yanhong Xu, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Wang, Yi Yu, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yuan Ou, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yaohui Li, Yi Zheng, Yuchen Zhu, Yunxian Ma, Ying Tang, Yukun Zha, Yuting Yan, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan, Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025. URL https://arxiv.org/abs/2501.12948.
- Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019. URL https://arxiv.org/abs/1810.04805.
- Dong et al. (2023) Ye Dong, Wen jie Lu, Yancheng Zheng, Haoqi Wu, Derun Zhao, Jin Tan, Zhicong Huang, Cheng Hong, Tao Wei, and Wenguang Chen. Puma: Secure inference of llama-7b in five minutes, 2023. URL https://arxiv.org/abs/2307.12533.
- Goldreich et al. (1987) O. Goldreich, S. Micali, and A. Wigderson. How to play any mental game. In Proceedings of the Nineteenth Annual ACM Symposium on Theory of Computing, STOC ’87, page 218–229, New York, NY, USA, 1987. Association for Computing Machinery. ISBN 0897912217. doi: 10.1145/28395.28420. URL https://doi.org/10.1145/28395.28420.
- Hao et al. (2022) Meng Hao, Hongwei Li, Hanxiao Chen, Pengzhi Xing, Guowen Xu, and Tianwei Zhang. Iron: Private inference on transformers. In Advances in Neural Information Processing Systems, volume 35, pages 15718–15731, 2022.
- Huang et al. (2022) Zhicong Huang, Wen jie Lu, Cheng Hong, and Jiansheng Ding. Cheetah: Lean and fast secure two-party deep neural network inference. In 31st USENIX Security Symposium (USENIX Security 22), pages 809–826, 2022.
- Kasneci et al. (2023) Enkelejda Kasneci, Kathrin Seßler, Stefan Küchemann, Maria Bannert, Daryna Dementieva, Frank Fischer, Urs Gasser, Georg Groh, Stephan Günnemann, Eyke Hüllermeier, Stephan Krusche, Gitta Kutyniok, Tilman Michaeli, Claudia Nerdel, Juergen Pfeffer, Oleksandra Poquet, Michael Sailer, Albrecht Schmidt, Tina Seidel, and Gjergji Kasneci. Chatgpt for good? on opportunities and challenges of large language models for education. Learning and Individual Differences, 103:102274, 01 2023. doi: 10.1016/j.lindif.2023.102274.
- Li et al. (2023) Dacheng Li, Rulin Shao, Hongyi Wang, Han Guo, Eric P. Xing, and Hao Zhang. Mpcformer: fast, performant and private transformer inference with mpc, 2023. URL https://arxiv.org/abs/2211.01452.
- Li et al. (2024) Zhengyi Li, Kang Yang, Jin Tan, Wen jie Lu, Haoqi Wu, Xiao Wang, Yu Yu, Derun Zhao, Yancheng Zheng, Minyi Guo, and Jingwen Leng. Nimbus: Secure and efficient two-party inference for transformers, 2024. URL https://arxiv.org/abs/2411.15707.
- Lin (2004) Chin-Yew Lin. ROUGE: A package for automatic evaluation of summaries. In Text Summarization Branches Out, pages 74–81, Barcelona, Spain, July 2004. Association for Computational Linguistics. URL https://aclanthology.org/W04-1013/.
- Liu et al. (2019) Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019. URL https://arxiv.org/abs/1907.11692.
- Luo et al. (2024a) Jinglong Luo, Guanzhong Chen, Yehong Zhang, Shiyu Liu, Hui Wang, Yue Yu, Xun Zhou, Yuan Qi, and Zenglin Xu. Centaur: Bridging the impossible trinity of privacy, efficiency, and performance in privacy-preserving transformer inference, 2024a. URL https://arxiv.org/abs/2412.10652.
- Luo et al. (2024b) Jinglong Luo, Yehong Zhang, Zhuo Zhang, Jiaqi Zhang, Xin Mu, Hui Wang, Yue Yu, and Zenglin Xu. Secformer: Towards fast and accurate privacy-preserving inference for large language models. arXiv preprint arXiv:2401.00793, 2024b.
- Moritz et al. (2018) Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael I. Jordan, and Ion Stoica. Ray: A distributed framework for emerging ai applications, 2018. URL https://arxiv.org/abs/1712.05889.
- Morris et al. (2023) John X. Morris, Wenting Zhao, Justin T. Chiu, Vitaly Shmatikov, and Alexander M. Rush. Language model inversion, 2023. URL https://arxiv.org/abs/2311.13647.
- Pang et al. (2023) Qi Pang, Jinhao Zhu, Helen Möllering, Wenting Zheng, and Thomas Schneider. BOLT: Privacy-preserving, accurate and efficient inference for transformers. Cryptology ePrint Archive, Paper 2023/1893, 2023. URL https://eprint.iacr.org/2023/1893.
- Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 311–318, 2002.
- Qwen et al. (2025) Qwen, :, An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, Huan Lin, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jingren Zhou, Junyang Lin, Kai Dang, Keming Lu, Keqin Bao, Kexin Yang, Le Yu, Mei Li, Mingfeng Xue, Pei Zhang, Qin Zhu, Rui Men, Runji Lin, Tianhao Li, Tianyi Tang, Tingyu Xia, Xingzhang Ren, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yu Wan, Yuqiong Liu, Zeyu Cui, Zhenru Zhang, and Zihan Qiu. Qwen2.5 technical report, 2025. URL https://arxiv.org/abs/2412.15115.
- Thirunavukarasu et al. (2023) Arun James Thirunavukarasu, Darren Shu Jeng Ting, Kabilan Elangovan, Laura Gutierrez, Ting Fang Tan, and Daniel Shu Wei Ting. Large language models in medicine. Nature Medicine, 29:1930–1940, 2023. doi: 10.1038/s41591-023-02459-w. URL https://www.nature.com/articles/s41591-023-02459-w. Review Article, Published: 17 July 2023.
- Thomas et al. (2025) Rahul Thomas, Louai Zahran, Erica Choi, Akilesh Potti, Micah Goldblum, and Arka Pal. An attack to break permutation-based private third-party inference schemes for llms. arXiv preprint arXiv:2505.18332, 2025.
- Wan et al. (2024) Zhipeng Wan, Anda Cheng, Yinggui Wang, and Lei Wang. Information leakage from embedding in large language models, 2024. URL https://arxiv.org/abs/2405.11916.
- Warner et al. (2024) Benjamin Warner, Antoine Chaffin, Benjamin Clavié, Orion Weller, Oskar Hallström, Said Taghadouini, Alexis Gallagher, Raja Biswas, Faisal Ladhak, Tom Aarsen, Nathan Cooper, Griffin Adams, Jeremy Howard, and Iacopo Poli. Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference, 2024. URL https://arxiv.org/abs/2412.13663.
- Xu et al. (2024) Hengyuan Xu, Liyao Xiang, Hangyu Ye, Dixi Yao, Pengzhi Chu, and Baochun Li. Permutation equivariance of transformers and its applications, 2024. URL https://arxiv.org/abs/2304.07735.
- Yao (1982) Andrew C. Yao. Protocols for secure computations. In 23rd Annual Symposium on Foundations of Computer Science (sfcs 1982), pages 160–164, 1982. doi: 10.1109/SFCS.1982.38.
- Yuan et al. (2024) Mu Yuan, Lan Zhang, and Xiang-Yang Li. Secure transformer inference protocol, 2024. URL https://arxiv.org/abs/2312.00025.
- Zheng et al. (2024) Fei Zheng, Chaochao Chen, Zhongxuan Han, and Xiaolin Zheng. Permllm: Private inference of large language models within 3 seconds under wan, 2024. URL https://arxiv.org/abs/2405.18744.
- Zhu et al. (2024) Wenhao Zhu, Hongyi Liu, Qingxiu Dong, Jingjing Xu, Shujian Huang, Lingpeng Kong, Jiajun Chen, and Lei Li. Multilingual machine translation with large language models: Empirical results and analysis, 2024. URL https://arxiv.org/abs/2304.04675.
Appendix A Integrating Cascade with SMPC
In Section 6.1, we concluded that Cascade should only be used when it is safe to reveal some number of tokens, but we mentioned that integration with SMPC could alleviate this requirement. Here, we describe how Cascade can be fused with a number of SMPC schemes to improve token security at the expense of computational and communication cost. The idea is to form an $L$ -layered Cascade-SMPC split, where a general SMPC scheme is executed on the first $L$ LLM layers and Cascade is carried out thereafter.
We begin with the SMPC stage for layers $≤ L$ . Denote the layer 0 input embeddings as $\bm{x}$ , and the layer $L$ hidden states as $\bm{h}$ . For the first $L$ layers, we execute any SMPC protocol $\Phi$ based on additive secret sharing. Suppose $\Phi$ involves $t$ nodes $\mathcal{N}_{1},...,\mathcal{N}_{t}$ . Then $\Phi$ begins by additively splitting $\bm{x}=\sum_{s=1}^{t}\bm{x}_{s}$ , with each share $\bm{x}_{s}$ given to $\mathcal{N}_{s}$ . By decomposing LLM layers into SMPC-friendly functions or approximations, $\Phi$ allows each node $\mathcal{N}_{s}$ to compute some additive share $\bm{h}_{s}$ of the layer $L$ hidden states, such that $\bm{h}=\sum_{s=1}^{t}\bm{h}_{s}$ . A key ingredient of SMPC schemes based on additive secret sharing is computational indistinguishability: no node gains any information about $\bm{x}$ during the execution of $\Phi$ .
Once each $\mathcal{N}_{s}$ gets $\bm{h}_{s}$ , we aim to execute Cascade for layers $>L$ in the LLM. We initialize a Cascade sharding scheme $\{R_{i}\},\{S_{j}\}$ and nodes $\{\text{CompNode}_{i}\},\{\text{AttnNode}_{jk}\}$ , which are some superset of the SMPC nodes That is, unless there are more SMPC nodes than Cascade nodes. But this is usually never the case in practice, because most SMPC schemes based on additive secret sharing involve 2 or 3 parties.. Then, for each $i$ , each $\mathcal{N}_{s}$ sends the slice $\bm{h}_{s}^{R_{i}}=\bm{h}_{s}[R_{i},:]∈\mathbb{R}^{|R_{i}|× d}$ to CompNode i. Finally, CompNode i adds its received sliced shares to get $\sum_{s=1}^{t}\bm{h}_{s}^{R_{i}}=\bm{h}^{R_{i}}$ . Since each CompNode i has $R_{i}$ -sharded layer $L$ hidden states, we can execute Cascade for all future layers of the LLM with the Cascade nodes.
In this process, by computational indistiguishability, SMPC nodes gain no information about the input, and thus have no direct token access. Likewise, computational indistiguishability ensures that sliced shares $\bm{h}_{s}^{R_{i}}$ give CompNode i no additional information about the input. That is, Cascade nodes no longer have direct access to tokens: information leakage only comes from the usual exposed shards at layers $>L$ . Thus, even though the execution of $\Phi$ on the first $L$ layers may be more expensive than Cascade, token security is improved. To determine the optimal choice of $L$ , further analysis should be conducted on how information leakage from all exposed Cascade shards varies across layers.
Appendix B CompNode Hidden Reversal Analysis on Layers & Llama
In Section 6.2, we described experiments performed on the hidden states of Gemma-2-2B-IT, where a bidirectional-attention model was trained to reverse the sharded hidden states into the original text prompt. In Table 3, we showed that the hiddens are largely secure to this attack for a suitable choice of $c$ and $\alpha$ .
In this section, we first analyze if this is also true for Llama 3.1 8B-Instruct. We run a similar experimental setup as described in Section 6.2, except we use Llama hidden representations, and we also use it as the reversal model; this also therefore tests if increasing the reversal model’s capacity is a suitable method for improving sharded reconstruction. Due to the computational constraints of training with a larger model, we examine this only for $c=8,\alpha=8$ and $c=8,\alpha=12$ . The reconstruction ROUGE-L scores are 0.1718 and 0.1443 respectively, significantly lower than those obtained with the same parameters for Gemma. We leave to future work the interesting question of whether this implies that Llama representations are inherently more resistant to decoding than Gemma representations.
Next, we analyze whether our results hold irrespective of the layer of the model used. We run additional experiments on the hiddens of layers 11 and 21 of Gemma-2-2B-IT. Our results are shown in Table 9. We see that there is no substantial difference in ROUGE-L score as the layer changes.
Table 9: ROUGE-L scores of text reconstruction from the hiddens of various layers of Gemma-2-2B-IT for different $(c,\delta)$ -sharding setups. We see that the reconstruction quality is similar across layers.
| 1 11 21 | 0.4268 0.4627 0.4021 | 0.2218 0.2467 0.2158 |
| --- | --- | --- |
Appendix C Computational Overhead Analysis
In Section 7.1, we claimed that Cascade has little overhead in computational costs compared to vanilla inference. We justify this statement below, by comparing CompNode and AttnNode steps against the vanilla forward pass.
Indeed, most operations performed by CompNodes will treat the (row) token dimension as the batch dimension. In the pre-pass, these are normalization and $Q,K,V$ -projection; and in the post-pass, these are attention value compilation (most of Algorithm 5), $O$ -projection, residual connection, and the MLP block. Except for attention value compilation, these steps all occur in the vanilla pass, so the CompNodes combined will perform the exact same operations as in vanilla inference: there are no extra computations performed.
The only extra operations thus come from (a) attention value compilation (linear weighting of partial attention outputs) by CompNodes in the post-pass, and (b) AttnNode floating point computations which do not appear in the vanilla pass, i.e. expsum s of shards of attention score rows. This is because all other steps of the Cascade self-attention either treat the tokens as batch elements, or involve splitting up matrix multiplication into multiplication of sharded matrices; and the latter is blocked matrix multiplication, which does not inherently change the operations performed.
Now, (a) only involves $\sim H|R_{i}|d$ operations for each CompNode i, since it involves a few steps of elementwise summation and multiplication of $H×|R_{i}|× d$ matrices (after broadcasting). Summing this over all $1≤ i≤\alpha$ gives $\sim\sum_{i}H|R_{i}|d=HNd$ extra operations. Also, (b) only involves $\sim H|S_{j}||S_{k}|$ operations for each AttnNode jk because expsum is done over rows of an $H×|S_{j}|×|S_{k}|$ shard of attention scores. Summing over all $1≤ j,k≤\beta$ , we see this requires $\sim\sum_{j,k}H|S_{j}||S_{k}|=HN^{2}$ extra operations in total. This means the total AttnNode computation overhead is $\sim HN(d+N)$ operations.
Importantly, this is cheaper than most computation-heavy steps in standard inference. Compared to the $\sim HN^{2}d$ operations from multiplication of $H× N× N$ attention scores with $H× N× d$ values, this overhead requires $\sim\frac{1}{N}+\frac{1}{d}$ times as many operations. Since $d$ is often in the hundreds, we can ensure for large $N$ , say $N≥ 256$ , that this ratio is quite small. Furthermore, if $N$ is not large, then the overhead is still limited compared to the $\sim HNd_{emb}d$ operations from $Q,K,V$ -projection, since it requires $\sim\frac{1}{d_{emb}}+\frac{N}{d_{emb}}$ times as many operations and $d_{emb}$ is in the hundreds or thousands. Essentially, the choice of sharding does not significantly affect the total computational overhead, and this overhead is quite modest compared to the computations performed in a vanilla forward pass.
Appendix D Communication Analysis
In Section 7.1, we gave the total communication byte and time overheads for performing a inference on a single layer of an LLM with Cascade. Here, we provide a full justification of these equations. Like in Appendix C, we assume symmetry of $S$ and $T$ sharding, so the superscript $T$ in sharded notation is replaced with $S$ .
Recall that in each layer, there are two communication stages: (A) the CompNodes send sharded query, key, value matrices to the AttnNodes between pre-pass and attention-pass, and (B) the AttnNodes send sharded attention outputs and expsums to the CompNodes between attention-pass and post-pass. We operate under the assumption that all CompNodes synchronize before (A) and all AttnNodes synchronize before (B), so that we can derive an exact expression for communication cost; this makes our communication cost derivation a worst-case analysis. See Appendix F for optimizations that can be made if this assumption is relaxed.
For single-layer inference, in stage (A), CompNode i must send each of the $|R_{i}|$ rows of the $H×|R_{i}|× d$ query matrix $\bm{q}^{R_{i}}$ to some AttnNodes. In particular, for a row index $r∈ R_{i}$ , it sends the row $\bm{q}[:,r,:]$ of $\bm{q}^{R_{i}}$ to all AttnNodes ${}_{j_{r}k}$ with $1≤ k≤\beta$ , where $j_{r}$ is the unique index satisfying $r∈ S_{j_{r}}$ . That is, CompNode i sends each of its $|R_{i}|$ rows to exactly $\beta$ AttnNodes. Since each row contains $Hd$ elements, then CompNode i must send out $\beta|R_{i}|Hd$ elements from sharded query states. A similar analysis shows CompNode i sends out $2\beta|R_{i}|H_{KV}d$ elements from sharded key and value states, so it sends out a total of $\beta|R_{i}|d(H+2H_{KV})$ elements. Summing this over all $i$ and noting $\sum_{i=1}^{\alpha}|R_{i}|=N$ gives the total bytes communicated in (A):
$$
\text{CommBytes}_{A}=\beta Fd(H+2H_{KV})\cdot N.
$$
Assuming perfect parallel transport and uniform bandwidth $B$ across nodes, i.e. all communication overhead comes from CompNode with the most elements to send (plus latency $\tau$ ), the communication time in stage (A) is
$$
\text{CommTime}_{A}=\tau+\frac{\beta Fd(H+2H_{KV})}{B}\cdot\max_{i}|R_{i}|.
$$
Next, in stage (B), each AttnNode jk must send each CompNode i some rows of its partial post-value attention outputs $\bm{u}^{S_{j}S_{k}}$ , partial attention score row maximums $\bm{m}^{S_{j}S_{k}}$ , and partial attention score row subtract-max expsums $\bm{e}^{S_{j}S_{k}}$ . These matrices are of shapes $H×|S_{j}|× d,H×|S_{j}|× 1,H×|S_{j}|× 1$ , respectively, and CompNode i receives $|R_{i}\cap S_{j}|$ out of the $|S_{j}|$ rows from each. This means the total number of elements that AttnNode jk sends to all CompNodes is
$$
(d+2)H\cdot\sum_{i=1}^{\alpha}|R_{i}\cap S_{j}|=(d+2)H\cdot|S_{j}|.
$$
Since $\sum_{j,k=1}^{\beta}|S_{j}|=\beta\sum_{j=1}^{\beta}|S_{j}|=\beta N$ , this means the total number of bytes sent by all $\beta^{2}$ AttnNodes is
$$
\text{CommBytes}_{B}=\beta F(d+2)H\cdot N.
$$
And, again under the parallel transport and uniform bandwidth assumption, the communication time in (B) is
$$
\text{CommTime}_{B}=\tau+\frac{F(d+2)H}{B}\cdot\max_{j}|S_{j}|.
$$
Combining these costs, we obtain the following total communication byte and time overheads for a single layer:
| | CommBytes | $\displaystyle=\beta F(2dH+2dH_{KV}+2H)· N,$ | |
| --- | --- | --- | --- |
Finally, we compute the number of communication rounds per layer. Stage (A) has each of the $\alpha$ CompNodes send results to at most $\beta^{2}$ AttnNodes, which is at most $\alpha\beta^{2}$ rounds. Stage (B) has each of the $\beta^{2}$ AttnNodes send results to at most $\alpha$ CompNodes, which is at most $\alpha\beta^{2}$ rounds. In total, the rounds per layer are bounded above by $2\alpha\beta^{2}$ . This can be quite large, but we can guarantee a tighter upper bound if our scheme involves $(c,\delta)$ -sharding for CompNodes with $m$ -splitting of AttnNodes. Here, $\beta=m\alpha$ since each of the $\alpha$ shards in $\{R_{i}\}_{i=1}^{\alpha}$ is split into $m$ pieces to form $\{S_{j}\}_{j=1}^{\beta}$ . Each CompNode sends results to $m\beta$ AttnNodes, and each AttnNode sends results to $1$ CompNode, so there are $m\alpha\beta+\beta^{2}=2\beta^{2}$ rounds. Essentially, the number of rounds scales linearly with the number of AttnNodes.
Appendix E Sufficiency of $(c,\delta)$ -Sharding
Now, we provide an argument that only considering Layer 0 shards, when $R$ and $S$ sharding both follow a $(c,\delta)$ -sharding scheme with sufficiently large $c,\delta$ , it is intractable for any CompNode to recover the input.
To do this, we will demonstrate a reduction to a variant of the subset sum problem with vectors, which is NP-complete. Our reduction relies on two assumptions (B1), (B2), which we highlight shortly. We begin by fixing $i∈[\alpha]$ , so that $R_{i}=\{ic,ic+1,...,ic+c-1,ic+\delta,ic+\delta+1,...,ic+\delta+c-1,...\}$ . As we mentioned earlier, CompNode i has access to $\bm{m}^{R_{i}S_{1}},\bm{e}^{R_{i}S_{1}},\bm{u}^{R_{i}S_{1}},...,\bm{m}^{R_{%
i}S_{\beta}},\bm{e}^{R_{i}S_{\beta}},\bm{u}^{R_{i}S_{\beta}}$ and $\bm{h}^{R_{i}}$ at Layer 0. Knowledge of the latter is equivalent to knowledge of tokens at indices $R_{i}$ , so we consider leakage from the former triples. In fact, at Layer 0, note each triple $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u}^{R_{i}S_{k}}$ only depends on tokens with indices in $R_{i}\cup S_{k}$ . Since all tokens in $R_{i}$ are known to CompNode i, then this triple only depends on unknown tokens $S_{k}\setminus R_{i}$ . Across all $k$ , these sets are disjoint. Thus, under the following assumption, we can conclude that the reversal problems from the $\beta$ different triples are independent.
(B1) For any $k_{1}≠ k_{2}$ , CompNode i has independent priors on tokens with indices in $S_{k_{1}}$ and tokens with indices in $S_{k_{2}}$ .
Thus, we now only need to consider leakage from one triple $\bm{m}^{R_{i}S_{k}},\bm{e}^{R_{i}S_{k}},\bm{u}^{R_{i}S_{k}}$ . Explicitly, $S_{k}=\{kc,kc+1,...,kc+c-1,kc+\delta,kc+\delta+1,...,kc+\delta+c-1,...\}$ . So, if $k=i$ and $S_{k}=R_{i}$ , then all these shards are functions only of known tokens at indices $R_{i}$ , and no information is leaked to CompNode i. Thus, without of loss of generality, we now assume $k<i$ , as $k>i$ is similar. From now on, when we say a shard depends on tokens, we ignore known tokens.
Recall that $\text{exp}(\bm{a}^{R_{i}S_{k}})\bm{v}^{S_{k}}=\bm{u}^{R_{i}S_{k}}\odot\bm{e}^{%
R_{i}S_{k}}\odot\text{exp}(\bm{m}^{R_{i}S_{k}})$ , and $\text{expsum}(\bm{a}^{R_{i}S_{k}})=\bm{e}^{R_{i}S_{k}}\odot\text{exp}(\bm{m}^{%
R_{i}S_{k}})$ by definition. Thus, this triple reveals the exact same information as the triple $\bm{m}^{R_{i}S_{k}},\text{expsum}(\bm{a}^{R_{i}S_{k}}),\text{exp}(\bm{a}^{R_{i%
}S_{k}})\bm{v}^{S_{k}}$ , which have shapes $(H,|R_{i}|,1),(H,|R_{i}|,1),(H,|R_{i}|,d)$ . We concatenate the last two to form $\bm{b}^{R_{i}S_{k}}$ of shape $(H,|R_{i}|,d+1)$ , so our original triple is equivalent to $\bm{m}^{R_{i}S_{k}},\bm{b}^{R_{i}S_{k}}$ . Now, note that because $S$ and $R$ are both $(c,\delta)$ -sharding and the attention mask is unidirectional, the first $c$ out of $|R_{i}|$ rows of $\bm{a}^{R_{i}S_{k}}$ will have $c$ elements followed by all $-∞$ , with the nonzero $c$ only depending on tokens at indices $kc,...,kc+c-1$ , respectively. The next $c$ rows have $2c$ elements that depend on tokens at indices $kc,...,kc+c-1,kc+\delta,kc+\delta+1,...,kc+\delta+c-1$ , followed by all $-∞$ . A similar pattern holds for each next $c$ rows. Thus, the $|R_{i}|$ rows of $\bm{b}^{R_{i}S_{k}}$ follow a similar token dependence pattern: the first $c$ depend on tokens $kc,...,kc+c-1$ , the next $c$ depend on tokens $kc,...,kc+c-1,kc+\delta,kc+\delta+1,...,kc+\delta+c-1$ , and so on.
Now, for each position $x∈[N]$ , we define a list of $V$ possible concatenations $[k_{x},v_{x}]∈\mathbb{R}^{H× 2d}$ , where $k_{x}$ denotes the $x$ th row across all heads of the layer 0 query states, and likewise for $v_{x}$ . Denote this list as $\mathcal{A}_{x}$ , which is distinct for different values of $x$ due to the effects of positional embeddings; it is indexed in the same order as the vocabulary. For each $r∈ R_{i}$ and position $x∈[N]$ with $x≤ r_{i}$ , we denote $\mathcal{B}_{r,x}$ as the list of $(d+1)$ -dimensional vectors whose first entry is $\text{exp}(\bm{q}[:,r,:]k_{x}^{T})$ and last $d$ entries are $\text{exp}\left(\bm{q}[:,r,:]k_{x}^{T}\right)v_{x}$ , over all $[k_{x},v_{x}]∈\mathcal{A}_{x}$ . Note this is also of size $V$ , and is indexed in the same order as the vocabulary. It is simple for a node to compute $\mathcal{B}_{r,x}$ for all $r∈ R_{i}$ and $x≤ r_{i}$ , as they know $\bm{q}[:,r,:]$ , and can directly compute and iterate through $\mathcal{A}_{x}$ .
The key point, building on our observation from earlier, is that $(d+1)$ -dimensional rows of $\bm{b}^{R_{i}S_{k}}$ are sums of vectors from sets $\mathcal{B}_{r,x}$ . In particular, there are $i_{1},...,i_{c}$ such that for all $1≤ r≤ c$ , the $r$ th row equals $\mathcal{B}_{r,kc}[i_{1}]+...+\mathcal{B}_{r,kc+c-1}[i_{c}]$ . Then, there are $i_{c+1},...,i_{2c}$ such that for all $c+1≤ r≤ 2c$ , the $r$ th row equals $\mathcal{B}_{r,kc}[i_{1}]+...+\mathcal{B}_{r,kc+c-1}[i_{c}]+\mathcal{B}_{r,%
kc+\delta}[i_{c+1}]+...+\mathcal{B}_{r,kc+\delta+c-1}[i_{2c}]$ . A similar pattern holds for $2c+1≤ r≤ 3c$ , and so on. The final indices $i_{1},i_{2},...$ correspond to the positions of input tokens in the vocabulary, for all tokens with indices in $S_{k}$ .
Therefore, reversal of the input from $\bm{b}^{R_{i}S_{k}}$ alone is at least as difficult as the following problem: given $c$ different sets, each containing $V$ vectors derived from token and position embeddings, select some vector from each set so that they sum to a given vector. This is a variant of the subset sum problem for vectors, which is computationally intractable for large enough $V,c$ . Particularly, since $V$ is in the hundreds of thousands, this is likely already intractable for $c≥ 8$ .
Finally, we consider additional leakage from $\bm{m}^{R_{i}S_{k}}$ . In the worst case, this reveals at most one of the tokens with indices in $S_{k}$ for each $k$ , as a maximum of a row of $\bm{a}^{R_{i}S_{k}}$ reveals at most one element of the row, with upper bound constraints placed on the other elements. Given the large number of possibilities for these elements, it is unlikely that the upper bound constraints provide much information, which we incorporate in the assumption below. Therefore, leakage from $\bm{m}^{R_{i}S_{k}}$ is at worst equivalent to effectively reducing the parameter $c$ above by one, since we have one less set to select a vector from. This means if $c≥ 8$ was secure for reversal from $\bm{b}^{R_{i}S_{k}}$ , we need $c≥ 9$ to ensure security for reversal from $\bm{b}^{R_{i}S_{k}},\bm{m}^{R_{i}S_{k}}$ . Note that the variant of the subset sum problem for reversal from $\bm{b}^{R_{i}S_{k}}$ , together with the inequality constraints from $\bm{m}^{R_{i}S_{k}}$ leakage, form a linear program. So, we have implicitly used the following assumption to claim security.
(B2) The linear program described above is computationally intractable.
Therefore, assuming choices of $c,\delta$ which satisfy (B1) and (B2), we have shown Cascade is completely secure for any CompNode only executing Layer 0. In practice, these assumptions may not always hold. In fact, learning-based attacks attempt to violate the independence assumption of (B1), through token-infilling between gaps of a $(c,\delta)$ -sequence. Furthermore, even as the subset sum problem is NP-complete, this does not mean that the particular choice of sets in our setup are not susceptible to, say, an approximation algorithm. To truly test (B2), future work should explicitly enumerate the linear program and test the efficacy of state-of-the-art solvers on it, perhaps with token-infilling priors integrated. Nevertheless, the above explicit reduction shows that for large enough $c,\delta$ , it is quite likely that (B1) and (B2) nearly hold.
Appendix F Cost Optimizations
In Section 5, we gave a high-level overview of Cascade, and deferred discussions about optimization. Here, we discuss a few cost and communication optimizations.
Caching After a new token is generated in Cascade, the CompNode holding that token will send it back to one of the existing CompNodes, and single-token generation will repeat to get the next token. To speed up generation after the first new token, the CompNodes and AttnNodes can store their partial intermediate states, and only the $1$ CompNode and $\beta$ AttnNodes associated with the most recent token will need to participate in the single-token generation: this means KV-caching naturally extends to Cascade. Formally, suppose $n$ is the index of the most recently generated token, and it belongs to the hidden shard $R_{i}$ and the query shard $S_{j}$ . Only CompNode i needs to perform new computation in generating the $(n+1)$ st token, along with each AttnNode jk for $1≤ k≤\beta$ : this is because only these AttnNodes require the $n$ th query row. Furthermore, these $\beta+1$ nodes, having stored intermediate results from previous forward passes, can avoid repeat computation of attention scores and earlier hidden states. Essentially, this results in token-sharded KV-caching.
Symmetry Reduction We see that AttnNode jk and AttnNode kj actually have the exact same information in the worst-case: they both have access to indices $S_{j}\cup S_{k}$ . Thus, at no loss of security, we can combine AttnNode jk and AttnNode kj into one node, thereby approximately halving the number of AttnNodes required, and reducing communication byte overhead.
Synchronization A key assumption in our communication analysis from Appendix D was that nodes synchronize between stages. That is, AttnNodes wait until they all finish before sending information to CompNodes in parallel; and likewise for the CompNodes sending information to AttnNodes. But in practice, depending on the sharding scheme, synchronization is not necessary; and relaxing it can allow some nodes to proceed earlier than others. For instance, in a sharding scheme where CompNode 1 holds only the first $k$ tokens, because the first $k$ logits do not depend at all on tokens $k+1,...,n$ in a unidirectional model, then CompNode 1 can proceed through all its forward passes without waiting for any information from other nodes. Future work could analyze the tradeoff between such synchronization relaxations, which are not possible with schemes like $(c,\delta)$ -sharding, and token security.
Appendix G Collusion Resistance
How many nodes in Cascade must share information to break the security of the scheme? The acceptable level of collusion depends on the sharding scheme and the desired hidden state security. Formally, let $\mathcal{N}$ be the set of all $\alpha+\beta^{2}$ nodes. For each node $\mathbf{n}∈\mathcal{N}$ , denote $t(\mathbf{n})$ as the set of (sequence dimension) indices of hidden states that $\mathbf{n}$ directly has access to. This is explicitly
$$
t(\text{CompNode}_{i})=R_{i},\quad t(\text{AttnNode}_{jk})=S_{j}\cup T_{k}
$$
because CompNode i directly receives hidden states with indices in $R_{i}$ , and AttnNode jk directly receives $Q,K,V$ -projections of hidden states with indices in $S_{j}$ and $T_{k}$ .
When nodes share information, they gain access to the union of their shards. Suppose a subset $\mathcal{S}⊂eq\mathcal{N}$ of nodes jointly colludes. Then they each gain access to hidden states with indices in the set $\cup_{\mathbf{n}∈\mathcal{S}}t(\mathbf{n})$ . This can be used to reveal a significant portion of the prompt if $\mathcal{S}$ contains many nodes, so in general, Cascade is not fully collusion-resistant. For example, if all $\alpha$ CompNodes collude, then they gain access to all hidden states with indices in $\cup_{i=1}^{n}t(\text{CompNode}_{i})=\cup_{i=1}^{n}R_{i}=[N]$ , so the entire prompt is revealed by reversing layer $0 0$ hidden states into tokens.
Collusion also raises security concerns in vocab-matching. Recall the generalized vocab-matching attack from Section 4 can be used to break security if gaps between non-consecutive hidden states in each shard are not sufficiently large. From our prior discussion in Section 6.2 and Theorem 6.2, we know that non-consecutive hidden states held by a node must be $>\rho$ positions apart to prevent the attack, where $\rho$ is the vocab-matching threshold. Thus, it may be the case that much fewer than $\alpha$ CompNodes need to collude to break security: even though the union set of their shards may be much smaller than $[N]$ (i.e. reveal little of the actual prompt), it could contain many pairs of non-consecutive hidden states $≤\rho$ positions apart, and the vocab-matching attack would reverse many unknown intermediate hidden states. This means the choice of vocab-matching threshold also significantly affects the collusion limit.
Finally, we note a potential solution to security concerns from collusion: further splitting of shards assigned to nodes, akin to $m$ -splitting of AttnNodes in Section 6.3. By decreasing the number of hidden states each individual node has access to, we increase the number of nodes that need to collude to break security. For instance, in the extreme case where each CompNode has access to $1$ token and the prompt contains sufficiently many filler words, a large number of CompNodes could collude and still be unable to extract private user data from the prompt. This is also resistant to the vocab-matching attack as long as CompNodes with tokens within $\rho$ positions of each other do not collude.