Token-Level Accept or Reject: A Micro Alignment Approach for Large Language Models

1 Introduction

The alignment of Large Language Models (LLMs) with human values and preferences has emerged as a crucial challenge in AI development [Wang et al., 2023b]. The alignment is extremely important for LLMs that operate safely and ethically. Among various alignment approaches, Reinforcement Learning from Human Feedback (RLHF) [Ouyang et al., 2022] and Direct Preference Optimization (DPO) [Rafailov et al., 2024b] have emerged as two dominant paradigms. RLHF fine-tunes language models using a reward model trained on human preference datasets, while DPO directly optimizes the models through pairwise comparisons without relying on explicit reward modeling.

∗Corresponding Author <tangbo@mail.ustc.edu.cn>.

Although RLHF, DPO and their variants [Perez et al., 2023; Bai et al., 2022b; Lee et al., 2023; Zeng et al., 2024; Azar et al., 2024; Wang et al., 2023a] have demonstrated impressive capabilities in aligning performance, they face a critical challenge: excessive computational resource consumption. These methods require fine-tuning billion or even hundred-billion parameter of language models, typically demanding hundreds of GPU hours and substantial memory resources. Such computational intensity poses a significant barrier for real-world applications [Anwar et al., 2024], particularly in scenarios requiring rapid alignment updates or resource-constrained environments.

To address this challenge, we pose a fundamental question: Can we develop a micro alignment approach that operates independently of the language model while maintaining excellent alignment performance?

As shown in Figure 1, unlike aligning operated on the language model, our key insight is that the alignment process can be simplified into accepting or rejecting a token using a streamlined alignment model. Specifically, the alignment model is the first to be implemented as a micro fullyconnected network. Given the prompt, the partially generated response, and a candidate set of tokens that are sorted by sampling probabilities from a supervised fine-tuned model, our alignment model sequentially determines whether to accept or reject each candidate token during the generation process. This approach transforms the alignment task into a simple binary classification problem, significantly reducing the computational overhead while ensuring effective alignment performance.

A concurrent work called Aligner [Ji et al., 2024a] shares a similar motivation of decoupling alignment from the upstream language model. However, our approach differs significantly in both scale and methodology. While Aligner employs a seq2seq model with 2B-13B parameters, our method utilizes a micro fully-connected network with merely millions of parameters. Moreover, our execution process is fundamentally simpler, focusing on token-level accept/reject decisions rather than reformulating entire output sentences.

In summary, we propose Micro token-level Accept-Reject Aligning (MARA) approach, which offers three significant advantages:

• Computation-friendly: Our approach requires only a three-layer fully-connected network as the alignment model, reducing the computation overhead significantly. Specifically, for aligning an 8B-parameter language model, MARA requires training only a 4Mparameter alignment model, whereas RLHF, DPO, and Aligner necessitate training language models with over 20M parameters, even when employing parameterefficient Low-Rank Adaptation (LoRA).
• Effectiveness: While operating at a micro level, our approach achieves superior alignment performance. Specifically, with Llama 3.1-8B as the base model, MARA demonstrates substantial improvements across different benchmark datasets: +31.8% over RLHF, +18.8% over DPO, and +8.8% over Aligner.
• Compatibility: The decoupled architecture of MARAtrained alignment models enables seamless integration with various LLMs, i.e., an alignment model trained on one LLM can be effectively transferred to other LLMs while maintaining strong performance. For instance, an alignment model trained on Mistral-7B-v0.3 enhances the alignment performance of Llama 3-8B by 25.5% on average across three evaluation datasets.

2 Preliminaries

This section introduces the standard RLHF framework, which consists of two primary phases: supervised fine-tuning (SFT) and reinforcement learning-based optimization. We first formalize these phases and then present a token-level decomposition of the alignment process.

2.1 RLHF Framework

SFT Phase: The pre-trained language model is initially finetuned on high-quality human demonstrations to generate appropriate responses. This process is optimized through the following objective:

where x and y denote the input prompt and model response, respectively, sampled from the supervised fine-tuning dataset DSFT . πref denotes the fine-tuned language model serving as the reference model for subsequent optimization.

