Bayesian Data Scheduler for LLM Safety

• BDS is an adaptive defense mechanism that uses Bayesian inference over latent safety attributes to modulate influence during LLM fine-tuning.
• It introduces both Bayesian Scalar and Amortized Bayesian Neural Schedulers to assign soft data weights, ensuring harmful examples are attenuated without attack simulation.
• Empirical results demonstrate that BDS significantly lowers harmful scores while maintaining fine-tuning accuracy compared to other defense methods.

Bayesian Data Scheduler (BDS) is an adaptive tuning-stage defense for large language models (LLMs) against harmful fine-tuning in fine-tuning-as-a-service settings. It formulates defense as Bayesian inference over latent per-example “safety attributes” (w_i), conditioned on the user’s fine-tuning dataset (\mathcal{D}{\rm ft}) and the provider’s alignment dataset (\mathcal{D}{\rm safe}). These latent variables act as soft data weights during optimization, so that harmful or safety-eroding examples are attenuated without requiring attack simulation, while benign task-relevant examples can still drive customization. The framework has two versions: a Bayesian Scalar Scheduler, which infers one latent scalar per training point, and an Amortized Bayesian Neural Scheduler, which predicts weights with a neural network for transfer to new datasets without retraining the whole scheduler [2510.27172].

1. Problem setting and adaptive-defense rationale

BDS is designed for the fine-tuning-as-a-service regime, where a provider holds a base aligned model and a user uploads a dataset for customization. The threat model is harmful fine-tuning: the uploaded dataset is a mixture
[
\mathcal{D}{\rm ft} = \mathcal{D}{\rm ft}^{\rm benign} \cup \mathcal{D}_{\rm ft}^{\rm harmful},
]
with harmful ratio (p). In this setting, even a small amount of harmful or safety-eroding data can weaken alignment and induce unsafe outputs.

The method is explicitly motivated by two limitations of attack-simulation-based defenses. First, extending attack simulation beyond bounded threat models is infeasible because unknown attacks are inherently difficult to anticipate. Second, simulation-based robustness is not well adapted to varying attack settings, since fixed simulations do not capture their variability and complexity. BDS therefore defines “adaptive defense” in a post hoc sense: the defense is conditioned on the actual encountered fine-tuning dataset, rather than on a precommitted simulated attack distribution [2510.27172].

This design makes BDS neither a detector in the hard-filtering sense nor a preemptive immunization strategy. It is a posterior-weighting method that reacts to the specific uploaded dataset. The provider’s alignment dataset remains essential, but it is used jointly with the user dataset rather than in a strictly separate alignment-then-fine-tuning pipeline.

2. Bayesian formulation

The core latent variable in BDS is the per-example safety attribute (w_i \in \mathbb{R}), with “safety attribute” and “data weight” used synonymously. The inference target is the joint posterior
[
p(\boldsymbol{\theta}, \boldsymbol{w} \mid \mathcal{D}{\rm ft}, \mathcal{D}{\rm safe}),
]
where (\boldsymbol{\theta}) denotes the LLM parameters and (\boldsymbol{w} = (w_1,\dots,w_{|\mathcal{D}{\rm ft}|})). The construction assumes i.i.d. data, so (p(\boldsymbol{w}) = \prod_i p(w_i)), and the conditional independence (\mathcal{D}{\rm safe} \perp (\boldsymbol{w},\mathcal{D}_{\rm ft}) \mid \boldsymbol{\theta}).

The posterior is decomposed as
[
p\left(\boldsymbol{\theta}, \boldsymbol{w} \mid \mathcal{D}{\rm ft}, \mathcal{D}{\rm safe}\right) \propto p(\mathcal{D}{\rm safe}\mid\boldsymbol{\theta}) \cdot p(\mathcal{D}{\rm ft}\mid\boldsymbol{\theta},\boldsymbol{w}) \cdot p(\mathcal{D}_{\rm ft}\mid \boldsymbol{w})^{-1} \cdot p(\boldsymbol{\theta},\boldsymbol{w}).
]

