PaperlayerView exact revision on arXiv

Source-linked AI summary

Modular Pretraining Enables Access Control

Ethan Roland, Murat Cubuktepe, Erick Martinez, Stijn Servaes, Keenan Pepper, Mike Vaiana, Diogo Schwerz de Lucena, Judd Rosenblatt, Addie Foote, Cem Anil, Alex Cloud

arXiv:2607.08077v1cs.LG

TL;DR

Dual-use AI capabilities can bring benefits and misuse risks, but separately training and deploying models for different access profiles is prohibitively costly. GRAM routes capability-specific updates into auxiliary modules within one training run. Across the reported evaluations, it closely tracks data filtering, and its ablated capabilities resist malicious finetuning better than post-hoc unlearning.

  • Problem

    Separate data-filtered models are a gold standard for serving different capability profiles, but training and deploying N models is prohibitively costly at frontier scale.

    Source: §1 · p. 1

  • Method

    GRAM adds capability-specific auxiliary modules to Transformer MLP layers and routes data-dependent updates so inference-time ablation can select capability profiles.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    On realistic dual-use data, GRAM nearly matches the all-data baseline on core data, matches filtering and FT-LoRA on capability retention and removal, and resists malicious finetuning better than post-hoc unlearning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    GRAM offers multiple capability profiles from one training run, while supporting composable capabilities and performing better under partial labeling than FT-LoRA.

    Source: §9 · p. 9

    The scaling trends have been shown up to 5B parameters, but not verified at frontier scales or in production settings.

    Source: §8 · p. 9

Abstract

from arXiv · show

AI developers face a dual-use dilemma. An AI capability that helps one user cure a disease can help another synthesize one. This dilemma could be resolved with access control, limiting dual-use AI capabilities to trusted deployments with a legitimate need. A gold standard for access control would be to serve separate models with different capabilities to different users. However, training and deploying multiple models is prohibitively expensive. To address this challenge, we propose gradient-routed auxiliary modules (GRAM), a pre-training method that adds modules to a neural network and selectively updates them to induce specialization. Ablating a module at inference time removes its capability from the network, approximating a model trained on filtered data. We evaluate GRAM on synthetic stories and realistic dual-use data spanning virology, cybersecurity, nuclear physics, and specialized code. These experiments show that GRAM disables targeted capabilities while preserving the rest, and resists their recovery under finetuning better than post-hoc unlearning. Most importantly, a Chinchilla-optimal scaling analysis from 50M to 5B parameters shows that the gap between data-filtered and full-data models widens with scale on removed capabilities but stays small on retained ones, and that GRAM closely tracks data filtering. GRAM's training cost is independent of the number of supported capability profiles, yielding a 5x reduction over data filtering in our 5-profile setting.

1. Introduction

Dual-use capabilities create pressure to balance misuse risk against beneficial use, while separate models for different capability profiles are prohibitively costly. GRAM aims to provide multiple profiles from one training run by routing capability-specific updates into auxiliary modules.

1. Introduction

  • Supporting N capability profiles with separately trained and deployed data-filtered models is prohibitively costly at frontier scale.The paper identifies these separate models as the gold standard for access control.Source: §1 · p. 1
  • Refusal training, output filters, and classifiers restrict outputs but leave capabilities intact, while post-hoc unlearning can be rapidly reversed by finetuning.Model branching reduces pretraining cost but still requires multiple deployments and may lack pretraining-scale capacity.Source: §1 · pp. 1–2
  • GRAM routes data-dependent updates to auxiliary modules, which can be ablated to remove their capabilities while retaining a single training run.The ablated model resists malicious finetuning about as well as data filtering and substantially better than post-hoc unlearning.Source: §1 · p. 2
  • In Simple Stories, GRAM closely matches data filtering on both capability retention and removal.The comparison uses a synthetic dataset partitioned into core and auxiliary subsets.Source: §1 · p. 2
  • On realistic dual-use data, GRAM nearly matches the all-data baseline on core data and matches filtering and FT-LoRA on capability retention and removal.Targeted removal is more robust to malicious finetuning than post-hoc unlearning; GRAM also performs better than filtering and branched finetuning with partially labeled data.Source: §1 · p. 2
  • From 50M to 5B parameters, GRAM closely tracks data filtering, whose suppression of forgotten capabilities widens with scale.The paper suggests this advantage persists at the model scaling frontier.Source: §1 · p. 2

