PaperlayerView exact revision on arXiv

Source-linked AI summary

Is Weight Tying Still Beneficial for Decoder-Only LLMs in Private Settings Under DP-SGD?

Razan El Mais, Ali Chehab, Ibrahim Issa, Razane Tajeddine

arXiv:2609.40335v1cs.LG

TL;DR

The effect of weight tying on differentially private optimization of decoder-only LLMs, and its compatibility with ghost clipping, remains poorly understood. The paper analyzes and evaluates tied and untied GPT2 and DistilGPT2 models, finding that untied embeddings improve DP-SGD utility and stability while retaining ghost clipping’s memory savings.

  • Problem

    How weight tying affects DP-SGD utility and the assumptions behind ghost clipping remains poorly understood, despite tying’s common use in decoder-only LLMs.

    Source: §I · p. 1

  • Method

    The paper analyzes shared-parameter effects on ghost norm computation and evaluates tied and untied GPT2 and DistilGPT2 models under DP-SGD.

    Source: F. Discussion · p. 8, §VII · p. 10

  • Results

    Across both models, untied embeddings improved DP-SGD utility and stability; on QNLI, DistilGPT2 accuracy rose from 72.59% to 73.57% versus weight tying with standard clipping.

    Source: F. Discussion · p. 9, B. Experimental Results of QNLI and QQP · p. 11

  • Takeaways & Limitations

    Untied embeddings provide a more favorable privacy–utility–efficiency trade-off, and ghost clipping retains its memory savings when its assumptions are satisfied.

    Source: F. Discussion · p. 9, §VII · p. 10

    The empirical evaluation covers only GPT2 and DistilGPT2 and focuses on embedding weight tying, leaving effects in larger or newer models and other forms of parameter sharing unestablished.

    Source: §VI · p. 9

Abstract

from arXiv · show

Differentially Private Stochastic Gradient Descent (DP-SGD) is a leading approach for privacy-preserving fine-tuning of large language models (LLMs). Many decoder-only LLMs employ weight tying between input and output embeddings, a design choice originally introduced for parameter efficiency and improved language modeling performance in the non-private setting. However, the impact of weight tying under differentially private training remains largely unexplored. In this work, we investigate the role of weight tying in the DP setting using GPT2 and DistilGPT2 as representative decoder-only architectures. Interestingly, we find that untied embeddings consistently outperform weight-tied models under DP-SGD, achieving gains of up to 4.74% points in accuracy on SST-2, QNLI, and QQP. Beyond improved utility, untying embeddings enables the use of memory-efficient ghost clipping for DP-SGD. By contrast, weight tying introduces shared-parameter interactions that complicate standard ghost norm computation and largely negate its computational advantages. As a result, untied models achieve over 60% lower memory usage while preserving the benefits of ghost clipping. Our results indicate that untied embeddings provide a more effective and scalable design for differentially private training of decoder-only LLMs and highlight the need to revisit standard LLM architectural choices in the privacy-preserving setting.

I. INTRODUCTION

The paper examines whether weight tying remains beneficial for differentially private fine-tuning of decoder-only language models, and how tying interacts with ghost clipping. It studies GPT2 and DistilGPT2 across GLUE tasks and reports that untied embeddings improve utility and optimization stability under DP-SGD.

A. Differentially Private Stochastic Gradient Descent

  • DP-SGD clips per-example gradients and adds calibrated noise, but per-example clipping can impose substantial memory and runtime costs for Transformer LLMs.Ghost clipping reduces these costs by computing clipping norms without explicitly materializing individual gradients.Source: §I · p. 1

I. INTRODUCTION

  • Ghost clipping assumes gradient contributions across parameter groups are block-separable, an assumption challenged when decoder-only models share input and output embedding parameters.Source: §I · p. 1
  • Weight tying is widely used for parameter efficiency and often improves generalization in non-private training, but its interaction with DP optimization remains poorly understood.Source: §I · p. 1
  • The study compares tied and untied GPT2 and DistilGPT2 variants on SST-2, QNLI, and QQP under non-private SGD, standard DP-SGD clipping, and ghost clipping.Source: §I · pp. 1–2
  • Untied embeddings improve utility and optimization stability under DP-SGD, while enabling memory-efficient ghost clipping.The paper also identifies an additional cross-term in per-example gradient norms when input and output embeddings share parameters.Source: §I · p. 2