RL-based Optimization Phase: Following SFT, the model undergoes reinforcement learning optimization to align with human preferences. The objective function for this phase is:

where πθ denotes the model undergoing RL optimization, and DRL represents the optimization dataset. The coefficient λ controls the KL divergence penalty between the reference model and the aligned model.

The reward model r(x, y) evaluates the alignment between model outputs and human preferences. Given a preference dataset D containing triples (x, yw, yl), where x is the prompt, yw is the preferred response, and yl is the less preferred response, the reward model is trained to minimize:

where σ denotes the Sigmoid function.

2.2 Token-Level Decomposition

At its core, the RLHF alignment process can be decomposed into a sequence of token-level decisions. The sentence generation process is formally expressed as:

where Ti denotes the vocabulary space of available tokens at the i-th position in the output sequence y, and H denotes the sequence length.

This token-level decomposition reveals the possibility of performing alignment at the granularity of individual tokens, enabling more precise control over the generation process. Furthermore, this granular perspective can be refined into explicit accept-or-reject decisions for each token, which motivates our proposed method detailed in the following sections.

3 Our Approach

In this section, we present MARA, a novel token-level acceptreject approach for LLM alignment. Building upon the foundation established by the SFT phase, MARA introduces a language model-independent alignment mechanism, replacing the original language model optimization in RLHF. Specifically, we first formalize the alignment task as a Markov Decision Process (MDP), then detail our accept-reject mechanism, and finally discuss implementation specifics.

3.1 MDP Formulation of Alignment

Given the reference model πref from the SFT phase, we formulate the alignment process as a MDP, characterized by the tuple (S, A, ρ, P, r), where:

• State space (s ∈ S): For the i-th token generation at time step τ, the state sτ = {x, y<i, tki } comprises the input prompt x, previously generated response y<i, and the current candidate token tki ∈ Ti, where Ti denotes a truncated candidate token set (detailed below), and k indexes the candidate token in Ti.
• Action space (a ∈ {0, 1}): Binary action where atki = 1 indicates accepting the candidate token tki to concatenate with the generated response y<i, while atki = 0 indicates rejection. To facilitate understanding, we use atki instead of aτ referring to the action on the candidate token.
• Initial state distribution (ρ): Defines the distribution over initial states.
• Reward function (r : S × A → R) will be detailed in Section 3.3.
• State transition (P): The dynamics follow:

To enhance computational efficiency, we employ a hybrid vocabulary truncation strategy combining two widely-used methods: top-k sampling, which retains only the k most probable tokens, and top-p (nucleus) sampling, which selects the smallest set of tokens whose cumulative probability exceeds p. This combination clips the vocabulary space Ti at each generation step to form a reduced candidate set Ti. This hybrid approach ensures both diversity and quality of the candidate tokens while maintaining computational tractability.

3.2 Token-level Accept-Reject Aligning Mechanism

Unlike traditional RLHF that fine-tunes the entire language model, our approach introduces a micro accept-reject model that operates at the token level. This model is implemented as a compact three-layer fully connected neural network. The alignment process follows four key steps:

Step 1: Candidate Generation. Given the prompt x and the previously generated tokens y<i, we use the reference model πref to generate a candidate token set Ti =

{t1i , t2i , . . . , t|Ti| i } through hybrid sampling that combines top-p and top-k strategies.

Step 2: Probability-based Sorting. Sort the candidate token set Ti in descending order according to their conditional probabilities πref (tki |x, y<i), resulting in an ordered set ˜Ti.

Step 3: Token Evaluation. For the first token in the ˜Ti, apply the accept-reject model to evaluate whether to accept or reject the token. The input of the accept-reject model consists of the last hidden state of the token sequence: prompt x, previously generated tokens y<i, and the first candidate token in ˜ Ti. The output is a binary decision: accept or reject. If accepting the token, add it to the output sentence y, and return to Step 1 until the EOS token is selected or the output length reaches its predefined limit. Otherwise, proceed to Step 4.