2. Background

The paper formalizes capability profiles as sets of capabilities to retain or restrict, with a fixed compute budget equal to one full-data baseline run. It evaluates these profiles using compute ratios and an elicited-forget test.

2. Background

  • A capability profile retains the core dataset and selected auxiliary capabilities while restricting the remaining auxiliary capabilities.The goal is low validation loss on retained datasets and high validation loss on restricted ones, using the compute budget of one full-data baseline run.Source: §2 · p. 2
  • Compute ratio expresses the compute needed to reach a loss relative to the full-data baseline, with ideal profiles scoring at least 1 on retained datasets and 0 on restricted datasets.The metric uses baseline learning curves to map validation loss to equivalent baseline compute.Source: §2 · p. 2
  • Core, retain, and forget performance are evaluated separately; higher compute ratios indicate better core and retain performance, while lower ratios indicate better forget performance.The forget metric is interpreted in the opposite direction from core and retain metrics.Source: §2 · pp. 2–3
  • Elicited Forget Performance tests whether excluded capabilities return after finetuning on 512 sequences, about 0.5M tokens, from each forget category.Lower post-finetuning compute ratios indicate more robust removal; the fixed budget represents a limited adversary, and greater elicitation may partially recover capabilities even from filtered models.Source: §2 · p. 3

3. Gradient-Routed Auxiliary Modules

GRAM augments Transformer MLP blocks with a shared core and capability-specific auxiliary modules. Data-dependent activation and gradient routing localize updates, so inference-time module ablation selects a capability profile.

3. Gradient-Routed Auxiliary Modules

  • Unlike a mixture-of-experts model, GRAM has no learned router or per-token routing; data labels direct gradient updates to specific modules.The architecture augments each standard Transformer MLP block with smaller auxiliary MLP modules.Source: §3 · p. 3
  • Each GRAM layer pairs a large core MLP for general knowledge with smaller auxiliary modules for capability-specific knowledge.There is one auxiliary module for each auxiliary dataset.Source: §3 · p. 3
  • GRAM isolates capabilities through dataset-dependent forward activations and gradient updates over disjoint core and auxiliary parameter sets.Separate AdamW optimizers are used for the parameter subsets.Source: §3 · pp. 3–4
  • On core batches, the core MLP is always active and a random auxiliary module is activated with probability pcr, promoting robustness to auxiliary-module activation.Setting pcr too high can degrade both core and retain performance through interference.Source: §3 · p. 4
  • On auxiliary batches, the core and corresponding auxiliary module are active; the auxiliary module is always updated, while core parameters update with probability pas.Increasing pas can improve capability competence while reducing modularization; paf controls dataset sampling frequency.Source: §3 · p. 4
  • At inference, a profile activates the core MLP and the auxiliary modules for capabilities to retain, ablating the rest.The core remains active for every forward pass.Source: §3 · p. 4

4. GRAM Approximates Multiple Data-Filtered Models in a Single Run

In synthetic stories, GRAM approximates multiple data-filtered models in one run: retained capabilities stay near baseline while forgotten capabilities fall well below it. This separation persists as the number of auxiliary categories grows from 4 to 20.

4. GRAM Approximates Multiple Data-Filtered Models in a Single Run

  • The Simple Stories evaluation uses four auxiliary categories among 48 topic categories, with the remaining 44 forming the core dataset.The dataset contains 2M children's stories.Source: §4 · p. 4
  • The comparison matches active parameters per forward pass and uses three random-seed runs for confidence intervals.GRAM uses an always-active core MLP and smaller auxiliary modules; filtering trains separate models on the core dataset plus one auxiliary topic.Source: §4 · pp. 4–5
  • Across five capability profiles, GRAM approximates data filtering: retained auxiliary performance is 0.95 versus 0.96, while core performance is 0.94 versus 0.96.Forgotten categories remain well below retained ones; for example, a forgotten category scores 0.61 with GRAM versus 0.65 with filtering.Source: §4 · p. 5
  • As auxiliary categories grow from 4 to 20, retain performance stays near baseline, forget and elicited-forget remain substantially lower, and core performance stays roughly constant around 0.88.The fixed auxiliary budget is split across progressively more topics while the total training budget and core/auxiliary token mixture remain fixed.Source: §4 · p. 5, Figure 3 · p. 5