B. Efficient Per-Example Gradient Computation

Per-example gradient clipping is a practical bottleneck for private training, motivating methods that compute gradient norms without materializing individual gradients. Ghost clipping reduces memory use, but its additive norm decomposition relies on separable parameter groups and becomes problematic with tied embeddings.

B. Efficient Per-Example Gradient Computation

  • DP-SGD’s per-example gradient computation and clipping can impose prohibitive memory and computational costs, especially in billion-parameter models.Source: B. Efficient Per-Example Gradient… · p. 2
  • Ghost clipping computes per-example gradient norms from activations and backpropagation signals without explicitly materializing gradients, reducing GPU memory use for Transformer-scale LLMs.Source: B. Efficient Per-Example Gradient… · p. 2

C. Ghost Clipping Mechanism

  • For a linear layer, ghost clipping uses the outer-product form of per-example weight gradients and Frobenius-norm properties to compute their norms indirectly.Source: C. Ghost Clipping Mechanism · pp. 2–3
  • Ghost clipping obtains the full per-example gradient norm by summing layer-wise contributions, then reweights the loss so a second backward pass computes the clipped aggregate without storing individual gradients.Source: C. Ghost Clipping Mechanism · p. 3
  • The additive norm computation assumes layer-wise gradient contributions can be decomposed independently, which is problematic when weight tying shares parameters.Source: C. Ghost Clipping Mechanism · p. 3

D. Weight Tying in Decoder-Only LLMs

Weight tying shares parameters between the input embedding and output projection, reducing parameter count and often improving generalization in non-private language modeling. That sharing couples gradients from both pathways, making its implications for DP-SGD and efficient clipping important to examine.

D. Weight Tying in Decoder-Only LLMs

E. Ghost Clipping Assumptions

Ghost clipping computes per-example gradient norms efficiently by summing contributions across independent parameter groups. Tied embeddings violate that separability, changing the true norm and compromising clipping; untied embeddings restore the intended structure and support a better privacy–utility–efficiency trade-off.

E. Ghost Clipping Assumptions

A. Understanding the Weight Tying–Ghost Clipping Conflict

  • Tied input and output embeddings receive gradient contributions from both pathways in the same shared parameter, introducing a cross-term that violates additive separability.The true norm therefore includes interaction between the input-embedding and output-projection gradients.Source: A. Understanding the Weight… · pp. 3–4
  • Untied embeddings achieve higher utility, improved optimization stability, and substantial memory savings under DP-SGD with ghost clipping across GPT2 and DistilGPT2.The findings indicate that weight-tying benefits in non-private training do not necessarily carry over to the differentially private setting.Source: A. Understanding the Weight… · p. 4

B. Corrected Ghost Norm for Tied Embeddings

The corrected tied-embedding norm explicitly incorporates the interaction between input and output pathways, restoring the true per-sample norm. But computing the cross-term removes ghost clipping’s core computational advantage, while untying restores additive separability.

B. Corrected Ghost Norm for Tied Embeddings

  • The corrected tied-gradient norm incorporates the cross-term to recover the exact per-sample gradient norm for the shared embedding matrix.It removes the discrepancy introduced by weight tying and restores consistency with the norm used for DP-SGD clipping.Source: B. Corrected Ghost Norm for Tied… · p. 4
  • Evaluating the cross-term requires materializing corresponding gradients, partially reintroducing the memory and computational overhead ghost clipping is designed to avoid.The correction therefore does not preserve ghost clipping’s core computational advantage.Source: B. Corrected Ghost Norm for Tied… · p. 4
  • Untying input and output embeddings restores additive gradient separability, allowing ghost clipping to operate as intended while retaining its scalability advantages.Source: B. Corrected Ghost Norm for Tied… · p. 4

IV. EXPERIMENTAL SETUP

The experiments compare tied and untied versions of GPT2 and DistilGPT2 on three GLUE classification tasks. All tasks use autoregressive label generation, preserving the models’ standard training objective while isolating embedding sharing.

A. Models

  • The study uses DistilGPT2 (approximately 82M parameters) and GPT2 (approximately 124M parameters), each evaluated with tied and untied embedding variants.The untied setup keeps the input embedding matrix and output projection as independent tensors.Source: A. Models · p. 4

