PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimal Hydra: Model Heterogeneous Federated Learning for LLMs using Data-Informed Hydra-Nets

Gwen Legate, Rahaf Aljundi, Eugene Belilovsky

arXiv:2610.20753v1cs.LGcs.AI

TL;DR

Federated fine-tuning of large language models is constrained by client resources, while earlier model-heterogeneous FL work focused on pre-training and did not tailor sub-networks to local data. Optimal Hydra combines shared layer selection with client-specific attention-head selection. Across three datasets, it outperformed adapted heterogeneous-FL baselines in the most resource-heterogeneous settings while reducing local resource use.

  • Problem

    Large-model resource demands challenge federated training, and prior model-heterogeneous FL work focused on pre-training rather than fine-tuning or client-specific data distributions.

    Source: §1 · pp. 1–2

  • Method

    Optimal Hydra forms data-informed Hydra-Nets through common layer selection and client-specific gradient-based attention-head selection.

    Source: §1 · pp. 2–3, §3.2.2 · p. 6

  • Results

    Across three datasets, Optimal Hydra consistently outperformed adapted model-heterogeneous FL baselines in the most resource-heterogeneous settings, reducing local computation by up to 32.0% and memory by up to 29.3%.

    Source: §5 · p. 10

  • Takeaways & Limitations

    The results challenge the assumption that parameters must be trained equally across clients during fine-tuning.

    Source: §1 · p. 3, §5 · p. 10

    The individual-head formulation imposes hard bounds for an even distribution of local parameters, although fine-tuning behavior can differ from pre-training.

    Source: §3.2.2 · p. 7

Abstract

from arXiv · show

Large language models (LLMs) trained on vast amounts of data have demonstrated impressive performance over a range of tasks, but future performance gains will require increasing amounts of data and compute. Federated Learning (FL) facilitates training on decentralized private data, but client resource limitations can make training impractical, particularly as model sizes continue to expand. These resource constraints make the assumption of a shared server-client architecture impractical. Model heterogeneous FL offers a promising alternative by adapting model size to client capabilities; however, previous work is limited to the pre-training setting, and conclusions on sub-network selection do not necessarily hold for fine-tuning. Furthermore, prior approaches create client sub-networks without considering their unique data distributions, thereby overlooking the opportunity to tailor sub-network architecture to local data. Motivated by evidence that attention heads exhibit task dependency within layer defined feature spaces, we propose Optimal Hydra, a transformer based heterogeneous FL framework for fine-tuning that combines global layer selection with client specific attention head selection, aligned with local data distributions. Optimal Hydra makes explicit use of client data to produce client-specific sub-networks we call Hydra-nets. To our knowledge, this is the first model heterogeneous FL approach tailored to LLM fine-tuning. Experiments on three datasets show significant improvements over baseline methods while also reducing local resource demands.

1 Introduction

Optimal Hydra addresses the resource demands and architectural mismatch that complicate federated fine-tuning of large models, while targeting client-specific data distributions. It combines shared layer selection with local attention-head selection, and reports improved baseline performance alongside lower resource use.

1 Introduction

  • Large model resource demands can force FL to limit global model complexity or exclude low-resource clients, while homogeneous architectures require clients to store and train the full model.The resulting exclusion can bias training toward data from better-resourced clients.Source: §1 · p. 1
  • Earlier model-heterogeneous FL work focused on pre-training and did not tailor sub-networks to heterogeneous client data distributions.The paper notes that conclusions about equal parameter training from pre-training need not hold for fine-tuning.Source: §1 · p. 2
  • Hydra-Nets combine collaboratively selected common layers with client-specific, gradient-scored attention-head selection to capture distribution-specific features.Common layers retain all their attention heads; clients select heads from the remaining layers.Source: §1 · pp. 2–3
  • Up to 42.5% improvement over the best performing baseline accompanies compute and memory reductions of up to 32.0% and 29.3%, respectively.The reported gains are for challenging model-heterogeneous conditions.Source: §1 · p. 3
  • The authors argue that strong initial representations from pre-trained models allow fine-tuning to relax pre-training assumptions about evenly distributing parameters.They also motivate local head selection by the task-dependent nature of attention heads within layer feature spaces.Source: §1 · p. 3

2 Background and Related Work