5. GRAM Isolates Dual-Use Capabilities Learned From Real Data

On realistic dual-use data, GRAM closely approximates data-filtered performance while supporting multiple capability profiles from one training run. Its modularization also performs well when capabilities are combined and when only half the training data is labeled.

Dataset and setup

  • The realistic evaluation uses web text, general arXiv papers, and code as core data, with virology, cybersecurity, nuclear physics, and Lisp as auxiliary domains.The 800M-parameter models train for one epoch on 16B core tokens plus 160M auxiliary tokens.Source: §5 · p. 5

Comparison methods

  • The comparison includes filtering, branched FT-LoRA and FT-Full, and post-hoc MaxEnt; DEMix is omitted from the main results because it performs poorly under class imbalance without a shared core.The shared core avoids requiring each module to learn basic language structure from limited data.Source: §5 · pp. 5–6

Results

  • GRAM and FT-LoRA reach roughly equal compute-normalized losses to filtering on core, retain, and forget sets, while MaxEnt’s forget-set degradation is mostly reversed by finetuning.The Figure 4 comparison averages performance across five capability profiles.Source: Figure 4 · p. 6, §5 · p. 6

Training and inference costs

  • Filtering requires five times the training FLOPS, while GRAM and FT-LoRA match baseline training compute and use smaller increases in active and total parameters.Filtering, FT-Full, and MaxEnt also require five times the baseline’s total parameters in this setting.Source: §5 · p. 6

Arbitrary capability subsets

  • GRAM’s retained-class performance stays constant as auxiliary modules are added, whereas FT-LoRA generally degrades when multiple adapters are summed.GRAM achieves this composability despite never training with all modules active.Source: §5 · p. 6

Partial labeling

  • With only 50% of training data labeled, GRAM achieves significantly lower forget performance (0.60 ± 0.01) than filtering and FT-LoRA (each 0.85 ± 0.01), indicating superior modularization.On unlabeled batches, all parameters—the core MLP and every auxiliary module—participate in forward and backward passes.Source: §5 · p. 6

6. Capability Isolation Improves with Scale

Across model sizes from 50M to 5B, GRAM and FT-LoRA closely track data filtering while requiring less amortized compute. Data filtering’s removal and resistance to adversarial finetuning strengthen with scale, while GRAM’s 5B retain result is anomalous.

Experimental setup

  • The experiment spans 50M to 5B parameters, with Chinchilla-optimal datasets and auxiliary data fixed at 1% of core dataset size.It compares GRAM, FT-LoRA, and Core + Virology data filtering for models retaining core and virology.Source: §6 · p. 7

Results

  • As data-filtered models scale, they learn excluded capabilities less well while maintaining core and virology performance and resisting adversarial finetuning more strongly.The authors speculate that domain-specific learning contributes a growing share of the baseline’s improvement on the excluded domain, widening the gap.Source: §6 · p. 7
  • GRAM and FT-LoRA track data filtering closely at every scale despite requiring five times less compute in amortized terms.The 5B GRAM retain result is off-trend and attributed to a scale-dependent hyperparameter interaction.Source: §6 · p. 7

7. Related Work

Related work addresses robust capability removal, neural modularity, and learning across tasks. GRAM is distinguished by gradient routing, asymmetric module sizing, and a shared core, with evaluation focused on broad capabilities across scales.

Data filtering for robust capability removal

  • Data filtering can robustly remove targeted capabilities, but its pretraining compute scales linearly with the number of supported capability profiles.Prior work also applies filtering to harmful CBRN information and undesired traits.Source: §7 · p. 7

Training for modularity

  • Unlike modular architectures that constrain how representations are organized, GRAM controls which parameters receive gradients from which data.DEMix assigns experts to domains, but symmetric sizing and no shared core limit its performance under class imbalance.Source: §7 · p. 8

Concurrent work

  • GRAM restricts gradients from forget data reaching shared parameters, while its scaling experiments isolate a few broad capabilities rather than many individual facts.Concurrent work by Ghosal et al. (2026) uses sparse masks to localize individual facts for removal.Source: §7 · p. 8

Modularity without supervision

  • Unsupervised routing and clustering approaches do not reliably produce domain-level separation or increased task specialization, whereas GRAM combines supervised routing with asymmetric modules and an always-active core.The authors connect these design choices to robust removal by module ablation.Source: §7 · p. 8

