PaperlayerView exact revision on arXiv

Source-linked AI summary

BOReFT: Manifold Steering of Language Models for Black-box Optimization

Dhruv Agarwal, Rico Angell, Kavitha Srinivas, Tahira Naseem, Horst Samulowitz, Willie Neiswanger, Andrew McCallum

arXiv:2609.33722v1cs.LGcs.AIcs.CL

TL;DR

Language-model search produces discrete proposals, while Bayesian optimization needs a continuous domain, and existing approaches do not search a continuous space decoded directly into solutions. BOReFT learns a low-dimensional intervention space in a frozen language model and applies Bayesian optimization to it. It finds more hidden Semantle targets than the next-best EM baseline and leads on two of three molecular objectives, though performance remains tied to training-data quality.

  • Problem

    Language models generate discrete, stochastic sequences, while Bayesian optimization requires a continuous domain, and existing approaches do not search a continuous space decoded directly into solutions.

    Source: §1 · pp. 1–2

  • Method

    BOReFT learns a low-dimensional intervention space in a frozen language model and searches that fixed domain with Bayesian optimization using external scores.

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

  • Results

    BOReFT recovers 8/15 held-out Semantle targets versus 3/15 for the next-best EM count before SFT, leads JNK3 and GSK3β before SFT, and trails three baselines on DRD2.

    Source: §5.2 · p. 8

  • Takeaways & Limitations

    Searching a learned manifold with Bayesian optimization can generalize beyond training-task performance and outperform other LLM search strategies on Semantle and molecule optimization.

    Source: §5.2 · p. 8

    Performance remains coupled to training-data quality, which may prevent access to high-scoring distributions such as on DRD2.

    Source: §5.2 · p. 8

Abstract

from arXiv · show

Language models are increasingly used as proposal models for black-box search, from program optimization to molecular design. Existing approaches typically improve proposals through iterative prompting or parameter updates, offering limited control over how completely and efficiently the model's search space is explored. Continuous optimization methods, such as Bayesian optimization, provide a principled way to search but require a suitable domain to operate over. To address this, we introduce BOReFT, which learns a compact, low-dimensional space of hidden-state interventions in a frozen language model, and uses this space as the search domain for Bayesian optimization with an external scoring function. Empirically, we find that the learned domain spans semantic regions and exhibits smoothness properties that support search. Theoretically, we show that semantic coverage and interpolation control the best score available in the learned space, and that decoding from this space yields a standard stochastic-bandit observation model for adaptive search. We evaluate BOReFT on the interpretable word search task "Semantle" and on three more real-world discovery tasks in de novo molecule property optimization. Compared to strong LLM baselines, BOReFT finds in Semantle a higher number of hidden targets and, on two out of three molecular objectives, achieves higher property scores. Consequently, our method provides a principled new bridge between discrete proposal spaces of LLM-based search and continuous black-box optimization.

1 INTRODUCTION

Language-model search proposes discrete candidates, while Bayesian optimization requires a continuous domain. BOReFT learns such a domain inside a frozen language model and searches it with external scores.

1 INTRODUCTION

  • Language-model search produces discrete, stochastic sequences, whereas Bayesian optimization requires a continuous decision domain.In many generative-search settings, external evaluations provide only scalar feedback.Source: §1 · p. 1
  • Prior approaches bridge the mismatch through prompting or candidate representations, rather than searching a continuous space decoded directly into solutions.The cited approaches include Yang et al. (2024), Agarwal et al. (2025a), Chen et al. (2024), Kristiadi et al. (2024), and Ranković & Schwaller (2025).Source: §1 · pp. 1–2
  • BOReFT learns a compact hidden-state intervention manifold in a frozen language model and searches that continuous domain with Bayesian optimization using external scores.Each code specifies an intervention that the frozen decoder maps to a distribution over discrete sequences.Source: §1 · p. 2
  • The paper analyzes representational and search gaps and evaluates BOReFT on Semantle and molecular property optimization.Its analyses examine semantic coverage, interpolation, stochastic-bandit observations, representation-training data, loss components, and downstream search.Source: §1 · p. 2