Step 4: Candidate Update. Upon rejection, remove the token from the candidate token set ˜Ti, and return to Step 3.

To maintain generation quality, we order candidate tokens by their sampling probabilities in descending order, as higher probability tokens typically better align with the reference model's learned patterns, thus reducing the risk of generating grammatically incorrect or contextually inconsistent content.

For cases when all tokens in the candidate set are rejected by the accept-reject model, we employ a fallback mechanism to ensure generation continuity. Specifically, we enforce the acceptance of the final candidate token, formally expressed as atki ∈ {1} when k = ∣∣∣˜Ti ∣∣∣. This strategy ensures generation proceeds by selecting lower-probability tokens when necessary, thereby balancing quality and diversity.

According to the above token-level aligning approach, the sampling probability for any token tki , representing the k-th candidate token in the candidate set ˜Ti, is formalized as:

Note that we still use πθ to denote the accept-reject model parameterized by θ.

From above, the complete aligned response generation process can thus be expressed as:

3.3 Implementation

Reward Design

A key challenge in MARA is bridging the gap between sentence-level rewards in standard RLHF and token-level decisions. Our reward function is designed with two objectives: (1) maximizing the reward model score r(x, y) to maintain alignment with human preferences as in RLHF, and (2) preventing reward hacking by constraining the token distribution divergence between the aligned model and the reference model, thereby avoiding unnatural token selections [Eisenstein et al., 2023].

Specifically, for any action atki at state sτ , we define the token-level reward as:

where r(x, y) denotes the evaluation score by reward model trained by Eq. 3.

This formulation represents a token-level decomposition of the RLHF objective in Eq. 2. We adopt the distributed Soft Actor-Critic (SAC) algorithm [Haarnoja et al., 2018] instead of the conventional Proximal Policy Optimization (PPO) algorithm used in RLHF due to SAC's higher computation efficiency, especially in distributed training environment. PPO requires synchronization between distributed nodes for sample collection and model updates, creating waiting periods. In contrast, SAC enables continuous sample collection across all nodes with real-time parameter updates via a global replay buffer, eliminating waiting periods.

Training Objectives

The alignment model serves as the actor in our framework, with its loss function defined as:

where H represents the entropy term of the alignment model that promotes diverse token selection, while V1(s) and V2(s) are two critic heads for robust value estimation and overfitting prevention.

The critic network, parameterized by φ, is trained using the following loss:

where πφ denotes the critic network with parameters φ. V ′ 1 (s′ ) and V ′ 2 (s′ ) denote the two heads of target critic network.

The entropy coefficient αh is automatically adjusted according to:

where H = 2 is the minimum entropy threshold based on the action space dimension.

4 Experiments

We conduct comprehensive experiments to evaluate MARA's effectiveness across diverse datasets, reward models, and base LLMs.

4.1 Experimental Setup

Datasets. Our experiments mainly leverage three comprehensive evaluation benchmarks for utility, safety and ethical assessment: PKU-SafeRLHF (SafeRLHF) [Ji et al., 2024b], which serves as our primary training dataset and contains 83.4K preference pairs with safety meta-labels and human preferences for helpfulness and harmlessness; BeaverTails [Ji et al., 2024c], covering 14 safety categories including sensitive topics like abuse and political discussions; and HarmfulQA [Bhardwaj and Poria, 2023], comprising 10 themes of ChatGPT-generated conversations for evaluating responses in potentially harmful scenarios.

Upstream LLMs. We evaluate our alignment model on two widely-adopted open-source LLM families with various model scales and versions. Specifically, the Llama family includes Llama-3-8B, Llama-3.1-8B, Llama-3.2-3B, and Llama-3.2-1B [AI@Meta, 2024], while the Mistral family comprises Mistral-7B-v0.1, Mistral-7B-v0.2, and Mistral7B-v0.3 [Jiang et al., 2023]. Due to computational constraints, we do not include larger models such as Llama-370B in our experiments. For brevity, we omit the suffix 'Instruct' from all model names.