Federated learning can support collaborative training without sharing data, but large foundation models raise local compute and memory demands. Prior model-heterogeneous approaches and attention-head studies motivate adapting sub-networks to both client resources and task-specific data.

2 Background and Related Work

  • Modern foundation-model parameter counts make local computation and memory major FL bottlenecks, while homogeneous architectures can force smaller global models or exclude resource-limited clients.Stable aggregation of heterogeneous local architectures into one global model also requires careful design.Source: §2 · p. 3
  • HeteroFL and FedDrop use static or random sub-model selection, while FedRolex and FedBRB train parameters more evenly but do not account for local-data suitability.These methods target pre-training, often smaller models and computer-vision tasks, rather than transformer fine-tuning.Source: §2 · p. 3
  • Attention-head studies report task specificity and redundancy during fine-tuning, and targeted heads can outperform random selection or layer-probing approaches.Some heads are reported as universally important across tasks.Source: §2 · pp. 3–4

3 Optimal Hydra

Optimal Hydra adapts federated fine-tuning to heterogeneous clients by combining shared layer selection with client-specific attention-head selection. Its Hydra-Nets use local data to tailor client sub-networks while preserving a common structural backbone.

3 Optimal Hydra

  • Optimal Hydra maps heterogeneous client updates back to global coordinates and weights parameter averaging by the samples of clients contributing to each parameter.This lets the global model accumulate specialized, data-informed gradients from specific heads without the gradient interference of full-parameter federated averaging.Source: §3 · p. 4

3.1 Problem Formation

  • Clients train Hydra-Nets with parameter sets smaller than the server model, while sharing at least the globally common layers.Different clients may share all, some, or none of the selected heads in the remaining layers.Source: §3.1 · p. 4

3.2 Hydra-Net Selection

  • Hydra-Net selection combines common layer selection with client-specific head selection because attention heads can be task-specific while some are universally important.Clients retain heads in selected common layers, then choose heads sensitive to their local data in the remaining layers.Source: §3.2 · p. 5

3.2.1 Common Layer Selection

  • The method retains all heads in low- and medium-redundancy layers, which empirically comprise fewer than 15% of layers.These common layers preserve important functionality and provide a shared structural backbone intended to defend against client drift.Source: §3.2.1 · pp. 5–6

3.2.2 Individual Head Selection

  • For remaining layers, a gradient-based score identifies attention heads important to each client’s local data, without the computational overhead of second-order methods.The head-assignment objective also balances assignments across clients; λ controls how strictly balance is enforced, and the method relaxes hard even-allocation bounds with slack variables.Source: §3.2.2 · pp. 6–7

4 Experiments

Experiments compare Optimal Hydra with transformer-adapted HeteroFL and FedRolex across three datasets and resource settings. Results vary by setting: gains are consistent under heterogeneous client resources, while parameter-balancing preferences differ by dataset.

4 Experiments

  • Experiments use Alpaca, Natural Instructions, and Dolly-15K, with configurations spanning full, three-quarter, half, and quarter attention-head retention.Dolly-15K and Alpaca use Dirichlet-distributed client data with α = 0.5; NI assigns 738 tasks uniquely to clients and holds out 119 tasks for global evaluation.Source: §4 · p. 7
  • 42.5% improvement over FedRolex on Natural Instructions in the 3Q(2)H(2)Q(16) configuration; Optimal Hydra outperforms both baselines across heterogeneous-resource configurations.Figure 2 reports mean Rouge-L over four seeds; the advantage is most pronounced on Natural Instructions, which has the most heterogeneous task distribution.Source: Figure 2 · p. 8, §4 · p. 10
  • At the most extreme setting, retaining 25% of attention heads saves 32.0% in computation and 29.3% in memory versus full FedAvg fine-tuning.At sequence length ≈12000, computational savings reach 50%.Source: §4 · p. 9, Table 1 · p. 9
  • With equal client resource budgets, Optimal Hydra is best in most configurations or competitive with the strongest baseline, but its advantage is relatively modest.The importance of retained-head choice increases as available model capacity decreases.Source: §4 · p. 9
  • The degree of balanced head assignment that works best varies by dataset: Dolly-15K often scores best at λ = 0 or 1, while NI generally favors larger λ values.NI sometimes performs best at λ = 500 or 1000; the authors relate the difference to head-score diversity and pretrained-representation strength.Source: §4 · p. 10