2 RELATED WORK

Related work covers Bayesian optimization with language models and activation interventions. BOReFT uses learned intervention coordinates as the search variables inside a frozen generator.

BO with language models

  • BOReFT learns continuous coordinates from target solutions and searches them inside the frozen generator.Related approaches include latent-space BO (Gómez-Bombarelli et al., 2018; Maus et al., 2022), LLM feature or kernel methods, prompt-based proposals, and preference-trained actors.Source: §2 · p. 2

Activation steering and representation interventions

  • BOReFT uses activation interventions as black-box search variables, building on steering methods and low-rank task-adaptation interventions.The passage names ReFT and LoReFT (Wu et al., 2024) among prior intervention methods.Source: §2 · p. 2

3 BOREFT: BAYESIAN OPTIMIZATION VIA REPRESENTATION FINETUNING

BOReFT learns a bounded continuous intervention space in a frozen language model, then uses Bayesian optimization to search it against an external objective. The space is trained from target-description pairs with losses designed to recover targets, constrain code spread, and retain richer output behavior.

3.1 LOW-RANK INTERVENTIONS AS A SEARCH SPACE

  • BOReFT uses low-rank interventions on a frozen model, sharing intervention matrices across targets while assigning each target a code.For fixed activation, the intervention is affine in the code, and the remaining model layers decode from the intervened activation as usual.Source: §3.1 · pp. 2–3

3.2 LEARNING THE SEARCH SPACE

  • Training maps each target-description pair to a Gaussian code posterior and combines reconstruction, variational prior regularization, and on-policy self-distillation.Reconstruction anchors targets, the prior limits posterior movement, and self-distillation encourages a richer distribution of outputs beyond the target itself.Source: §3.2 · p. 3
  • Figure 2 presents a PCA projection of the 3072 Semantle training posterior means, with color encoding interpolation weight.The figure’s caption identifies it as an interpolation view of the learned code space.Source: Figure 2 · p. 4
  • On-policy self-distillation uses rollouts from the posterior, with the intervened model as student and a frozen, description-conditioned model as teacher.The teacher supplies token-level supervision, and forward KL is used to cover all behaviors exhibited by the teacher.Source: §3.2 · p. 4
  • After training, the search domain is a bounded continuous box formed from the component-wise ranges of learned posterior means.The domain includes more than the training means themselves.Source: §3.2 · p. 4

3.3 SEARCHING THE LEARNED SPACE WITH BAYESIAN OPTIMIZATION

  • At inference, Bayesian optimization selects codes in a fixed learned domain, the frozen decoder generates sequences, and an external black-box objective scores them.A Gaussian-process surrogate is fit to observed code–score pairs, and the next code maximizes an acquisition function.Source: §3.3 · pp. 4–5

4 WHAT CAN SEARCH OVER THE LEARNED CODE SPACE RETURN?

The section separates limitations imposed by the learned code space from those imposed by the search procedure. It relates the representational gap to semantic coverage and interpolation, then gives conditions under which adaptive search has standard bandit observations and discovery guarantees.

4.1 THE SEARCH PROBLEM AND A DECOMPOSITION OF ITS REGRET

  • The regret splits into a representational gap fixed by the model, rank, and training, and a search gap controlled by the inference procedure.The learned code space sets the best expected score available; further search cannot reduce the representational gap.Source: §4.1 · p. 5, §4 · p. 5

4.2 HOW TRAINING SHAPES THE SPACE THAT IS SEARCHED

  • Reconstruction anchors training targets but does not ensure held-out reachability; the prior bounds the search domain, while low expected forward KL makes the student track the teacher on sequence sets.The prior bounds the domain according to the intervention rank and prior scale.Source: §4.2 · p. 5
  • Interpolating between learned anchors can produce gradual semantic changes and meanings absent from individual training targets, but this alone does not guarantee similar task values.The Semantle-space illustration uses temperature-sampled decodes along linear interpolations; task-value smoothness requires an additional condition.Source: §4.2 · p. 5
  • When task scores are Lipschitz in semantic embeddings, coverage and interpolation error bound the representational gap; this condition holds for Semantle but remains an assumption for molecular predictors.The bound scales with the task’s Lipschitz constant.Source: §4.2 · p. 6