Evaluation Metrics. Following the methodology of aligner [Ji et al., 2024a], we adopt the preference rate metric to evaluate model performance across two critical dimensions: helpfulness and harmlessness, which respectively assess the utility and safety aspects of generated responses. The preference rate is defined as:

where Nw, Ne, and Nl denote the numbers of wins, ties, and loses respectively in pairwise comparisons between different alignment approaches.

Reward Models. Our alignment training utilizes a combination of two specialized models: beaver-7b-v1.0-reward for utility evaluation and beaver-7b-v1.0-cost [Dai et al., 2024] for safety evaluation (where a negative cost score indicates a safe response). The final reward r(x, y) is computed as:

where R(x, y) and C(x, y) denote the evaluation scores from the reward and cost models respectively, and αr and αc are their corresponding weights.

For the Llama family of models, we employ a balanced weighting scheme with αr = αc = 1. For the Mistral family, we adopt an asymmetric weighting (αr = 2, αc = 1) to counteract their inherent bias towards safety over utility, thereby achieving a better utility-safety trade-off.

Computing Resources. Our experiments are performed on one Nvidia H800-80GB GPU. The machine is equipped with 192 Intel(R) Xeon(R) Platinum 8468v processors and has a CPU memory of 1584 GB.

4.2 Experiment Results

Performance on Different Evaluation Datasets.

Table 1 shows that MARA significantly improves the alignment performance of various upstream LLMs with different scales and versions in all three challenging datasets. For the Llama family, MARA achieves remarkable improvements, with gains of up to +45.23% on SafeRLHF, +36.86% on BeaverTails, and +62.09% on HarmfulQA. The improvements are particularly pronounced for Llama 3.1-8B, which shows the strongest overall performance gains. For the Mistral family, while the improvements are more modest, MARA still demonstrates consistent positive impacts, with average gains of +10.89%, +15.91%, and +8.81% across the three datasets respectively. On average, MARA yields substantial improvements of +26.28%, +25.88%, and +33.64% across SafeRLHF, BeaverTails, and HarmfulQA respectively, indicating its robust generalization capability across different evaluation scenarios.

Comparison with Different Baselines.

Table 2 presents a comprehensive comparison between MARA and three representative baselines: RLHF, DPO, and Aligner. For the implementation of RLHF and DPO, we utilize the source code from [Zheng et al., 2024]; For the implementation of Aligner, we utilize the aligner-7b-v1.0 alignment model from [Aligners, 2024]. Note that we don't include recent variants of RLHF and DPO (e.g., RLAIF, RLHAIF, TDPO, IPO) as baselines since they mainly improve specific aspects (e.g., feedback sources, overfitting) rather than proposing fundamentally different alignment mechanisms. The experimental results demonstrate that MARA consistently outperforms RLHF and DPO across all three datasets, achieving average improvements of +18.65% and +12.59% respectively.

Compared to Aligner, MARA shows marginally lower performance on certain Llama-based models, likely because Aligner was trained using Llama models as base models. Nevertheless, MARA demonstrates competitive advantages on SafeRLHF (+5.60%) and BeaverTails (+4.37%) evaluation datasets. More importantly, MARA shows significant computational advantages: it requires only a micro three-layer alignment model (4M parameters) compared to Aligner's 7B parameters, achieving higher inference speed (31.41 vs. 20.63 tokens/s).

Compatibility Analysis.

To verify the compatibility of our approach, we conduct comprehensive experiments examining how alignment models trained on one LLM generalize to other inference LLMs. Table 3 presents the results for two alignment models, trained on Llama-3.1-8B and Mistral-7B-v0.3 respectively, when applied to various inference LLMs. The results show that the alignment model trained on Llama-3.1-8B demonstrates strong generalization capability, achieving substantial improvements across all evaluation datasets. Even when applied to different model families, it maintains robust performance, with improvements ranging from +4.02% to +28.14% across different Mistral versions. Similarly, the Mistral-7B-v0.3trained alignment model shows consistent improvements across both model families, with particularly strong performance when applied to Llama 3-8B (SafeRLHF: +18.09%, BeaverTails: +20.71%, HarmfulQA: +37.70%). Overall, our approach demonstrates strong cross-model compatibility with average improvements of +13.82%, +13.91%, and +17.87% across SafeRLHF, BeaverTails, and HarmfulQA respectively.

