PaperlayerView exact revision on arXiv

Source-linked AI summary

D-FROST: Decentralized Federated pRompt-tuning via Optimal tranSporT for Non-IID and Imbalanced Data

Quan Minh Nguyen, Hoang M. Ngo, Trong Nghia Hoang, My T. Thai

arXiv:2609.01802v1cs.LG

TL;DR

Decentralized prompt tuning lacks a suitable way to coordinate heterogeneous, unordered prompt sets while preserving consensus and progress toward a shared objective. D-FROST models prompts as Wasserstein measures and merges neighborhood prompts through optimal transport, with theory and experiments supporting its effectiveness under heterogeneous data.

  • Problem

    Prompt tuning in DFL lacks convergence analysis and suitable aggregation for heterogeneous, index-wise misaligned prompt sets.

  • Method

    D-FROST formulates decentralized prompt tuning over prompt measures and uses OT-based merging to combine neighborhood prompts into compact representative sets.

  • Results

    Across eight datasets, D-FROST consistently outperforms decentralized baselines, with the largest gains under extreme non-IID data.

  • Takeaways & Limitations

    Wasserstein consensus and stationarity analysis provide theoretical guarantees for decentralized prompt tuning, while experiments support D-FROST under heterogeneous client data.

  • Takeaways & Limitations

    The analysis assumes Wasserstein smoothness and bounded local Wasserstein gradient variance.

Abstract

from arXiv · show

Prompt tuning provides a parameter-efficient way to adapt foundation models (FMs) by freezing the pretrained backbone and updating only a small set of learnable prompts. This property makes prompt tuning especially suitable for decentralized federated learning (DFL), where exchanging full-model updates can be prohibitively expensive. However, prompt tuning in DFL introduces new challenges. Prompt sets learned from heterogeneous local data may not be index-wise aligned, making standard decentralized averaging unsuitable. In addition, the algorithm should be theoretically guaranteed to achieve consensus and make progress toward the shared objective. In this work, we provide the first study of prompt tuning in DFL. We formulate decentralized prompt tuning as a Wasserstein-based optimization problem over prompt measures, which captures the set-valued structure of prompts. We then propose D-FROST, an optimal-transport-based (OT-based) decentralized prompt-tuning algorithm that merges neighborhood prompts into compact representative prompt sets through transportation-based matching. We further analyze D-FROST by bounding the Wasserstein consensus error across clients, and establishing convergence of the network-level prompt barycenter to a neighborhood of stationarity. Experiments under heterogeneous client data demonstrate the effectiveness of D-FROST for decentralized prompt tuning.

1 Introductions

Prompt tuning reduces federated adaptation costs by updating only learnable prompts, but decentralized settings require methods that handle heterogeneous, unordered prompt sets. D-FROST addresses this with OT-based neighborhood merging and provides convergence analysis and empirical evaluation.

  • Prompt tuning freezes the pretrained backbone and updates only a small set of learnable prompts, reducing the parameters exchanged in federated adaptation.This makes prompt tuning attractive when communication costs are a bottleneck.
  • DFL clients communicate only with graph neighbors, creating a server-free setting in which local prompt states must reach consensus and progress toward a shared objective.Existing full-model convergence analyses rely on coordinate-aligned vectors and Euclidean averaging.
  • Heterogeneous client data can produce index-wise misaligned prompts, so direct index-wise averaging may combine unrelated prompt directions.This is the central prompt-specific challenge motivating transportation-based matching.
  • D-FROST updates local prompts and then applies an OT-based Merge function to summarize neighborhood prompts into a compact representative prompt set.The merge matches prompts according to their geometry in prompt embedding space rather than their indices.
  • The convergence analysis bounds Wasserstein consensus error and shows the network-level prompt barycenter converges to a neighborhood of stationarity.The analysis also establishes stability of the local OT-based Merge solver as inner OT steps increase.
  • Experiments across eight diverse vision datasets show D-FROST is effective under data imbalance and extremely heterogeneous decentralized settings.It is evaluated against decentralized federated prompt-tuning baselines based on existing DFL techniques.

2 Related Works

