Technical Summary: TriAlign: Towards Universal Truth Consistency in Personalized LLM Alignment

Problem Statement

Personalized Large Language Models (PLLMs) adapt responses to user preferences and social attributes (e.g., age, gender, occupation). However, this personalization introduces a critical flaw: Universal Truth Inconsistency. While the presentation of a response may vary, the underlying objective truth (e.g., mathematical facts, scientific principles) should remain invariant across social groups. Existing evidence suggests that PLLMs can systematically provide less accurate responses to specific social groups compared to others, creating "truth gaps."

Current alignment methods fail to address this adequately:

1. "One-size-fits-all" models ignore personalization entirely.
2. Existing PLLM alignment focuses primarily on subjective preferences, overlooking truth invariance.
3. Prompt-based mitigation (e.g., instructing models to ignore demographics) only steers surface-level generation without modifying the demographic correlations encoded in model parameters, often yielding marginal improvements or degrading personalization quality.

There is a need for a training-based alignment method that jointly optimizes universal truth accuracy, cross-group truth consistency, and personalization.

Methodology: TriAlign

The authors propose TriAlign, the first offline Multi-Agent Reinforcement Learning (MARL) framework designed for Truth-Invariant Alignment (TIA).

1. Formulation as MARL

TriAlign models each social group as an independent agent interacting within a shared environment.

• Agents: Each agent $i$ corresponds to a social group $g_i$ (defined by demographic attributes like gender or age).
• State: $s_{i,t} = (x, g_i, \tau_{i,t})$, where $x$ is the query, $g_i$ is the group profile, and $\tau$ is the conversation history.
• Action: The agent generates a response $a_{i,t}$.
• Environment: Collects joint actions from all $m$ agents for the same query $x$, compares them to identify truth inconsistencies, and provides feedback.
• Training/Execution: Training is centralized (agents interact to compare responses), while execution is decentralized (each agent generates a response independently based on its specific group profile).

2. Objective and Reward Design

TriAlign optimizes a fairness-aware objective inspired by Nash Social Welfare (NSW) to prevent the model from favoring dominant groups. The objective function $J(\theta)$ maximizes the sum of returns while incorporating a logarithmic term for the worst-performing groups:
$$J(\theta) = \mathbb{E} \left[ \sum_{i=1}^m (R_i + \lambda_{NSW} \log(R_i + \epsilon)) \right]$$
where $R_i$ is the cumulative return for agent $i$.

The Reward Function $r_i$ for each action combines three components:

1. Truth Accuracy ($r_{truth}$): Correctness against the universal truth $y^\star$.
2. Personalization ($r_{pref}$): Alignment with group-specific constraints (style, tone).
3. Consistency Penalty ($r_{cons}$): An explicit penalty based on the absolute difference between the agent's truth accuracy and the average truth accuracy across all social groups for the same query. This directly discourages deviations in universal truth.

3. Offline RL Implementation

To avoid the high cost of online RL (which requires human annotations for fairness and truth), TriAlign uses offline RL on pre-collected trajectories:

• Data Construction: A multi-turn interaction simulation is run where $m$ agents (representing different social groups) answer the same questions. The environment compares answers and assigns rewards.
• Optimization: The framework uses a weighted Supervised Fine-Tuning (SFT) objective with implicit KL constraints. The weights are derived from the advantage function and the NSW objective, prioritizing actions that improve "worst-case" group performance while maintaining high collective accuracy.

4. Persona Augmentation

To ensure diverse training data, the authors constructed a social-group space $G$ with 75 demographic category-value pairs. They utilized a persona-to-persona augmentation pipeline to generate a balanced pool of 7,500 personas, ensuring underrepresented groups were adequately simulated.

Key Contributions

1. Problem Formulation: The work is the first to formulate Truth-Invariant Alignment (TIA) as an offline MARL problem, explicitly addressing the trade-off between personalization and universal truth consistency across social groups.
2. TriAlign Framework: Proposes a novel MARL framework where social groups act as agents. It introduces a fairness-aware objective (NSW-inspired) and an explicit cross-group inconsistency penalty to jointly optimize truth, consistency, and personalization.
3. Interaction Simulation: Develops a multi-social-group, multi-turn interaction framework to construct offline RL trajectories, allowing agents to iteratively refine responses through comparative feedback.
4. Empirical Validation: Demonstrates through extensive experiments that TriAlign outperforms strong baselines (including prompt-based debiasing and standard SFT) in reducing truth disparities while maintaining or improving personalization quality.

Experimental Results

Experiments were conducted on diverse benchmarks including AIME25 (math), StereoSet (social bias), MMLU-PRO, and TruthfulQA.

• Truth Consistency: TriAlign significantly reduced the Standard Deviation (Std), Worst-Group Gap, and Coefficient of Variation (CV) of truth accuracy across social groups compared to baselines. For instance, on StereoSet, TriAlign reduced the gap between the best and worst-performing groups to 0.015, compared to 0.100 for the base model.
• Truth Accuracy: The method improved the Mean universal truth accuracy (e.g., 0.788 on StereoSet vs. 0.317 for the base model) and the Worst-Group accuracy, indicating that marginalized groups received significantly more accurate responses.
• Personalization: Unlike prompt-based methods that often sacrifice personalization for fairness, TriAlign maintained or improved Preference Alignment Accuracy (Pref) (e.g., 0.791 on StereoSet).
• Ablation Studies: Removing the NSW objective or the consistency modeling component led to degraded performance in both fairness metrics and personalization, confirming the necessity of both components.

Significance and Claims

The paper claims that TriAlign represents a foundational step toward fairer and more reliable personalized AI systems. Its significance lies in:

• Moving beyond "surface-level" debiasing (prompting) to training-based alignment that modifies the model's internal behavior regarding truth and fairness.
• Demonstrating that universal truth consistency and personalization are not mutually exclusive; they can be jointly optimized using a multi-agent perspective.
• Providing a scalable, offline RL solution that does not rely on expensive human-in-the-loop annotations for fairness, making it applicable to large-scale deployment.

The authors modestly acknowledge limitations, noting that their simulated personas may not fully capture the complexity of real-world intersectional identities and that balancing these objectives remains a challenging, ongoing research problem. They emphasize that TriAlign is a step toward, not a complete solution for, universal truth consistency in PLLMs.