Partitioning for Intrinsic Model Inversion Resistance in Collaborative Inference

Conference: ICML 2026
arXiv: 2506.15412
Code: https://github.com/GoldenPartitionZone/GoldenPartitionZone
Area: AI Security / Collaborative Inference / Model Inversion Defense
Keywords: Model Inversion Attack, Collaborative Inference, Partitioning Strategy, Information Entropy, Label Smoothing

TL;DR

Departing from traditional defense paradigms that "add noise or masks to shallow intermediate representations," this paper proves from an information-theoretic perspective that in edge-cloud collaborative inference, the model should be partitioned at the layer where the "feature → decision" mutation occurs (termed the Golden Partition Zone, GPZ). The intra-class mean square radius \(R_c^2\) is identified as the key variable for localizing the GPZ and can be actively contracted through label smoothing training dynamics.

Background & Motivation

Background: Collaborative Inference (CI) partitions deep networks into an edge component \(f_{\text{edge}}\) and a cloud component \(f_{\text{cloud}}\). The edge device uploads the intermediate representation \(z = f_{\text{edge}}(x)\) to the cloud. This deployment mode is widely used in drones, IoT, and private cloud inference. However, Model Inversion Attacks (MIA) can train a generator \(g \approx f_{\text{edge}}^{-1}\) to reconstruct the original input \(x\) from \(z\), leading to sample-level privacy leakage.

Limitations of Prior Work: Existing MIA defenses almost exclusively focus on perturbing shallow \(z\) (adding noise, masking, bottleneck layers, homomorphic encryption, etc.). The cost is either a sacrifice in downstream accuracy or the introduction of additional computational overhead, essentially struggling with the "privacy-utility" trade-off.

Key Challenge: The authors argue that the question is framed incorrectly. The real question should be: at which layer should the network be partitioned so that \(z\) naturally and irreversibly loses input information before transmission? This shifts defense from a "post-hoc patch" to the "partitioning location" itself.

Goal: (1) Theoretically characterize the relationship between "where to partition" and MIA difficulty; (2) Provide a computable and monitorable layer-wise metric to allow users to actively locate the optimal split point; (3) Actively shape this metric during the training phase to enter the inversion-resistant zone earlier.

Key Insight: Conventional intuition suggests that "deeper is safer." However, in ViT, even partitioning at the last layer allows patch tokens to retain detailed information for each sample, resulting in successful inversion. In IR-152/ResNet-50 with residual connections, increasing depth actually causes \(I(X; Z)\) to decay more slowly due to skip connections. These counter-examples force a shift in focus from "depth" to the "essential mutation of representation form."

Core Idea: The transition from "feature-level → decision-level" representation is used as a necessary condition for intrinsic defense. The intra-class mean square radius \(R_c^2 = \frac{1}{N_c} \sum_{i:y_i=c} \|z_i - \mu_c\|^2\) is employed as the sole computable proxy variable to locate this transition zone (GPZ). Furthermore, \(R_c^2\) is actively contracted during training via label smoothing.

Method

Overall Architecture

The paper follows a chain of "Theory → Metric → Training Dynamics → Experimental Validation." First, it derives the lower bound of \(H(X \mid Z)\), revealing that this bound is dominated by the global variance of the feature layer \(\sigma_{\text{feat}}^2\) and becomes dominated by the much smaller intra-class radius \(R_c^2\) at the decision layer, causing the lower bound to jump upward. Next, \(R_c^2\) is refined into a usable layer-wise probe. Finally, the training dynamics of \(R_c^2\) are backward-regulated via label distribution (termed the Neural Vortex), making \(R_c^2\) at the decision layer smaller and more resistant to inversion.

Key Designs

1. GPZ Localization Criterion: \(R_c^2\) Probe Derived from \(H(X\mid Z)\) Lower Bound:

   • Function: Calculates a scalar for each layer to find the layer where the representation undergoes a "decisional" mutation.
   • Mechanism: By treating \(z\) as a continuous variable and using the maximum entropy principle with determinant-trace inequalities, the upper bounds for differential entropy at the feature and decision levels are obtained. At the feature level, \(h(Z_{\text{feat}}) \le \frac{d}{2}\ln(2\pi e \sigma_{\text{feat}}^2)\), primarily depending on dimension \(d\) and global variance \(\sigma_{\text{feat}}^2\). At the decision level, after conditioning by class, \(h(Z_{\text{dec}} \mid Y=c) \le \frac{D}{2}\ln(2\pi e R_c^2/D)\), where the decisive variance term switches from "global variance" to "intra-class mean square radius \(R_c^2\)." Substituting these into the mutual information identity yields \(H(X\mid Z) \ge H(X) - h(Z) - \kappa_\Delta\). Thus, when the representation enters the decision zone, \(h(Z)\) drops significantly, and the lower bound jumps. In practice, one only needs to sweep \(R_c^2\) across candidate layers without estimating mutual information.
   • Design Motivation: To transform the engineering problem of "where to cut" into an observable and automatable problem of "at which layer \(R_c^2\) shows an abrupt drop," bypassing the high variance and cost of MI estimators (like MINE).