Prior work reduces federated communication through prompt aggregation, but centralized methods and coordinate-aligned decentralized averaging do not fit unordered, misaligned prompt sets. D-FROST replaces parameter averaging with OT-based merging for Wasserstein consensus.

  • Prompt-tuning methods such as FedPrompt reduce communication and storage by aggregating prompts instead of full models.PFPT uses probabilistic prompt aggregation for non-IID and imbalanced data.
  • These federated prompt methods use centralized server aggregation and do not address decentralized communication among graph neighbors.The related-work gap is both architectural and communication-based.
  • Classical decentralized learning methods typically combine local optimization with neighbor averaging under coordinate-aligned model parameters.Decentralized SGD analyses establish convergence under suitable mixing conditions, and DFedAvgM adapts multiple local SGD iterations to DFL.
  • Decentralized prompt tuning violates the coordinate-alignment assumption because client prompts form unordered, potentially misaligned sets.Index-wise averaging can therefore combine mismatched prompt directions.
  • D-FROST replaces parameter averaging with OT-based merging to achieve Wasserstein consensus for prompt states.The replacement matches the set-valued structure of prompts rather than assuming fixed indices.

3 Preliminaries

DFL optimizes a shared objective over a graph while clients retain heterogeneous private data and exchange states with neighbors. For prompt tuning, client states are unordered prompt measures, making Wasserstein distance a natural matching geometry.

  • 3.1 Decentralized Federated Learning (DFL): DFL connects clients through a communication graph, with each client holding private data whose distribution may differ from other clients.The graph defines which client pairs can communicate directly.
  • 3.1 Decentralized Federated Learning (DFL): A mixing matrix represents topology, and its connectivity factor ρ captures information mixing and consensus strength.For connected graphs, smaller ρ indicates faster information mixing and stronger consensus.
  • 3.1 Decentralized Federated Learning (DFL): Each decentralized round performs a local update followed by neighbor exchange and aggregation; classical full-model Merge uses mixing-matrix-weighted parameter averaging.Prompt tuning changes both the client state and the Merge operation.
  • 3.2 Measure Space and Wasserstein Distance: Prompt-tuning clients maintain prompt sets rather than full parameter vectors, and these sets can be represented as distributions over prompt embedding space.This treats prompts as unordered support points rather than indexed coordinates.
  • 3.2 Measure Space and Wasserstein Distance: P2(Rd) is the space of probability measures on the prompt embedding space with finite second moment.It supplies the measure space used to model prompt states.
  • 3.2 Measure Space and Wasserstein Distance: A coupling between two measures is a joint measure whose marginals are the compared measures, and Wasserstein distance minimizes transport cost over such couplings.The discrete form uses a transport matrix between empirical support points.
  • 3.2 Measure Space and Wasserstein Distance: Wasserstein distance compares empirical distributions by optimizing over possible support-point matchings instead of assuming a fixed ordering.This property suits prompt sets that can be misaligned across clients.

4 Decentralized Wasserstein Prompt Tuning

D-FROST formulates decentralized prompt tuning as optimization over Wasserstein prompt measures and alternates local adaptation with OT-based neighbor merging. Its merge solver transports neighborhood prompts into compact representatives while supporting consensus and shared-objective progress.

  • 4 Decentralized Wasserstein Prompt Tuning: The decentralized Wasserstein prompt-tuning section defines the setup, a shared prompt-measure objective, and an OT-based algorithm for approximate optimization.The formulation keeps the pretrained backbone fixed while adapting local prompt sets.
  • 4.1 Setup: Each client maintains a local prompt set represented as an empirical measure, treating prompts as unordered support points in embedding space.Neighboring sets can become misaligned because clients update using different data distributions.
  • 4.2 Decentralized Wasserstein Prompt Tuning Objective: The global objective minimizes average client loss over a shared prompt measure rather than separate personalized objectives.This preserves the shared-state structure of classical DFL while replacing Euclidean consensus with Wasserstein consensus.
  • 4.2 Decentralized Wasserstein Prompt Tuning Objective: Local empirical prompt measures approximate the shared state, while their Wasserstein barycenter summarizes collective network behavior.Consensus is measured by the Wasserstein distance of local measures from the network-level barycenter, with smaller ε(t) indicating tighter concentration.
  • 4.3 OT-Based Decentralized Algorithm: Each D-FROST communication round combines a local prompt update, which improves local loss, with neighbor merging, which promotes Wasserstein consensus.The algorithm targets shared-objective progress while maintaining consensus among local prompt measures.
  • 4.3 OT-Based Decentralized Algorithm: Clients exchange prompt states with neighbors to form a neighborhood collection, then summarize it into a compact prompt set to limit storage and computation.The compact set approximates the received collection under transportation-based geometry.
  • 4.3 OT-Based Decentralized Algorithm: The OT merge solves a regularized matching problem in which neighborhood prompts receive uniform source mass and representative prompts can receive different mass according to geometry.The objective combines spatial matching, entropy regularization, and L2 regularization.
  • 4.3 OT-Based Decentralized Algorithm: Alternating transport and barycenter updates softly assign neighborhood prompts and move representatives toward dominant directions without index-wise averaging.The transport softness is controlled by ε, while λ controls regularization strength; the solver uses closed-form block updates.