5 Conclusion

Optimal Hydra is a model-heterogeneous federated fine-tuning framework that forms client-specific, data-informed Hydra-Nets. Across three datasets and resource configurations, it outperforms adapted baselines in the most heterogeneous settings while reducing local compute and memory requirements.

5 Conclusion

  • Optimal Hydra creates data-informed Hydra-Nets that adapt model parameters to resource-limited clients and their local data.The framework targets fine-tuning large language models under heterogeneous client resource constraints.Source: §5 · p. 10
  • Across three datasets and multiple resource configurations, it outperforms adapted model-heterogeneous FL baselines in the most resource-heterogeneous settings.Local computational and memory requirements fall by up to 32.0% and 29.3%, respectively.Source: §5 · p. 10
  • The degree of parameter balancing needed during fine-tuning depends on pretrained-representation strength and local-data heterogeneity.Source: §5 · p. 10

A Memory and Compute Calculations

This section explains the rationale for the paper’s memory and compute calculations. Table 3 summarizes variables and values used in the appendix calculations.

A Memory and Compute Calculations

A.1 Compute

The compute calculation combines attention, feed-forward, and output-map operations over transformer layers. For DataJuicer-1.3B, the resulting expression is given as a function of attention-head scaling and sequence length.

A.1 Compute

  • Matrix multiplication of shapes (M × K) and (K × N) is counted as 2MKN flops, and α scales attention computation according to head reduction.The appendix defines α as the ratio of retained to original attention heads.Source: §A.1 · p. 13
  • Per transformer block, attention includes Q, K, V, and output projections; head reduction scales projection cost to 8d^2α and score cost to 2sdα.Source: §A.1 · p. 13
  • The SwiGLU feed-forward block costs 6di flops per layer, while the output map contributes 2dv once.The feed-forward total covers the gate, up, and down projections.Source: §A.1 · p. 14
  • The one-forward-pass formula combines these contributions over L layers and is expressed for DataJuicer-1.3B as a function of α and sequence length s.Table 3 summarizes variables and values for the DataJuicer-1.3B architecture.Source: §A.1 · p. 14, Table 3 · p. 14

A.2 Memory

The memory estimate separates model-weight memory from KV-cache memory, assuming 16-bit precision. It combines static and dynamic contributions to estimate bytes required for a forward pass.

A.2 Memory

  • Memory calculations assume 16-bit precision, or 2 bytes per parameter, and divide requirements into static model-weight memory and dynamic KV-cache memory.The combined contributions yield an estimate of bytes required for one forward pass.Source: §A.2 · pp. 14–15
  • Attention weights contribute 4d^2α × 2 bytes per layer, while feed-forward weights contribute 3d_i × 2 bytes per layer.Source: §A.2 · p. 14
  • The KV-cache contributes 2s(h′ × d_h) × 2 bytes per layer, giving total dynamic memory of 4Ls(h′d_h).Source: §A.2 · p. 15
  • For DataJuicer-1.3B, the memory expression is parameterized by α and sequence length s using architecture-specific constants.The supplied expression is approximately 24(8 × 2048^2 × α + 6 × 2048 × 5504) + 4 × 32000 × 2048 + 4 × 24 × s × 2048 × α.Source: §A.2 · p. 15

B Experimental Settings

Experiments ran across multiple compute clusters and GPU types, using two or four GPUs per experiment depending on availability.

B Experimental Settings

  • Experiments used L4OS, RTX8000, A100L, A100, V100, A6000, or H100 GPUs, with 2 or 4 GPUs per run depending on compute availability.Source: B Experimental Settings · p. 15

C Hyper-parameter Searches

Hyperparameters were selected through grid search, with Optimal Hydra adding λ to govern equal-parameter training. Final settings are reported for each dataset and heterogeneity combination.

C Hyper-parameter Searches

D ∇(Wv)2 Dominates Head Scoring

The head-scoring discussion motivates simplifying the scoring function because the score from Eq. 11 is largely dominated by the squared gradient of Wv. Selected client, layer, and round examples illustrate this pattern.

D ∇(Wv)2 Dominates Head Scoring

Related papers

Loading 2610.20753v1…