4.3 Ablation Studies

Ablation on Training Datasets

The experiments introduced before utilized PKU-SafeRLHF datasets for training the alignment model. To verify MARA's effectiveness across different training datasets, we evaluate its performance against SFT, RLHF, DPO and Aligner on two additional training datasets: HH-RLHF [Bai et al., 2022a] and Ultra-Feedback [Cui et al., 2023]. Table 4 presents the results using Llama 3.1-8B as the upstream LLM. The experimental results demonstrate MARA's consistent superiority, with substantial average improvements over all baselines: +23.25% over SFT, +31.75% over RLHF, +18.75% over DPO, and +8.75% over Aligner. Notably, MARA achieves particularly strong performance on HH-RLHF, with improvements of up to +43.00% over RLHF, while maintaining robust gains on Ultra-Feedback across all baseline comparisons.

Ablation on Reward Signal Distribution

To analyze the impact of different reward signal distributions, we conduct experiments with various ratios between reward and cost model signals (αr:αc). As shown in Table 5, we observe distinct trade-offs between helpfulness and harmlessness across different configurations. When solely using the reward model (αr:αc=1:0), the model shows improved helpfulness (up to +37.57%) but decreased harmlessness. Conversely, using only the cost model (αr:αc=0:1) leads to significant gains in harmlessness (up to +65.83%) but often at the expense of helpfulness. A balanced ratio of αr:αc=2:1 achieves the best overall performance across datasets, with preference rate improvements reaching +17.09% on SafeRLHF, +16.00% on BeaverTails, and +11.07% on HarmfulQA using Mistral-7B-v0.3. This suggests that while both reward signals are important, slightly emphasizing the reward model over the cost model leads to optimal balance between helpfulness and safety.

4.4 Visualisation of Accepted Tokens Distribution

To visualize the accept or rejection process for candidate tokens under our approach, and evaluate the impact of sampling parameters on token selection, we present an instance of the experiment performed on PKU-SafeRLHF dataset with Llama-3.2-1B as the upstream LLM.

Figure 3 shows the distribution of accepted tokens across the candidate token set under different sampling parameter configurations. The analysis reveals three key findings:

1. Dominance of first token: Across all parameter settings, the first candidate token demonstrates strong dominance, accounting for 65.9-74.0% of all acceptances. This suggests the model maintains high confidence in its primary predictions with the greatest sampling probability.
2. Parameter sensitivity: Increasing the Top-p value and Top-k value results in a notable decrease in first token acceptance, while simultaneously increasing the acceptance rates of second and third tokens. This indicates that higher sampling thresholds promote more diverse token selection.
3. Optimal decision window: The acceptance distribution demonstrates MARA has learned to concentrate its safetyalignment decisions within a compact token space, where over 95% of all acceptances occur within the top three positions. Such a focused decision mechanism helps maintain generation efficiency while ensuring safety controls.

5 RelatedWorks

Prior approaches to language model alignment can be categorized along two primary dimensions: the utilization of RL and the granularity of alignment.

5.1 Alignment Methodology

RL-based Alignment: RL has demonstrated remarkable effectiveness across various language tasks, including question-answering, machine translation, and text summarization [Kang et al., 2020; Ziegler et al., 2019; Stiennon et al., 2020; Nakano et al., 2021]. Recent years have witnessed its successful application in aligning language models with human values and preferences. State-of-the-art language models, notably InstructGPT [Ouyang et al., 2022] and Llama2 [Touvron et al., 2023], leverage RL-based finetuning. A prominent paradigm in this domain is RLHF, which constructs reward models from human preference data to guide model fine-tuning. Building upon RLHF, RLxF extends this framework to incorporate diverse feedback sources beyond human responses. The variable x in RLxF encompasses AI feedback (RLAIF) [Bai et al., 2022b; Lee et al., 2023] and hybrid human-AI feedback (RLHAIF) [Wu et al., 2021; Saunders et al., 2022; Perez et al., 2023].