5 Theoretical Analysis of D-FROST

D-FROST is theoretically analyzed through local OT-merge stability, Wasserstein consensus control, and convergence of the network-level prompt barycenter toward stationarity.

  • Algorithm and assumptions: D-FROST alternates transport and barycenter updates to merge neighborhood prompt measures into compact representative prompt sets.The algorithm combines local prompt adaptation with OT-based neighbor merging, while its analysis assumes smooth losses, bounded gradient variance and prompt support, and graph mixing.
  • Local merge stability: After S alternating OT steps, at least one inner iterate has discrepancy bounded by O(1/S), making the local merge progressively more stable.The discrepancy is Δ(s) = ∥P(s) − P(s−1)∥F + ∥Φ(s) − Φ(s−1)∥F, and the bound depends on the initial objective gap and µ = min(ε, λ/σ2).
  • Global consensus: D-FROST controls Wasserstein consensus error by keeping local prompt measures close to a network-level barycenter.The consensus analysis is stated under Assumptions 1–4 and gives an expected network consensus bound.
  • Stationarity: The network-level prompt barycenter converges to a neighborhood of Wasserstein stationarity under the stated smoothness, variance, support, and mixing assumptions.The stationarity result requires η ≤ 1/L and α < 1; its bound contains optimization, local-variance and consensus, and OT-merge approximation terms.
  • Global consensus: The consensus neighborhood reflects OT-merge approximation error and heterogeneous local updates, while better-connected graphs yield tighter neighborhoods through smaller mixing factor ρ.The approximation term decreases with more accurate finite-prompt merging, and graph connectivity affects the denominator of the bound.
  • Stationarity: More accurate local merging, achieved by increasing prompt budget n or OT steps S, leads to a tighter stationarity neighborhood.The merge approximation error is attributed to finite prompt budget and finite OT iterations.

6 Experiments

Experiments evaluate D-FROST under heterogeneous data, varied topologies, and adverse network conditions. Across these settings, D-FROST consistently achieves strong accuracy, convergence, and robustness.

  • Performance and Convergence: D-FROST reaches 81.95%, 79.36%, and 70.46% accuracy across FiveDataset splits, improving over baselines by 6.67, 8.48, and 18.57 points.The performance gap widens under greater heterogeneity, especially in the extreme non-IID split.
  • Performance and Convergence: D-FROST separates from baselines within a few rounds and converges at O(1/T) to a neighborhood determined by merge error and gradient variance.Index-wise averaging has larger consensus error under severe data skew and slower convergence.
  • Topology-aware performance: D-FROST outperforms all baselines across Ring, Grid, Erdős-Rényi, Regular, and fully connected topologies.Performance improves for all methods as connectivity increases, while D-FROST retains a substantial advantage.
  • Robustness to network conditions: With link-drop probability p = 0.4, D-FROST declines only mildly from 79.36% to 77.29% accuracy.The stress tests vary unreliable links, network scale, and partial participation.
  • Robustness to network conditions: With 10 clients active per round, D-FROST retains 64.88% accuracy, nearly 12 points above DFedAvgM.With 100 clients, D-FROST leads the strongest baseline by 11%.

7 Conclusion

The paper frames decentralized prompt tuning as optimization over prompt sets rather than aligned parameter vectors. D-FROST uses optimal transport to merge neighborhood prompts and provides theoretical and empirical support for this formulation.

  • Theory and conclusion: The analysis bounds Wasserstein consensus error and shows that the network-level prompt barycenter converges to a neighborhood of stationarity.The local OT solver is also shown to become stable as inner OT steps increase.
  • Empirical conclusion: Across eight datasets, D-FROST consistently outperforms decentralized baselines, with the largest gains under extreme non-IID data.The conclusion reports effectiveness across heterogeneous decentralized prompt-tuning settings.
  • Prompt tuning formulation: Prompt tuning freezes the pretrained backbone and optimizes only a small set of prompt parameters, reducing trainable parameters for decentralized learning.The local prompt-tuning objective is optimized while backbone parameters remain fixed.
  • Prompt-set aggregation: Index-wise averaging is unsuitable because heterogeneous clients can learn prompt sets without direct positional correspondence.Local prompts may drift toward client-specific directions, unlike coordinate-aligned full-model parameters.