Continual learning

  • Unlike continual learning, which handles sequentially arriving data and seeks to avoid forgetting, this work uses all data during training to produce multiple capability profiles from one pretrained model.The paper evaluates continued pretraining and LoRA as related baselines.Source: §7 · p. 8

8. Discussion

The discussion frames capability control as a multi-objective trade-off and explains GRAM’s filtering-like behavior through module localization and cross-domain training effects. It outlines possible deployment uses alongside limitations in scope, implementation, and understanding.

Trade-offs

  • Capability-control methods trade off core and retained performance, forgetting, robustness, training and inference costs, memory, and implementation complexity, with priorities depending on use case.The authors compare methods at fixed compute cost and account for total parameters.Source: §8 · pp. 8–9

Why single-run methods match filtering

  • GRAM trains over all auxiliary domains, which the authors suggest improves retained-capability performance through cross-domain transfer, unlike filtering each model to one auxiliary domain.Branched methods introduce auxiliary data in a separate finetuning phase rather than mixing it throughout pre-training.Source: §8 · p. 9

Conditions on method comparisons

  • In the tested setting, differences between separating, ordering, and mixing data are small; larger auxiliary shares, multi-epoching, or different profiles may change that.Prior work cited by the authors finds that later sequential training can deteriorate pre-trained capabilities, whereas mixing domains from the start avoids this.Source: §8 · p. 9

Architectural separation

  • GRAM and FT-LoRA localize capabilities in dedicated modules, while post-hoc unlearning was mostly reversed by finetuning in the 800M MaxEnt experiments.Scaling a module’s forward weight from zero to one restores its capability up to baseline while leaving core performance unchanged; ablation-related loss increases concentrate on sparse domain-specific tokens.Source: §8 · p. 9

Partial labeling

  • With half the labels missing, GRAM removes capabilities more effectively than filtering or FT-LoRA, whose shared parameters learn from unlabeled auxiliary examples.The authors attribute GRAM’s smaller degradation to absorption: updates from similar unlabeled examples concentrate in a module already specialized on that domain.Source: §8 · p. 9

Deployment implications

  • Module control could support serving dual-use capabilities only in authorized deployments, while pre-release ablation may support open-weight models with purely beneficial capabilities.The proposed access-control mechanism is to control which modules are active at inference time.Source: §8 · p. 9

Limitations

  • Entangled target and desired capabilities may prevent clean separation; the favorable scaling trends are tested only through 5B parameters, not at frontier scale or in production.The authors identify substantial implementation effort as another production concern.Source: §8 · p. 9

Open questions

  • The authors lack a clear account of GRAM’s compositional advantage, have not directly tested the absorption explanation, and do not know how post-training affects separation.They also note that cross-entropy loss may not provide a complete performance picture.Source: §8 · p. 9

9. Conclusion

GRAM aims to provide the equivalent of many data-filtered models for the cost of one training run. Across scales, GRAM and FT-LoRA track data filtering; GRAM also composes capabilities and performs better under partial labeling, though those advantages need further study.

9. Conclusion

  • GRAM may provide the equivalent of many data-filtered models for the cost of a single training run.Source: §9 · p. 9
  • Across scales, both GRAM and FT-LoRA track data filtering, and the results establish that data filtering itself scales favorably.Source: §9 · p. 9
  • GRAM supports composable capabilities and performs better under partial labeling than FT-LoRA, though both benefits need further study.Source: §9 · p. 9

Impact Statement

The impact statement frames capability control as a possible way to reduce misuse while preserving beneficial deployment. The Simple Stories appendix specifies the data, model, and comparison methods used in its experiments.

Impact Statement

  • The authors present capability control as a possible means to reduce misuse while enabling broader deployment of capable models for beneficial applications.They encourage further study of limitations, adversarial users, and integration with organizational and policy safeguards.Source: Impact Statement · p. 10