4.3 GUARANTEES FOR ADAPTIVE SEARCH

  • Each query is an unbiased, conditionally sub-Gaussian observation for any past-dependent proposal policy; standard bandit and Bayesian optimization guarantees require further assumptions and control the best sequence emitted.A cumulative-regret bound translates to a discovery guarantee; subtracting coverage and interpolation errors additionally requires the Lipschitz condition.Source: §4.3 · pp. 6–7

5 EXPERIMENTS

The experiments evaluate BOReFT on Semantle and molecular property optimization against LLM-search baselines, including after supervised fine-tuning on shared training data. BOReFT performs strongly overall, while ablations and coverage/interpolation analyses examine what supports search quality and how training-data quality limits it.

5.1 TASKS AND EVALUATION PROTOCOL

  • The evaluation uses 500-query budgets for Semantle and molecular optimization, with word and molecule warmstarts, and compares LLM-search strategies both before and after SFT on shared training data.Semantle uses held-out and training targets; molecular objectives are DRD2, GSK3β, and JNK3.Source: §5.1 · p. 7, §5 · p. 7

5.2 RESULTS AND ANALYSES

  • BOReFT recovers 8/15 held-out Semantle targets, with mean similarity 0.892 and novelty rate 0.691; the next-best exact-match count is 3/15 before SFT and 2/15 after SFT.Discrete BO finds no held-out words by construction, while finding every training word in the fewest evaluations.Source: §5.2 · p. 8, Table 1 · p. 8
  • On molecular objectives, BOReFT leads JNK3 before SFT with mean 0.370 and best score 0.940, and has the highest GSK3β score at 0.670; it trails three baselines on DRD2.After SFT, BOReFT ranks in the top two on all properties as all baselines except MiGrATe drop in performance.Source: §5.2 · p. 8, Table 2 · p. 9
  • BOReFT’s performance remains coupled to representation-training data quality, which may limit access to high-scoring distributions such as DRD2 and motivates continual learning as future work.The evaluation uses Semantle word search and three molecular objectives, comparing search baselines before and after SFT on shared data.Source: §5.2 · p. 8
  • Full BOReFT finds 14/30 hidden targets, versus 9/30 without self-distillation, 7/30 without reconstruction, 3/30 without variational training, and 5/30 without the shared encoder.These ablations indicate that each loss component and the encoder’s geometric bias matter for learning a search space.Source: §5.2 · p. 9, Figure 4 · p. 9
  • On 928 held-out words, median coverage error falls from 0.89 at N = 1 to 0.59 at N = 3072 as search improves; on 40 pairs, ∆interp is 0.338.Removing self-distillation, variational training, reconstruction, or the shared encoder raises ∆interp to between 0.355 and 0.396 and reduces search quality.Source: §5.2 · p. 9, Figure 4 · p. 9

6 CONCLUSION

BOReFT learns a low-dimensional continuous space of representation interventions in a frozen language model and searches it with Bayesian optimization. The paper characterizes when the space contains high-quality solutions and reports experimental gains over baselines on Semantle and molecular property optimization.

6 CONCLUSION

  • BOReFT learns a low-dimensional continuous intervention space in a frozen language model, then uses Bayesian optimization to search it against a black-box objective.The analysis characterizes conditions for high-quality solutions and well-behaved adaptive search; experiments report gains over baselines on Semantle and molecular property optimization.Source: §6 · p. 9

AI USE STATEMENT

  • The paper reports using generative AI to help produce task definitions and categories, assist theory proofs, and write or edit code; outputs were manually reviewed and edited.The authors take responsibility for the final content and artifacts.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The paper states that assumptions and full proofs are in §4 and §A, method details in §3, and model, data, and evaluation details in §B; anonymized code is supplementary.Source: REPRODUCIBILITY STATEMENT · p. 10