2. Neural Vortex: Training Dynamics for Active \(R_c^2\) Contraction via Label Smoothing:

   • Function: Allows \(R_c^2\) at the decision layer to continue being "pulled down" in the late training stages, further raising the GPZ lower bound.
   • Mechanism: Writing the one-step backpropagation as \(\Delta R_c^2 = -\frac{2\gamma}{N_c} \sum_{i\in c} (z_i - \mu_c)^\top \tilde g_i\) and substituting \(\tilde g_i = J_i^\top (p_i - y_i)\), it can be decomposed into a "correct class pulling term" \((p_{ic}-1)T_{\text{corr},i}\) and an "incorrect class interference term" \(\sum_{k\ne c} p_{ik} T_{k,i}\). Under one-hot supervision, the pulling term approaches zero as \(p_{ic} \to 1\), stopping the decline of \(R_c^2\). With Label Smoothing (LS), the correct class coefficient becomes \((p_{ic}-1+\alpha)\). Once \(p_{ic} > 1-\alpha\), the sign flips, \(T_{\text{corr},i}\) reverses geometrically, and \(\Delta R_c^2 < 0\) is maintained, continuously tightening the intra-class point cloud. The authors name this counter-intuitive coupling of "output entropy increase + intermediate entropy decrease" the Neural Vortex.
   • Design Motivation: Unlike simple IB regularization or post-hoc observation of neural collapse, this actively regulates from the perspective of training dynamics. Moreover, this regulation results in almost no drop in downstream accuracy (LS+ even slightly improves it in experiments), making it a free lunch.

3. Bi-directional Stress Testing for Decision-Layer Inversion Resistance: Entropy Enhancement + Inversion Model Enhancement:

   • Function: Verifies that GPZ is not a "paper tiger" sustained by weak attacks, confirming that even with "buffs" for the attacker, the decision layer remains significantly difficult to invert.
   • Mechanism: On the representation side, transmitted \(z\) is enriched using FFT residuals/concatenation, global normalization, and small NN modules with dropout. On the attacker side, Multi-Head Attention, Attention-as-Conv, SE, LSK, and MSCA are progressively inserted between deconvolutional blocks, following the principle of "shallow weak attention → deep strong decoupling." Additionally, inverse IR-152 residual blocks are tested. Both types of enhancements serve as "stress tests."
   • Design Motivation: Previous defenses often fail against stronger attack models. This paper uses symmetric bi-directional enhancement to solidify the conclusion of "decision-layer inversion resistance" by maintaining a significant gap even against enhanced attacks.

Loss & Training

Target models were trained on 7 datasets (\(64\times 64\)) including CIFAR-10, FaceScrub, and KMNIST, using three label distributions: one-hot, LS+ (\(\alpha=0.3\)), and LS- (\(\alpha=-0.05\), reverse smoothing as a control). Inversion models utilized deconvolutional backbones from Yang et al. (2019) and Zhang et al. (2023). Evaluation metrics included MSE / PSNR / SSIM / LPIPS (AlexNet default weights), with MSE \(<0.02\) as the empirical threshold for "high-fidelity reconstruction."

Key Experimental Results

Main Results: Impact of Representation Level on MIA Difficulty (IR-152, CIFAR-10)

Partition Point	Representation Type	MSE (Test)	PSNR (Test)	Reconstruction Readable?
Block 40	Feature-level	0.018	22.17	Yes
Block 48	Residual accumulation, still feature-level	\(<0.02\)	\(\approx 22\)	Yes
Block 50	Decision-level (GPZ)	0.057	17.22	No
Block 30→39 (VGG19)	Feature→Decision mutation	0.066 → 0.137	—	Significant degradation at mutation

IR-152 compresses spatial resolution to \(4\times 4\) at Block 49, where the representation suddenly becomes decisional, and MSE jumps from \(<0.02\) to \(0.057\). This is the source of the "GPZ has 4× higher MSE than shallow cuts" claim. The paper also notes that ViT never undergoes this transformation because it retains 256 patch tokens throughout, preventing the formation of a GPZ regardless of the partition point.

Ablation Study: Stability of GPZ after Representation/Attacker Enhancement (IR-152, Block 50 vs Block 40)