A. Details for the Simple Stories Results

  • Simple Stories contains about 2 million stories across 48 topics, with four designated auxiliary capabilities and the other 44 pooled as core data.The four auxiliary topics are A Deadline or Time Limit, Alien Encounters, Bygone Eras, and Cultural Traditions.Source: A. Details for the Simple Stories… · p. 14
  • The Simple Stories models use an 8-layer Transformer with 26M parameters, and the reported Table 1 results are means with errors.Models were trained for one epoch with batch size 128, sequence length 256, and three random seeds.Source: A. Details for the Simple Stories… · pp. 14–15, Table 1 · p. 15
  • GRAM adds four auxiliary MLP modules to each Transformer MLP block, with a 1,856-dimension core MLP and 192-dimension auxiliary modules.Each auxiliary module has about 1.6M parameters, roughly 6% of total model capacity; active parameters are capped to match the baseline's total parameter count.Source: A. Details for the Simple Stories… · p. 14
  • MaxEnt suppresses unwanted capabilities by maximizing output entropy on forget data while preserving performance on retain data.The experiments ran it for 2,000 steps per profile and tuned retain-loss weight and learning rate via Bayesian optimization.Source: A. Details for the Simple Stories… · p. 14
  • DEMix uses equal-sized domain-specific MLP experts, unlike GRAM's larger always-active core MLP and selectively updated auxiliary modules.The comparison also includes FT-Full branching and FT-LoRA adapters; training steps across FT-Full branches sum to the baseline's steps.Source: A. Details for the Simple Stories… · pp. 14–15

A.1. Extended Simple Stories Results

On Simple Stories, GRAM preserves core and retained capabilities while removing excluded ones, with stronger robustness to adversarial finetuning than post-hoc unlearning. The accompanying sections describe the realistic dual-use data and scaling-experiment configurations.

A.1. Extended Simple Stories Results

  • 0.95 retain compute ratio for GRAM closely matches data filtering (0.96), while core performance is 0.94 versus 0.96; GRAM uses one training run and deployment rather than N models.Among single-run approaches, GRAM has the strongest overall tradeoff.Source: §A.1 · p. 15
  • GRAM's compute advantage over data filtering grows linearly with auxiliary capability count because filtering trains and deploys N models, whereas GRAM uses one model with module ablation.FT-Full also requires deploying N separate checkpoints.Source: §A.1 · pp. 15–16
  • DEMix has core and retain compute ratios of 0.25 and 0.18; the authors attribute this to lacking a commonly active core expert in the class-imbalanced setting.Its experts must redundantly learn basic dataset understanding from small data volumes.Source: §A.1 · p. 16
  • 0.86 elicited forget compute ratio for GRAM is near data filtering's 0.87, while MaxEnt reaches 0.98 after adversarial finetuning.GRAM's pre-elicitation forget ratio is 0.77; FT-Full reaches 0.81 and FT-LoRA 0.77, while DEMix reaches 0.30 amid poor overall learning.Source: §A.1 · p. 16

B. Details for the Realistic Dual-Use Baseline

  • The realistic dual-use mixture pairs general-purpose web text, arXiv papers, and non-Lisp code with four auxiliary domains: code, virology, nuclear physics, and cybersecurity.General-purpose arXiv papers are included to control for academic formatting differences when comparing core and auxiliary domains.Source: B. Details for the Realistic Dual-Use… · p. 16

C. Scaling Experiment Details

  • The scaling experiments use four auxiliary modules for virology, cybersecurity, nuclear physics, and code, with auxiliary data set to 1% of core data.All models train for one epoch using AdamW and a warmup-stable-decay schedule.Source: Table 3 · p. 18

D. GRAM Scaling with Increasing Numbers of Auxiliary Capabilities

The study tests whether GRAM maintains capability separation as auxiliary categories increase, holding training budget and core/auxiliary token balance fixed. Across 4 to 20 categories, core and retain performance stay roughly stable while excluded-capability performance remains lower.

Experimental Setup

  • The experiments use Simple Stories with 4, 8, 12, 16, or 20 auxiliary categories because the realistic setting has only four domains.The auxiliary sets are nested by a fixed alphabetical ordering, with seeds varying training stochasticity.Source: D. GRAM Scaling with Increasing… · p. 18
  • For each configuration, auxiliary categories receive 20% of training tokens and core categories receive 80%, with total training budget fixed.Sampling is adjusted to preserve this split and isolate the effect of category count.Source: D. GRAM Scaling with Increasing… · p. 18