A PROOFS AND EXPERIMENTAL DEFINITIONS

A.1 RECONSTRUCTION ALONE GIVES NO HELD-OUT COVERAGE (PROPOSITION A.1)

Proposition A.1 shows that reconstructing training targets, even with smooth decoder maps, does not guarantee that held-out sequences are reachable at useful confidence. Its construction leaves held-out coverage without any positive lower bound based only on the number of targets and intervention rank.

A.1 RECONSTRUCTION ALONE GIVES NO HELD-OUT COVERAGE (PROPOSITION A.1)

  • Reconstruction and smoothness provide no positive lower bound on held-out oracle coverage as a function of training-target count and rank.For a finite held-out set, the proposition measures coverage by the fraction reliably elicitable at confidence θ.Source: §A.1 · pp. 14–15
  • The construction fits every training target at confidence θ while keeping each held-out sequence below θ at every code.It penalizes the first token where a held-out sequence leaves the training-prefix tree, while an alternative continuation remains available.Source: §A.1 · pp. 14–15
  • The result applies to smooth code-to-logit decoders, a class broader than transformers with rank-r interventions.The prefix-tree construction also covers held-out sequences sharing long prefixes with training targets.Source: §A.1 · pp. 14–15

A.2 THE PRIOR AND DISTILLATION TERMS (LEMMA A.2)

Lemma A.2 formalizes how the prior constrains the learned search domain and how self-distillation relates the posterior-averaged student distribution to the teacher. The latter bound applies to sets of sequences, including disjoint sets.

A.2 THE PRIOR AND DISTILLATION TERMS (LEMMA A.2)

  • The prior bounds the posterior means and hence the search box: if DKL(qi∥p0) ≤ κi, then B lies within [−R,R]^r for R = σ0√(2κmax).The result also constrains posterior spread relative to the prior; the covariance contribution to the divergence is nonnegative and grows as eigenvalues approach zero or infinity.Source: §A.2 · p. 15
  • If the expected sequence-level forward KL ηi is small, the posterior-averaged student distribution differs from the teacher by at most √(ηi/2) on any set of terminated sequences.For disjoint sequence sets, the same bound controls the difference in their assigned masses.Source: §A.2 · p. 15
  • The lemma’s ηi is an expected sequence-level forward KL, not the implemented loss that averages per-token divergences over teacher samples, student rollouts, and the target.The proof obtains the distributional bound using convexity of relative entropy and total variation.Source: §A.2 · pp. 15–16

A.3 PROOF OF PROPOSITION 4.1

The appendix explains how semantic interpolation and coverage enter the representational-gap bound, and states the observation-model conditions used for adaptive search. It also records implementation choices for the evaluated models and prompts.

A.3 PROOF OF PROPOSITION 4.1

  • The interpolation error measures displacement of mean decoded semantics from the interpolant, while coverage error is minimized at a nearest training target in the vertex protocol.The interpolation displacement is no larger than expected spread around the interpolant; restricting coverage to vertices can increase coverage error, and the two errors trade off as interpolation weights are added.Source: §A.3 · pp. 16–17

A.4 PROOFS OF PROPOSITION 4.2 AND COROLLARY 4.3

  • A query gives an unbiased, conditionally sub-Gaussian observation of expected score; combining regret with the representational bound additionally requires score Lipschitzness in the embedding.That condition holds for Semantle but remains an assumption for the molecular predictors.Source: §A.4 · p. 17

B.1 MODELS AND REPRESENTATION TRAINING

  • Semantle uses a frozen Llama-3.2-1B-Instruct model with dictionary-like descriptions, while molecular representation training uses a frozen MiST Qwen2.5-3B checkpoint and ChEBI-20 pairs.The molecular corpus contains 8,000 pairs, with a random 1,024-pair subsample used for the primary experiments.Source: §B.1 · p. 18