The safe-data likelihood is written with the standard loss-as-energy form,
[
p\left(\mathcal{D}{\rm safe}\mid \boldsymbol{\theta}\right) \propto \prod{i=1}^{|\mathcal{D}_{\rm safe}|}\exp{\left(-\ell\left(\boldsymbol{z}i^{\rm safe};\boldsymbol{\theta}\right) \right)},
]
while the fine-tuning likelihood is reweighted by transformed safety attributes,
[
p\left(\mathcal{D}{\rm ft}\mid \boldsymbol{\theta},\boldsymbol{w}\right) \propto \prod_{i=1}^{|\mathcal{D}_{\rm ft}|}\exp{\left(- \sigma({w}_{i}) \cdot \ell\left(\boldsymbol{z}_i^{\rm ft};\boldsymbol{\theta}\right) \right)}.
]

The term (p(\mathcal{D}_{\rm ft}\mid \boldsymbol{w})^{-1}) is approximated through a safety-agnostic model (\hat{\boldsymbol{\theta}}), fit to weighted fine-tuning data only. This contrast is important in the paper’s interpretation: examples that fit the user data but conflict with safety alignment should receive lower influence. The resulting weighted fine-tuning objective is therefore not a mere curriculum or sample-prioritization heuristic; it is embedded in a joint posterior over model parameters and latent safety weights [2510.27172].

3. Schedulers, inference, and weight dynamics

BDS uses Stochastic Gradient Langevin Dynamics (SGLD) to sample approximately from the joint posterior over ((\boldsymbol{\theta},\boldsymbol{w})). In practice, the algorithm alternates between updating the model with safe-data loss plus weighted user-data loss, and updating the scheduler variables that control those weights.

The Bayesian Scalar Scheduler assigns one latent scalar (w_i) to each fine-tuning example. The Amortized Bayesian Neural Scheduler replaces explicit per-point latent inference with a neural mapping
[
p(w_i\mid\boldsymbol{\phi},\boldsymbol{z}{i}^{\rm ft})\coloneqq \delta\left(w_i-\mathcal{N}(\boldsymbol{z}{i}^{\rm ft}\mid \boldsymbol{\phi})\right),
]
so that (w_i) is predicted by a neural network with parameters (\boldsymbol{\phi}). This amortization enables transfer to unseen datasets by forward prediction rather than relearning all (w_i) from scratch.

A central implementation detail is the transformation (\sigma) applied to the latent weights. The paper studies identity, sigmoid, and softmax transformations. Identity and sigmoid induce monotone decreasing updates, so all examples tend to lose weight. Softmax instead compares each example’s loss with the weighted average loss, producing bidirectional scheduling: examples with below-average loss can gain weight, while above-average examples lose weight. This relative normalization is treated as critical. In the paper’s theoretical language, Theorem 1 on “Time-Weighted Accumulation of Posterior Bias” is used to argue that poor early weight dynamics can accumulate over time, which helps explain the failure of identity-like transformations [2510.27172].

Operationally, BDS performs soft scheduling rather than hard filtering. It does not discard data through a threshold; it modulates influence continuously. This is also why the method is called a “data scheduler”: it redistributes gradient influence across the fine-tuning set according to inferred safety attributes.

4. Empirical behavior across attacks, tasks, and models

The evaluation uses SST2, AGNEWS, GSM8K, AlpacaEval, and GEM tasks; harmful data from BeaverTails; out-of-domain attacks from RealToxicityPrompts and AdvBench; and identity-shifting attack (ISA) settings. The main backbones are Llama2-7B, Gemma2-9B, and Qwen2-7B. The principal metrics are Finetune Accuracy (FA) and Harmful Score (HS), where HS is the fraction of unsafe outputs on 1000 BeaverTails test instructions [2510.27172].

Setting	Booster	BDS
Average over low harmful ratios (p \in {0,0.05,0.1,0.15,0.2})	HS 10.94, FA 93.03	HS 1.34, FA 93.83
High harmful ratio (p=0.9)	HS 75.9	HS 1.5
AlpacaEval at (p=0.1)	HS 36.70, FA 45.19	HS 2.30, FA 48.08
RealToxicityPrompts at (p=0.1)	HS 26.8	HS 1.5
ISA at (p=0.8)	HS 77.4	HS 1.7

These comparisons summarize the paper’s broader pattern. At high harmful ratios (p=0.3) to (1.0), Booster’s HS rises through (40.6, 68.4, 77.2, 76.9, 77.5, 76.3, 75.9, 76.6), whereas BDS remains near (1.2, 1.5, 1.2, 1.5, 1.5, 1.6, 1.5, 1.8). At (p=0.9), the paper describes the gap between Booster HS (=75.9) and BDS HS (=1.5) as a (74.4\%) improvement.