Results

  • Core compute ratios stay around 0.88 from 4 to 20 auxiliary categories, while retain performance remains approximately 0.89–0.91.The number of active parameters per forward pass is held constant across configurations.Source: D. GRAM Scaling with Increasing… · pp. 18–19
  • Forget compute ratios remain approximately 0.65–0.74, and elicited-forget ratios remain approximately 0.71–0.84, below retain performance.Forget and elicited-forget ratios rise modestly as the fixed 20% auxiliary-token budget is divided among more topics.Source: D. GRAM Scaling with Increasing… · p. 19
  • GRAM supports increasing auxiliary counts without degrading core performance or collapsing retain–forget separation, using one training run and module ablation.The study holds active parameters constant, while adding categories increases dormant module parameter sets.Source: D. GRAM Scaling with Increasing… · pp. 18–19

E. Architecture Ablation: MLP vs. LoRA

This appendix section examines how GRAM’s capability profiles relate to auxiliary-module architecture and training methodology, including labeling, gradient accumulation, compute matching, and scaling choices. It also describes power-law scaling of optimal learning rates and batch sizes and applies the resulting values across methods because their core MLPs share an architecture.

E. Architecture Ablation: MLP vs. LoRA

  • Within each training methodology, GRAM’s capability profiles varied little between MLP and LoRA auxiliary modules, while profiles varied more between training methodologies.The comparison includes GRAM (MLP), GRAM (LoRA), FT-LoRA, and FT-MLP, with equal compute and auxiliary capacity fixed at 10% of core MLP parameters.Source: E. Architecture Ablation: MLP vs. LoRA · p. 19, Figure 9 · p. 20

F. Hyperparameter Appendix

  • GRAM’s routing sweeps expose tradeoffs in shared-representation quality, core stability, and retained capability, whereas FT-LoRA’s ratio shifts capacity without trading forget against retain.As pas increased from 0 to 1, core and retain rose 0.90 and 0.88 →0.97 and 1.04, while forget and elicited-forget rose 0.69 and 0.78 →0.89 and 0.97. As pcr increased, core configurations converged near 0.93 at pcr ∈[0.4, 0.8], but the ablated configuration fell to 0.79 at pcr = 1; retain declined 1.05 →0.92, while forget and elicited-forget changed 0.77 →0.71 and 0.90 →0.83. Increasing pca raised core 0.82 →0.90 and lowered retain 1.06 →0.90, while forget and elicited-forget stayed roughly 0.69–0.74 and 0.79–0.83.Source: Figure 11 · p. 20, Figure 10 · p. 20, F. Hyperparameter Appendix · p. 21

G. Partial Labeling Details

  • With 50% of tokens labeled, GRAM’s forget compute ratio was 0.60 versus 0.85 for partially labeled filtering and FT-LoRA, while core performance was 0.96 versus 0.99.All methods exceeded perfectly labeled filtering on core performance (0.92).Source: G. Partial Labeling Details · p. 22, Figure 13 · p. 22

H. Heterogeneous Accumulation

  • Heterogeneous accumulation freezes modules excluded by each micro-batch’s backward mask, preventing unwanted gradient contributions from entering their accumulated updates.Uniform accumulation particularly suppresses GRAM’s retain performance as accumulation steps increase; 400M and 800M runs used uniform accumulation, so their GRAM retain results may be slight underestimates.Source: H. Heterogeneous Accumulation · pp. 23–24
  • At 200M, heterogeneous accumulation raised GRAM’s core and retain ratios to baseline level while leaving forget essentially unchanged; FT-LoRA changed little across regimes.GRAM’s core rose 0.94 →0.99, retain 0.75 →1.11, and forget 0.70 →0.69; FT-LoRA retain rose 0.83 →0.90. Table 4 reports these ratios against the heterogeneously trained baseline.Source: H. Heterogeneous Accumulation · p. 24, Table 4 · p. 24

I. Equal-Compute Training Length Adjustment

  • Compute-matched comparisons adjust routed-model training lengths to match baseline FLOPS while preserving the core-to-auxiliary token mixture.FT-LoRA adapters were matched in parameter count to GRAM auxiliary modules, with each domain adapter sized to 10% of baseline MLP width.Source: §I · p. 25