B.2 PROMPTS

  • The search prompts differ by task: Semantle generates an English word, whereas molecular search uses an objective instruction and a SMILES completion prefix.Self-distillation supplies the target description to the teacher while the student remains conditioned on the reconstruction prompt.Source: §B.2 · pp. 18–19

B.3 EVALUATION PROTOCOLS

The protocols allocate 500 evaluations to each task run and specify task-specific warm starts and sampling. The listed baseline comparison uses common budgets and warm-start observations unless otherwise noted.

B.3 EVALUATION PROTOCOLS

  • Semantle evaluates ten hidden targets with three runs each and 500 evaluations per run, using ten shared labeled warm-start pairs and greedy BOReFT decoding.The ten targets comprise five representation-training words and five held-out reconstruction-test words.Source: §B.3 · p. 19
  • Figure 5 plots best-so-far molecular property at each verification, with before-SFT and after-SFT runs in the top and bottom rows and a shared first ten evaluations.BOReFT is shown only after SFT; its scores are repeated in Table 2’s before-SFT columns.Source: Figure 5 · p. 20
  • Molecular optimization evaluates DRD2, GSK3β, and JNK3 over five seeds and 500 property evaluations, with ten shared Sobol warm-start proposals decoded at temperature 1.Invalid SMILES score 0, and all three oracle scores are class-1 probabilities in [0,1].Source: §B.3 · p. 20

B.4 SEARCH BASELINES

  • Search methods generally share the objective-evaluation budget and warm-start observations unless their baseline row specifies otherwise.Prompt-based methods sample at temperature 1; molecular completions use the MiST SMILES format.Source: §B.4 · p. 20

B.5 PROPOSAL BEHAVIOR ON SEMANTLE

The Semantle baselines differ in how their proposals evolve, but exploration or rising proposal similarity does not consistently yield hidden-word recovery. Scalar-only feedback also limits the additional supervision available to SDPO-TTT.

SDPO-TTT

  • SDPO-TTT repeatedly proposes “cloud”; mean proposal similarity falls from 0.633 in the first 50 proposals to 0.630 in the last 50, with no exact matches.“Cloud” accounts for at least 0.29 of post-warm-start proposals in every run; best similarity is 0.763 on training targets and 0.709 on held-out targets.Source: §B.5 · pp. 20–21

Search-method settings

  • AutoDiscovery is configured for branch-based proposals, while MiGrATe scores four new completions per round and reuses greedy group members from history.AutoDiscovery expands eight completions per node and scores one, leaving the rest as untried children.Source: Table 3 · p. 21

SDPO-TTT

  • SDPO-TTT’s scalar-only Semantle feedback may leave its teacher close to the unconditioned proposal distribution, supplying little additional supervision.The latest word and its scalar cosine similarity are the teacher’s only additional context, and proposal similarities stay flat across the budget.Source: §B.5 · p. 21

AutoDiscovery

  • AutoDiscovery’s proposals exceed the best score observed so far on only 0.006 of search steps, despite a 0.364 repetition rate and no exact matches.Final similarities are 0.772 on training targets and 0.757 on held-out targets.Source: §B.5 · p. 21

MiGrATe

  • MiGrATe’s mean proposal similarity rises from 0.637 to 0.742, but exact-match recovery remains 0/15 on training targets and 1/15 on held-out targets.Its search takes about 3.5× as long as BOReFT at the same 500-score budget; this is search time only and excludes BOReFT representation training.Source: §B.5 · pp. 21–22

B.6 SEARCH ON TRAINING TARGETS

These Semantle analyses compare training-target performance and probe how repetition, search-domain geometry, interpolation, decoding temperature, GP kernels, and acquisition choices affect search. Results favor the primary mean-box, greedy-decoding, Matérn-2.5, and log expected improvement settings among those tested.

B.6 SEARCH ON TRAINING TARGETS

  • BOReFT records 8/15 held-out exact matches, exceeding every fine-tuned baseline; discrete BO reaches 15/15 on training targets because those words are its candidate set.After fine-tuning, OPRO’s training-target exact-match count falls from 7/15 to 4/15, while random sampling rises from 0/15 to 3/15.Source: §B.6 · p. 22

