Bifrost: Steering Strategic Trajectories to Bridge Contextual Gaps for Self-Improving Agents

Autonomous agents excel in self-improvement through reflection and iterative refinement, which reuse successful task trajectories as in-context examples to assist subsequent reasoning. However, shifting across tasks often introduces a context mismatch. Hence, existing approaches either discard the trajectories or manipulate them using heuristics, leading to a non-negligible fine-tuning cost or unguaranteed performance. To bridge this gap, we reveal a context-trajectory correlation, where shifts of context are highly parallel with shifts of trajectory. Based on this finding, we propose BrIdge contextual gap FoR imprOvised trajectory STeering (Bifrost), a training-free method that leverages context differences to precisely guide the adaptation of previously solved trajectories towards the target task, mitigating the misalignment caused by context shifts. Our trajectory adaptation is conducted at the representation level using agent hidden states, ensuring trajectory transformation accurately aligns with the target context in a shared space. Across diverse benchmarks, Bifrost consistently outperforms existing trajectory reuse and finetuned self-improvement methods, demonstrating that agents can effectively leverage past experiences despite substantial context shifts.

💡 Research Summary

The paper introduces Bifrost, a training‑free framework that enables large language model (LLM) agents to reuse successful past trajectories even when the target task resides in a substantially different context. The authors first identify a “context‑trajectory correlation”: when a task’s context shifts, the latent hidden‑state representation of the agent shifts in a direction that is highly parallel to the shift observed in the corresponding successful trajectory. Under the Linear Representation Hypothesis, they model an LLM as an embedding‑unembedding system where the probability of the next token is proportional to the inner product between the hidden state h(x) of the input x and the unembedding vector g(y) of a candidate token y. By assuming that shared concepts across tasks occupy independent subspaces, they show that a context shift can be expressed as an additive vector ¯h_W scaled by a scalar α (Equation 3). Consequently, the logit for any binary concept changes linearly with α (Theorem 2), which they empirically validate on paired samples from the AQUA and GSM8K datasets.

Bifrost operationalizes this insight by computing, for each layer ℓ, the average hidden state ¯h_ℓ over all stored trajectories C = {(q_i, a_i)}. For a new query \hat{q}, the model extracts its hidden state h_ℓ(\hat{q}) and forms a steering vector Δ_ℓ = h_ℓ(\hat{q}) – ¯h_ℓ. The hidden states are then adjusted as h_ℓ^s = h_ℓ + α·Δ_ℓ for a user‑specified strength α, and the agent proceeds with in‑context learning using the prompt p = C ◦ \hat{q}. This procedure requires no gradient updates; it merely perturbs the internal representations in a direction that aligns past demonstrations with the target context. The authors note that middle‑to‑late layers are most effective for steering, consistent with prior work on reasoning‑oriented representations. They also demonstrate that alternative projection techniques (PCA, sparse autoencoders) can compress Δ_ℓ without sacrificing performance, offering a memory‑efficient variant.

The theoretical contribution frames in‑context learning as Bayesian inference. Past trajectories define a posterior distribution p(ϕ|C) over latent concepts ϕ. By steering hidden states, Bifrost shifts the posterior mean toward the target context while reducing posterior variance, thereby tightening the risk bound relative to a vanilla LLM that solves the task without any contextual guidance. The analysis extends to multi‑concept tasks, where the independence and causal separability assumptions guarantee that steering one concept does not interfere with others.

Empirically, Bifrost is evaluated on three domains: question answering (AQUA), mathematical reasoning (GSM8K), and code generation. Experiments span model sizes from 7 B to 70 B parameters. Baselines include (1) naïve trajectory reuse, (2) fine‑tuned self‑improvement, and (3) heuristic prompt engineering methods. Across almost all settings, Bifrost yields 2–7 percentage‑point accuracy gains, with the most pronounced improvements observed when transferring from GSM8K to AQUA—an extreme context shift scenario. Ablation studies confirm that performance scales smoothly with α and that compressed steering vectors retain most of the benefit.

In summary, Bifrost makes four key contributions: (i) it uncovers and formalizes the context‑trajectory correlation; (ii) it proposes a simple, training‑free hidden‑state steering mechanism that precisely adapts past trajectories to new contexts; (iii) it provides a Bayesian risk analysis showing reduced uncertainty and tighter generalization bounds; and (iv) it validates the approach on diverse benchmarks, consistently outperforming existing trajectory‑reuse and fine‑tuning baselines. The work opens avenues for further research on multi‑modal contexts, online adaptation, and more sophisticated subspace steering techniques.