J. Learning Rate and Batch Size Scaling

  • The appendix fits peak learning rate and effective batch size by grid-searching per-size optima and then fitting power laws across model sizes.Baseline models at 50M, 100M, 200M, and 400M parameters were searched using three seeds per configuration and a weighted loss combining core and auxiliary-category cross-entropy.Source: J. Learning Rate and Batch Size Scaling · p. 25
  • The fitted LR exponent is −0.308, close to Porian et al. (2024)’s −0.32, while predicted 5B batch sizes are approximately 2750 versus 2000.The paper attributes the batch-size gap primarily to holding Adam momentum parameters fixed at β1 = 0.9 and β2 = 0.95, whereas Porian et al. co-tune them with LR and batch size.Source: J. Learning Rate and Batch Size Scaling · p. 26
  • All methods use the same power-law-derived LR and batch size per model size because GRAM’s core MLP matches the baseline’s, while FT-Full and FT-LoRA scale finetune LR by plr and MaxEnt tunes unlearning LR separately.The paper lists the per-size values in Table 3 and gives the separate MaxEnt tuning in Appendix B.Source: J. Learning Rate and Batch Size Scaling · p. 26
  • Across 50M, 100M, 200M, and 400M models, optimal batch size trends upward and optimal learning rate trends downward with increasing model size.Figure 14 reports response contours based on mean validation cross entropy on core and auxiliary datasets, with optima selected by quadratic-surface fits.Source: Figure 14 · p. 26
  • Optimal learning rates and batch sizes are well approximated by power laws across 50M–400M models.Figure 15 shows observed and predicted optima; bars indicate the fit-surface curvature and the span of values with predicted loss no more than 0.1% above the optimum.Source: Figure 15 · p. 27

K. Sample Outputs From Different Methods

This section presents sample outputs and token-level loss patterns associated with capability removal, and defines compute-ratio estimation using pooled baseline learning curves. It also describes how GRAM modules support continuous capability access while core performance remains approximately unchanged.

K. Sample Outputs From Different Methods

  • Using a fixed prompt, the examples contrast retained and removed configurations; GRAM and FT-LoRA ablate capability-specific parameters, while filtering and MaxEnt compare model variants.GRAM and FT-LoRA are shown with three of four versus all four auxiliary subsets ablated; outputs use temperature 1.0 sampling.Source: K. Sample Outputs From Different Methods · p. 27

EXAMPLE 1: Alien Encounters

  • In the alien-encounter example, capability-removed configurations turn the mysterious whirring sound into something mundane rather than an alien spacecraft.GRAM and filtering produce a robot, while MaxEnt produces an empty pantry; the baseline and capability-retaining configurations produce spacecraft.Source: EXAMPLE 1: Alien Encounters · p. 27

EXAMPLE 2: A Deadline or Time Limit

  • Capability-removed configurations drop the deadline’s time pressure and relax into a leisurely party, while DEMix’s lone auxiliary expert degenerates through repetition.The baseline and capability-retaining configurations preserve the ticking-clock urgency.Source: EXAMPLE 2: A Deadline or Time Limit · p. 29

L. Capability Removal Concentrates Loss on Domain-Relevant Tokens

  • With capability ablated, loss increases concentrate on specialized domain tokens across virology, cyber, and nuclear examples under GRAM, FT-LoRA, and MaxEnt; activating the matching component returns those increases near zero.Function words, syntax, and domain-general vocabulary remain near zero; the examples are illustrative, while quantitative claims are established in Section 5.Source: L. Capability Removal Concentrates… · pp. 30, 32

M. Compute Ratio Methodology

  • Compute ratio expresses a variant’s step-equivalent loss relative to the baseline’s, with 1 meaning the variant reaches a loss attained only at the end of baseline training.Step-equivalents come from fitting and inverting the baseline loss-versus-step curve; this makes comparisons across datasets and model sizes more meaningful than raw loss differences.Source: M. Compute Ratio Methodology · p. 32
  • Compute ratios map variants’ final losses to step-equivalents on a shared power-law curve fitted to pooled baseline seeds, then divide by the baseline reference step-equivalent.The analysis uses three independent seeds and reports 90% t-intervals over seed means; pooling baseline seeds suppresses uncertainty in the shared reference.Source: M. Compute Ratio Methodology · pp. 32–33

N. Continuous Control of Capability Profiles

  • ≈0.98 for GRAM and ≈0.95 for FT-LoRA are the virology retain compute ratios at full module weight, while scaled capability compute ratios rise and core stays near one.In the 800M realistic-setting models, capability compute ratio rises monotonically as module weight increases from 0 to 1; core compute ratio stays approximately one throughout.Source: N. Continuous Control of Capability… · p. 33, Figure 16 · p. 34

Related papers

Loading 2607.08077v1…