B.7 PROPOSAL REPETITION

  • After fine-tuning, novelty drops for AutoDiscovery, BOPRO, and OPRO, while SDPO-TTT and MiGrATe stay near their base novelty rates.The respective drops are 0.736 to 0.154, 0.798 to 0.203, and 0.708 to 0.212.Source: §B.6 · p. 22

B.8 SEARCH-DOMAIN CONSTRUCTION

  • Expanding the mean box sharply reduces recovery: at k = 0.1 its volume is 4.71×10^5 times larger, with 3/15 training and 4/15 held-out exact matches.At k = 0.5 and k = 1, no exact matches are recovered; full expansion yields mean best similarities of 0.726 on training targets and 0.707 on held-out targets.Source: §B.8 · p. 23, Table 6 · p. 24
  • The covering ellipsoid has 4.08×10^-4 of the mean box’s volume and recovers 5/15 training and 5/15 held-out targets, versus 6/15 and 8/15 for the mean box.Replacing the box with the ellipsoid does not improve exact-match recovery; enlarging the box to include posterior spread reduces recovery and raises repetition.Source: §B.8 · p. 23, Table 6 · p. 24

B.9 GEOMETRIC PRIOR

  • Figure 6’s PCA plots show semantic representations used as the geometric prior, with semantic categories colored within each task’s row.The rows represent molecular optimization and Semantle; the columns compare Qwen3-Embedding-0.6B with the task language model.Source: §B.9 · p. 23, Figure 6 · p. 24

B.10 INTERPOLATION IN THE LEARNED CODE SPACE

  • BOReFT’s interpolation distance is 0.338, below the endpoint-only reference of 0.365; removing reconstruction or the shared encoder raises it to 0.396 or 0.382.Removing self-distillation or variational training yields 0.355 and 0.356, respectively.Source: §B.10 · p. 25

B.11 EFFECT OF TEMPERATURE ON SEARCH

  • At T = 0, exact-match recovery is highest at 6/15 training and 8/15 held out; increasing temperature to 0.75 lowers repetition but does not improve recovery.Repetition falls from 0.505 and 0.452 under greedy decoding to 0.234 and 0.231 at T = 0.75.Source: §B.11 · p. 25, Table 7 · p. 25

B.12 EFFECT OF GP KERNEL ON SEARCH

  • Matérn-2.5 recovers the most targets among tested GP kernels, at 6/15 training and 8/15 held out; Matérn-0.5 and RBF each fall to 4/15 held out.Mean best similarity follows the same ordering.Source: §B.12 · p. 25, Table 8 · p. 25

B.14 TRAINING ABLATION NUMBERS

The appendix reports how training ablations and intervention rank relate to Semantle search and interpolation. It also distinguishes BOReFT’s learned search domain from prior approaches that adapt models or describe existing activation structure.

B.14 TRAINING ABLATION NUMBERS

  • Table 10 compares exact matches and mean best similarity across training ablations, alongside interpolation distance measured against an endpoint-only reference of 0.365.The supplied passage describes the table’s measures but does not state the ablation outcomes.Source: §B.14 · p. 26, Table 10 · p. 27

B.15 EFFECT OF INTERVENTION RANK

  • Rank 64 gives the strongest exact-match result, 14/30, while performance is non-monotone in rank and dips sharply at rank 32.Figure 8 reports mean best semantic similarity and exact-match rate pooled over ten targets and three seeds.Source: §B.15 · p. 26, Figure 8 · p. 27

C ADDITIONAL RELATED WORK

  • Unlike test-time search or adaptation methods, BOReFT learns its search domain before inference and then holds both generator and domain fixed.It learns interventions rather than recovering existing verbalizable activation structure; its coordinates are selected by an external optimizer.Source: C ADDITIONAL RELATED WORK · p. 27

Related papers

Loading 2609.33722v1…