A.2 Decentralized Federated Learning Baselines

The baselines adapt standard decentralized optimization methods to frozen-backbone prompt tuning, while PFPT illustrates a centralized alternative whose naive decentralized extension causes prompt-pool growth.

  • Baseline protocol: D-PSGD-PT, DFedAvgM-PT, and DFedSAM-PT update and exchange only prompt tokens and the classification head while keeping ViT-B/32 frozen.These methods use the prompt-tuning protocol rather than full-model updates.
  • Baseline aggregation: All three baselines implement Merge as coordinate-wise weighted averaging through the doubly stochastic mixing matrix W.This assumes same-index prompt tokens represent comparable semantic directions.
  • Baseline methods: D-PSGD-PT performs one local SGD update per round, whereas DFedAvgM-PT uses multiple local SGD iterations with momentum before communication.DFedAvgM is designed to reduce communication rounds relative to D-PSGD.
  • Baseline methods: DFedSAM-PT replaces local SGD with two-step Sharpness-Aware Minimization to seek flatter loss neighborhoods under heterogeneous data.Its update perturbs parameters toward high loss before evaluating the gradient at the perturbed point.
  • PFPT comparison: PFPT treats prompt aggregation as distributed set modeling with a dynamically sized global pool that merges similar prompts centrally.Its naive decentralized extension causes local prompt pools to grow exponentially because clients repeatedly union divergent neighborhood pools.
  • D-FROST analysis: D-FROST analyzes local OT prompt merging and decentralized global behavior by representing client prompt sets as empirical measures in Wasserstein space.The analysis targets stability of the alternating solver and convergence of the network-level barycenter.

B.1 Stability of the Local OT-Based Merge

The local OT-based Merge solver has a quantitative stability guarantee: increasing inner OT iterations reduces successive-iterate changes. This stability limits merge approximation error, supporting network-level consensus and convergence analysis.

  • Objective lower bound: The local OT objective is bounded below, enabling control of total objective decrease across alternating transport and barycenter updates.The spatial matching and L2 terms are non-negative, while the entropy term is lower-bounded over the feasible transport simplex.
  • Solver stability: The alternating OT solver produces a stable representative prompt set as the number of inner iterations increases.Theorem 1 bounds the smallest successive-iterate discrepancy among the first S iterations by O(1/S).
  • Solver stability: Stability depends on the initial objective gap and effective strong-convexity parameter µ = min(ε, λ/σ2).These quantities determine the bound on the successive-iterate discrepancy.
  • Implications for D-FROST: Increasing S reduces inner-solver instability and the approximation error from replacing neighborhood prompts with a compact representative set.This local error control feeds into the subsequent Wasserstein consensus and network-barycenter convergence analysis.
  • Measure-space aggregation: D-FROST analyzes clients’ prompt sets as empirical measures and uses fixed-budget OT compression during communication.The compressed measures define the practical aggregation state used to track network consensus.

C.1 Datasets and Partitions

The experiments use composite image datasets spanning multiple visual domains and impose controlled class-distribution heterogeneity. Performance is evaluated under several communication topologies represented by symmetric doubly stochastic mixing matrices.

  • Datasets: FourDataset combines MNIST-M, Fashion-MNIST, CINIC-10, and MMAFEDB, while FiveDataset combines CIFAR-10, MNIST, Fashion-MNIST, SVHN, and notMNIST.The composite datasets create heterogeneity across visual domains and class distributions.
  • Data partitions: Dirichlet partitions use α = 0.1 and α = 0.5, with smaller α producing stronger class-distribution skew.Each client draws a class-proportion vector over the relevant domain classes.
  • Data partitions: The extreme non-IID split assigns 99% of each class to one designated client and distributes the remaining 1% among other clients.This construction creates maximally imbalanced local datasets within each sub-dataset.
  • Shared settings: All methods use a frozen ViT-B/32 backbone with 10 learnable prompt tokens of dimension d = 768, Adam, and batch size 16.The shared protocol keeps backbone, optimizer, prompt configuration, and evaluation settings consistent across methods.
  • Communication topologies: The topology evaluation covers Ring, Grid, Erdős-Rényi, Regular, and Fully Connected graphs with symmetric doubly stochastic mixing matrices.The standard topology is a time-varying κ-regular graph, using κ = 5 for FiveDataset and κ = 4 for FourDataset.