B. Tasks and Dataset

  • Experiments cover SST-2, QNLI, and QQP, three binary classification datasets from the GLUE benchmark.SST-2 is the primary benchmark; QNLI and QQP assess whether observed effects generalize across natural-language-understanding tasks.Source: B. Tasks and Dataset · pp. 4–5
  • Each task is framed as prompt-based autoregressive classification, with the model generating a label token rather than using a task-specific classification head.This preserves the standard autoregressive objective without modifying the underlying architecture.Source: B. Tasks and Dataset · p. 5
  • The evaluation uses downstream classification tasks to isolate the interaction between weight tying and ghost clipping under DP-SGD, not to maximize task-specific performance.Experiments use a maximum sequence length of 128 tokens.Source: B. Tasks and Dataset · p. 5

C. Training Configurations

The study compares non-private SGD with standard- and ghost-clipped DP-SGD across tied and untied configurations. It assesses utility, memory use, and variability, with configurations selected through validation-based tuning.

C. Training Configurations

  • The training comparison includes non-private SGD, DP-SGD with normal clipping, and DP-SGD with ghost clipping, using tied and untied embeddings where applicable.This permits analysis of the separate and combined effects of privacy training, ghost clipping, and embedding sharing.Source: C. Training Configurations · p. 5
  • Final configurations are selected by validation accuracy and optimization stability after tuning learning rates, clipping norms, batch sizes, and epochs.Main DP experiments use a target privacy budget of ϵ = 3; privacy loss is tracked with an RDP accountant.Source: C. Training Configurations · pp. 5–6

D. Evaluation Metrics

  • Evaluation reports mean validation accuracy across five runs, peak GPU memory use, and performance variability across independent random seeds.These measures capture predictive utility, memory scalability, and optimization robustness.Source: D. Evaluation Metrics · p. 6

E. Implementation Details

Under standard non-private training, tying preserves accuracy while reducing parameter and memory costs; under normal-clipping DP-SGD, untied models outperform tied models and show lower seed variability. The implementation also computes an explicit interaction term for tied-weight ghost clipping.

E. Implementation Details

  • Tied-weight ghost clipping uses a corrected norm that separately evaluates input-embedding, output-projection, and interaction contributions for each sample.The experiments ran on a single NVIDIA A100 40GB GPU.Source: E. Implementation Details · p. 6

A. Weight Tying in Standard Non-Private Training

  • 86.77 ± 0.76% for DistilGPT2 and 89.09 ± 0.69% for GPT2: tied and untied models achieve identical average accuracy under non-private training.Tying reduces trainable parameters by approximately 32% for DistilGPT2 and 24% for GPT2, and lowers peak GPU memory for both.Source: A. Weight Tying in Standard… · p. 6, Table I · p. 7

B. Weight Tying Under DP-SGD with Normal Clipping

  • 80.61 ± 1.32% tied versus 83.23± 0.25% untied for DistilGPT2 average accuracy; GPT2 records 78.44 ± 1.73% versus 81.48±1.07% under normal-clipping DP-SGD.The reported improvements are not attributed solely to increased model capacity; seed variability also falls for both untied models.Source: B. Weight Tying Under DP-SGD with… · p. 7

C. Weight Untying Under DP-SGD with Ghost Clipping

Across the evaluated privacy budgets and benchmark tasks, untied embeddings improve utility, and ghost clipping preserves normal-clipping utility under untied embeddings with substantially lower memory use. Correcting ghost clipping for tied embeddings restores clipping consistency but adds substantial computational overhead.

D. Weight Tying Under DP-SGD with Corrected Cross-Term Computation

  • Corrected cross-term computation restores clipping consistency for tied embeddings, but raises DistilGPT2 training time from approximately 110.39 ± 3.55 minutes to 877.07 ± 2.92 minutes.The authors report that no computational method is currently known for evaluating the cross-term without materializing the corresponding gradients.Source: D. Weight Tying Under DP-SGD with… · p. 7

C. Weight Untying Under DP-SGD with Ghost Clipping

  • 67% for DistilGPT2 and more than 64% for GPT2: ghost clipping reduces peak GPU memory under untied embeddings while matching normal-clipping accuracy and stability.For DistilGPT2, both methods achieve 83.23±0.25% accuracy; for GPT2, both achieve 81.48 ± 1.07%.Source: C. Weight Untying Under DP-SGD with… · p. 7, Table III · p. 8