Across architectures, the average comparison is Booster HS (7.03), FA (94.19) versus BDS HS (1.13), FA (94.30). Under varying fine-tuning set sizes (|\mathcal{D}_{\rm ft}|), Booster HS increases from (3.8) to (33.6) over (500,1000,1500,2000) examples, whereas BDS remains at (1.2,1.2,0.9,1.2). With only 100 safe examples, Booster HS is (62.2) and BDS HS is (1.7). The neural scheduler degrades slightly under transfer but remains strong: SST2 (\rightarrow) unseen SST2 gives HS (2.50), FA (93.23), and SST2 (\rightarrow) AGNEWS OOD gives HS (2.80), FA (89.20) [2510.27172].

5. Ablations, implementation, and efficiency

The ablation study identifies the softmax transformation as decisive. Softmax yields HS (1.20), FA (93.69); sigmoid yields HS (16.70), FA (92.23); and identity yields HS (16.50), FA (0.00). This is the clearest mechanistic result in the paper: the scheduler is not robust to arbitrary weight parameterizations.

The alignment dataset is also essential. Removing (\mathcal{D}{\rm safe}) gives FA (94.27) and HS (77.90) at (p=0.1), compared with FA (93.69) and HS (1.20) when (\mathcal{D}{\rm safe}) is used. Removing data weighting while retaining safety data is also insufficient: at (p=0.3), the unweighted variant gives FA (93.69), HS (11.20), whereas weighted BDS gives FA (93.32), HS (1.20). The paper further reports that a noninformative prior performs best, while Gaussian priors produce HS around (20)–(22), and initialization is robust over (0.001) to (10), with HS between (1.1) and (1.5) and FA around (93.6)–(93.7) [2510.27172].

Implementation uses LoRA with rank (32), LoRA alpha (4), FusedAdam, LLM learning rate (1\times10^{-5}), weight decay (0.1), batch size (10), and BF16 precision. The scalar scheduler learning rate is (5\times10^{-3}), the neural scheduler learning rate is (1\times10^{-6}), and the neural scheduler backbone is a lightweight 125M Fairseq-Dense model with a trainable linear head and hidden size (768). The default scheduler is the scalar version.

The paper positions BDS as nearly as cheap as vanilla fine-tuning and substantially cheaper than Booster. The reported benchmark is: Vanilla FT (2.01) min/epoch and (25.32) GB; Booster (6.42) min/epoch and (57.86) GB on H100-80G; BDS (2.04) min/epoch and (25.44) GB on A100-40G. The theoretical comparison is Booster (O(3n_1 f)) in alignment plus (O(n_2 f)) in fine-tuning, versus BDS (O(n_2(f+w))), where (w) is scheduler size and is usually much smaller than model parameter count (f) [2510.27172].

6. Interpretation, limitations, and acronym scope

The paper attributes BDS’s behavior to two mechanisms. First, posterior conditioning on the actual encountered fine-tuning dataset yields dataset-specific defense rather than robustness to a simulated threat distribution. Second, the joint safety-aware objective makes harmful examples incur larger loss, so their inferred safety attributes are reduced at the instance level. This suggests a defense that is adaptive both at the dataset level and at the example level [2510.27172].

Several misconceptions are therefore explicitly excluded by the method’s construction. BDS is not attack simulation, not threshold-based filtering, and not merely mixing safe data into fine-tuning. Its gains depend on posterior-weighted optimization. At the same time, the paper does report limitations: some over-refusal on safe prompts in XSTest, some unsafe prompts that receive non-refusal answers, extra computation from the neural scheduler, and open questions around better priors and model-level defenses.

The acronym is also ambiguous in the broader arXiv literature. “BDS” is used elsewhere for Broadband Dielectric Spectroscopy [2307.11537], BDS-3 in the BeiDou Navigation Satellite System [2112.13252], Beam Delivery System [1202.0562], and BDs, the Bayesian Dirichlet sparse score for Bayesian network structure learning [1704.03942]. In the present sense, however, Bayesian Data Scheduler denotes the 2025 LLM fine-tuning defense framework based on Bayesian inference over latent safety attributes.