D Additional Experiment Results

Additional experiments show that D-FROST maintains the strongest accuracy across heterogeneous FourDataset and FiveDataset settings, with gains that widen as heterogeneity increases and remain consistent across topologies.

  • Cross-dataset results: D-FROST achieves the best accuracy across Dirichlet and extreme non-IID settings on FourDataset and FiveDataset.Its margin over the strongest baseline widens as heterogeneity increases.
  • Topology robustness: D-FROST’s gains remain consistent across all evaluated communication topologies.The reported advantage of OT-based merging does not depend on a particular graph structure.
  • FourDataset: On FourDataset, D-FROST reaches 69.06%, 63.94%, and 58.24%, improving over DFedAvgM by 5.78, 4.22, and 12.19 points, respectively.These values correspond to the three reported non-IID settings.

E Factors Impacting Network Consensus Error

Network mixing, link reliability, and participation affect decentralized consensus behavior. The analysis and experiments connect weaker mixing with larger consensus neighborhoods and examine these effects across topology and network conditions.

  • Mixing and consensus: Theorem 2 bounds consensus error by ε(t) ≤ β/(1 −α) = O(...), with the denominator governed by the mixing factor ρ.Weaker mixing, represented by larger ρ, enlarges the consensus neighborhood.
  • Heterogeneity: Figure 7 compares test accuracy across FourDataset under Dirichlet α=0.5, Dirichlet α=0.1, and extreme non-IID partitions.The figure covers 40 clients and three increasing levels of non-IID heterogeneity.
  • Topology effects: Denser topologies mix faster and yield lower consensus error than sparse topologies such as Ring.The reported ordering places Grid below Ring in consensus error, consistent with the mixing-factor analysis.
  • Consensus-error evaluation: Figure 8 tracks total Wasserstein consensus error for D-FROST across clients on FiveDataset under Dirichlet α = 0.1.The experiment examines consensus error over communication rounds.

F Cost Analysis of D-FROST

D-FROST’s OT-based merge has modest computational overhead, no extra communication cost, and supports accuracy gains under heterogeneous settings. Its cost scales predictably with prompt budget and inner iterations, while larger prompt budgets improve attainable accuracy.

  • Computational complexity: The merge cost is quadratic in prompt budget n and linear in κ+1, embedding dimension d, and inner iterations S.Each inner iteration computes a cost matrix, transport plan, and barycenter update.
  • Computational complexity: 4.6 × 10^7 MACs make the OT merge negligible compared with a frozen ViT-B/32 forward/backward pass.With S=50, κ=5, n=10, and d=768, the neighborhood size is 60.
  • Wall-clock and communication cost: The OT merge takes 0.63 s versus 0.26 s for index-wise averaging, accounting for ≈1.5% of the 42.93 s round time.Local training dominates the per-round cost at approximately 42.3 s.
  • Wall-clock and communication cost: The OT merge adds no extra communication cost because clients exchange only updated prompt sets, not the locally computed cost matrix or transport plan.This matches the communication pattern of the parametric baselines.
  • Wall-clock and communication cost: D-FROST reaches 79.36% accuracy at 42.93 s/round and lies on the Pareto frontier, while DFedSAM-PT reaches 66.74% at 80.75 s/round.D-PSGD-PT and DFedAvgM-PT are faster per round but plateau below D-FROST in accuracy.
  • Prompt budget and robustness: Shrinking the prompt budget from n = 10 to n = 5 drops FiveDataset extreme non-IID accuracy from 70.46% to 63.35%.The paper uses n = 10 as the default and states that larger budgets tighten consensus and stationarity neighborhoods.
  • Improved baselines: Under extreme non-IID data, D-FROST reaches roughly 79%, exceeding DFedAvgM+ near 63%, DFedSAM+ around 58%, and D-PSGD+ around 30%.The comparison uses baselines augmented with client-specific prompt selection.
Loading 2609.01802v1…