Efficient Inference-time Control and Alignment

Abstract

Modern foundation models are typically trained in three broad stages. First, large-scale pre-training is performed using self-supervised learning on massive corpora. Second, models are adapted through mid-training using supervised fine-tuning or instruction tuning on labeled datasets. Finally, a post-training stage is often applied using preference data and reinforcement learning in order to align the model and improve its safety, reliability, and usefulness.

Although effective, post-training methods can be computationally expensive and inflexible once large models are deployed. This thesis explores an alternative paradigm: enforcing behavioral objectives at inference time rather than modifying model parameters during post-training. In this approach, smaller modular control models are combined with a base model to shape predictions during the decision process. Our aim is to design alignment mechanisms that are both mathematically grounded and empirically strong while remaining computationally efficient and easy to deploy.

We apply this perspective of inference-time control to three problems. First, we address reliability in neural classifiers. We introduce PreLoad, an inference-time mechanism that mitigates arbitrarily high confidence on inputs that lie outside the training support while preserving accuracy and training efficiency.

Second, we study reward-guided text generation (RGTG) in large language models as a form of inference-time alignment. We show that stable reward-guided decoding requires carefully designed token-level reward models and propose two algorithms, PARGS and FaRMA, that enable effective reward-guided generation.

Third, we address the computational cost of RGTG and propose an efficient algorithm that adds only a minor overhead during inference while preserving the performance and benefits of reward-guided decoding.

Together, these results demonstrate that inference-time control provides a flexible and computationally efficient framework for shaping the behavior of modern neural systems. By decoupling representation learning from the decision-time objectives, this work introduces new tools for improving the reliability, alignment, and efficiency of large-scale machine learning models without retraining them.