E. Complete DP Perspective

  • In main experiments at ϵ = 3 on SST-2, QNLI, and QQP, untied embeddings improve utility over weight tying; DistilGPT2 on SST-2 shows the same trend at ϵ = 1, 3, and 5.With untied embeddings, Ghost Clipping reproduces standard DP-SGD utility and stability while reducing GPU memory consumption by more than 60%.Source: E. Complete DP Perspective · p. 8

F. Discussion

Weight tying couples input-embedding and output-projection gradients, adding an interaction term that conflicts with standard ghost norm decomposition. The discussion contrasts tying's non-private efficiency benefits with its DP-SGD utility and stability costs, and emphasizes that architectural choices affect efficient private optimization.

F. Discussion

  • One shared matrix serves both input-embedding and output-projection pathways, adding a cross-term to the true per-example norm and invalidating standard independent-block ghost norm computation.The authors describe this as a change in the gradient geometry assumed by standard ghost norm computation.Source: F. Discussion · p. 8
  • Under DP-SGD, tied embeddings reduce utility and increase seed variability, unlike standard SGD, where tying preserves utility while improving parameter efficiency and memory use.The authors connect the DP-SGD difference to shared-gradient interactions under clipping and noise injection.Source: F. Discussion · pp. 8–9
  • Architectural design and privacy-preserving optimization should be considered together because architecture can alter gradient structure and affect private optimization's correctness and efficiency.Source: F. Discussion · p. 9
  • Explicit cross-term computation can restore mathematical correctness under tied embeddings, but substantially diminishes ghost clipping's computational advantages.Under untied embeddings, ghost clipping retains identical utility and stability to normal clipping while reducing GPU memory by more than 60% across both models.Source: F. Discussion · p. 9

VI. LIMITATIONS AND FUTURE WORK

The paper concludes that untied embeddings achieve better privacy–utility–efficiency trade-offs than tied embeddings in its GPT2 and DistilGPT2 experiments. It also identifies a cross-term that violates ghost clipping’s additive-separability assumption for tied embeddings and notes unresolved scope and efficiency limitations.

VI. LIMITATIONS AND FUTURE WORK

  • The empirical evaluation covers GPT2 and DistilGPT2 only, leaving utility effects in larger or newer decoder-only models unestablished.The identified cross-term follows from shared input-embedding and output-projection coordinates and is not specific to GPT2 scale.Source: §VI · p. 9
  • The analysis focuses on embedding weight tying; whether shared experts or parameter-efficient adaptation produce similar gradient-separability violations remains future work.Source: §VI · p. 9
  • Corrected norms for tied embeddings restore consistency, but computing the interaction term introduces substantial overhead, leaving efficient exact clipping an open challenge.Source: §VI · pp. 9–10

VII. CONCLUSION

  • Across GPT2 and DistilGPT2, untied embeddings achieved superior privacy–utility–efficiency trade-offs, improving predictive performance and optimization stability while preserving ghost clipping’s memory savings.The conclusion frames these findings as applying to the evaluated decoder-only models.Source: §VII · p. 10
  • Weight tying adds an interaction term to the per-sample gradient norm, so treating embedding pathways as independent blocks omits the cross-term needed for the exact shared-parameter norm.Including the cross-term restores exact norms but requires materializing gradient interactions that ghost clipping aims to avoid.Source: §VII · p. 10
  • The findings suggest that scalable DP optimization is not architecture-agnostic and may require co-design of privacy mechanisms, optimization algorithms, and parameter-sharing strategies.Source: §VII · p. 10, §VI · p. 10

APPENDIX

The appendix adds QNLI and QQP results and examines DistilGPT2 on SST-2 at additional privacy budgets. These results support the reported advantages of untied embeddings with Ghost Clipping across tasks, models, and the tested privacy budgets.

APPENDIX

  • The appendix reports additional QNLI and QQP results that reinforce the observations in the main paper.Source: APPENDIX · p. 11

A. Dataset Statistics

B. Experimental Results of QNLI and QQP

C. Sensitivity to the Privacy Budget for DistilGPT2 on SST-2

B. Experimental Results of QNLI and QQP

Loading 2609.40335v1…