Non-RL Alignment: Given the implementation complexity and computational demands of RL-based approaches, alternative methodologies have emerged. RAFT [Dong et al., 2023] and RRHF [Yuan et al., 2023] employ selective fine-tuning on high-quality samples. Rain [Li et al., 2024] and Aligner [Ji et al., 2024a] adopt output rectification strategies, with Rain utilizing an evaluation-rewind mechanism and Aligner implementing a supervised, decoupled alignment model. DPO [Rafailov et al., 2024b] establishes a theoretical mapping between optimal policies and reward functions, eliminating explicit reward modeling. Building upon DPO, recent extensions-IPO [Azar et al., 2024], Token-DPO [Zeng et al., 2024], and DPO-f [Wang et al., 2023a]-address challenges in overfitting, fine-grained alignment, and reverse KL regularization constraints, respectively. Our approach maintains alignment quality comparable to RLbased methods while achieving computational efficiency superior to existing non-RL techniques.

5.2 Alignment Granularity

Sentence-Level Alignment: In most alignment research, the alignment environment is modeled as a bandit environment, treating complete sentences as atomic actions. This paradigm encompasses RLHF [Ouyang et al., 2022] and its variants (RLAIF [Bai et al., 2022b; Lee et al., 2023], RLHAIF [Wu et al., 2021; Saunders et al., 2022; Perez et al., 2023]), as well as DPO-based methods [Rafailov et al., 2024b; Azar et al., 2024; Wang et al., 2023a]. Additional approaches in this category include Rain and Aligner's comprehensive sentence correction mechanisms, and RAFT and RRHF's high-quality sentence filtering strategies.

Token-Level Alignment: Recently, there has been increased interest in token-level alignment, decomposing the task into sequential token generation decisions for enhanced granularity and control. Both RLHF [Zhong et al., 2024] and DPO [Zeng et al., 2024; Rafailov et al., 2024a] have evolved to support token-level alignment through MDP formulations, enabling more precise intervention at each generation step. A concurrent work, IVO [Liu et al., 2025], shares a similar idea with MARA in terms of token-level control. While they focus on improving value estimation, our primary advantage lies in the accuracy of value estimation. MARA employs a lightweight 4M-parameter alignment model, compared to the Llama-3.2-1B model used in [Liu et al., 2025]. This advantage stems from our novel approach of simplifying alignment to binary selection (accept/reject) rather than token sampling.

6 Conclusion

This paper proposes MARA, a micro alignment approach that enhances LLMs' adherence to human preferences through token-level control. Our key innovation lies in introducing a micro alignment model that operates independently from the base language model, making Accept/Reject decisions for candidate tokens to achieve fine-grained alignment. Implemented as a compact three-layer fully-connected network, MARA significantly reduces computational overhead compared to existing SOTA approaches while maintaining superior alignment performance. Extensive experiments across multiple evaluation datasets and LLM architectures demonstrate MARA's effectiveness.

Ethical Statement

Our research on MARA presents a novel approach to language model alignment that significantly reduces computational requirements while maintaining alignment effectiveness. Our experimental validation utilizes publicly available datasets like PKU-SafeRLHF, BeaverTails and HarmfulQA, which are designed to promote helpful, harmless, and honest AI responses. While our efficient alignment approach could theoretically be misused for harmful purposes, we strongly oppose any malicious applications and advocate for the responsible development of alignment techniques.

To promote transparency and reproducibility, we release our complete implementation code under open-source licensing. To maintain anonymity during the review process, we temporarily withhold our trained models and training logs to avoid disclosing author information. Upon acceptance, we will release all materials on the Hugging Face platform. We encourage the AI community to build upon our work while maintaining strict ethical standards, with the goal of making AI alignment more accessible while serving human values and social good.

Acknowledgments

This work was supported by the Innovation and Technology Commission - Mainland-Hong Kong Joint Funding Scheme (Grant No. MHP/038/23).

Contribution Statement

Yang Zhang and Yu Yu contribute to this work equally.