Configuration	Block 50 (GPZ) MSE	Block 40 (Feature) MSE	GPZ Relative Disadvantage Narrowed?
Baseline Attack	0.057	0.018	—
Representation side: Normalize+Dropout-Concat	0.052	0.014	No (Gap maintained ~3.7×)
Attacker side: Attention-as-Conv+SE+LSK+MSCA	0.051	—	No
Hybrid Attack/Rep + Inversion-IR152	0.049–0.052	0.012	No (Gap still ~4×)

Key Findings

• True intrinsic defense is not about "partitioning deep," but "partitioning after the representation undergoes a morphological transformation." Residual connections and ViT patch tokens delay or erase this transformation, making them largely ineffective against MIA.
• Decision-level representations maintain an average 66% inversion resistance advantage over feature-level ones even under dual buffs to the attacker and representation, proving GPZ is not a fragile attack artifact.
• Data distribution determines the GPZ location: On FaceScrub, the GPZ is earlier and narrower; on KMNIST, feature extraction continues deeper due to many zero pixels, pushing the GPZ to around Block 43. This aligns with the joint effect of \(H(X)\) and \(R_c^2\) in the lower bound.
• Inversion models trained on MNIST to invert KMNIST bias towards "0," while those trained on EMNIST bias towards "D," suggesting that reconstructions after GPZ no longer reflect private content but rather priors from auxiliary data, further confirming the stripping of private information.
• VGG19 is the most cost-effective for deployment: only 2.5% of parameters remain at the edge to reach GPZ, whereas IR-152 requires over 78%. Moving from depth-26 to depth-30 in VGG adds almost no latency but halves the transmission payload.

Highlights & Insights

• Treating "where to partition" as a defense dimension is more upstream and fundamental than "how to perturb," avoiding the constant push-pull on the privacy-utility curve. It feels like "changing the question instead of the answer."
• Seamlessly bridging an information-theoretic bound with an engineering-computable metric (\(R_c^2\)) makes the theory both explanatory and actionable. The "theory provides the pointer, pointer sweeps the layers" combination is extremely practical.
• The Neural Vortex section is an "Aha!" moment: output entropy increase (label smoothing flattening softmax) paradoxically leads to intermediate entropy decrease (tighter intra-class clouds), a contradiction naturally resolved by the sign flip of \((p_{ic}-1+\alpha)\). Such detailed training dynamics analysis is rare in MIA literature.
• Transferability: The \(R_c^2\) probe could be applied to any scenario aiming to make intermediate representations harder to invert, such as gradient leakage defense in Federated Learning or state hiding in Multi-Party Computation. The idea is to use \(R_c^2\) to find the intrinsic transition layer and then apply selective perturbations.

Limitations & Future Work

• Primarily validated on vision models. Whether GPZ exists in text and sequence models depends on whether patch/tokens retain sample-level information in deep layers; the "no GPZ in ViT" result suggests Transformers may be harder for forming clear transition zones.
• For minimalist data (MNIST/KMNIST), the GPZ moves significantly backward and narrows, implying that low-entropy edge cases still require additional perturbation reinforcement.
• The \(R_c^2\) probe requires available labels and clear categories; in self-supervised representation or multi-task output scenarios, the definition of a "class" is not obvious and requires extension to prototype or cluster perspectives.
• Synergy with active defenses (like noisy IB, HE) has not been systematically compared. A composite defense of GPZ partitioning + lightweight perturbation might be superior to either alone.

Related Work & Insights

• vs Information Bottleneck methods (Wang et al., 2021; Duan et al., 2023): IB explicitly penalizes \(I(X;Z)\) but requires MI estimation and significantly degrades accuracy; Ours achieves equivalent or stronger intrinsic defense via split-point selection + LS without changing the loss or estimating MI.
• vs Neural Collapse (Papyan et al., 2020): Neural Collapse is a post-hoc observation of geometric finality; Neural Vortex provides a controllable dynamic description during training and links it to the privacy lower bound.
• vs Shallow Perturbation Defenses (Wang et al., 2022; Ding et al., 2024): Shallow noise/masking treats the symptom; pushing the split point to the decision side treats the cause. The two are actually stackable.

Rating

• Novelty: TBD
• Experimental Thoroughness: TBD
• Writing Quality: TBD
• Value: TBD

Related Papers

• [ICML 2026] Local Hessian Spectral Filtering for Robust Intrinsic Dimension Estimation
• [ICML 2026] Amortized Simulation-Based Inference in Generalized Bayes via Neural Posterior Estimation
• [ICML 2026] How Does Bayesian Sampling Help Membership Inference Attacks?
• [ICML 2026] Correcting Split Selection in Online Decision Trees via Anytime-Valid Inference
• [ICML 2026] Local and Mixing-Based Algorithms for Gaussian Graphical Model Selection from Glauber Dynamics