COALA: Convex Optimization for Alignment and Preference Learning on a Single GPU

Aligning large language models with human preferences underpins systems such as ChatGPT and Gemini, but the dominant approaches are heavy. Reinforcement Learning from Human Feedback (RLHF) is computationally expensive and complex, while Direct Preference Optimization (DPO) — though simpler — suffers from inconsistent ranking accuracy, high GPU requirements (it keeps a frozen reference model in memory), and expensive hyperparameter tuning.

We propose COALA (Convex Optimization for Alignment and preference Learning Algorithm): a lightweight, reference-free strategy with strong theoretical guarantees. By leveraging the convex reformulation of two-layer ReLU networks, COALA removes the reference model and obtains a significant reduction in both training time and VRAM, enabling efficient end-to-end training on a single consumer GPU. Across four datasets — including a 26,621-sample synthetic Educational Feedback set — and six models up to Llama-3.1-8B, COALA matches or beats DPO and ORPO while using as little as ~17.6% of DPO's total TFLOPs, and exhibits stable, monotonically increasing rewards. To the best of our knowledge, this is the first time convex optimization has been effectively applied to preference fine-tuning of LLMs.

COALA stacks a convex two-layer ReLU network (cvxNN) on the frozen features of a pre-trained backbone $f_{\theta_\mathrm{pre}}$ and trains it in two convex stages: cvxNN training (Phase I), followed by reference-free convex preference fine-tuning (Phase II).

Phase I — CRONOS: cvxNN training

We solve the convex reformulation of the two-layer ReLU network (Pilanci & Ergen, 2020) with CRONOS, an ADMM solver using preconditioned conjugate gradients (Feng, Frangella & Pilanci, 2024). Once the two subproblems are solved, each ADMM iteration reduces to a vector addition and a single matrix–vector product, which JAX parallelizes efficiently on one GPU. This phase outputs convex-layer weights $(\Theta_1, \theta_2)$ with an $\mathcal{O}(1/K)$ ergodic convergence guarantee.

Phase II — convex preference fine-tuning

We freeze the base model $f_{\theta_\mathrm{pre}}$ and the first-layer weights $\Theta_1$. The reference-free COALA loss is then convex in $\theta_2$:

$$ \min_{\theta_2}\;\; \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}}\, \log\!\Big(1 + \exp\!\big(-\beta\, y_w\, \theta_2^\top (\Theta_1 f_{\theta_\mathrm{pre}}(x))_+ + \gamma\big)\Big). $$

This is a logistic regression in $\theta_2$ — smooth, convex, with a Lipschitz-continuous gradient — so accelerated gradient descent attains an $\mathcal{O}(1/k^2)$ convergence rate to the global optimum; in practice we minimize it with AdamW. Because only the last layer is trained, COALA fits on a single RTX 4090 (24 GB), whereas DPO and ORPO in our experiments require an A100 (40 GB).

For the three models in the figure at the top of this page, each evaluated with and without SFT initialization, COALA empirically shows stable, steadily increasing reward margins, consistent with its convex formulation and convergence guarantees; DPO trajectories are noisier and more hyperparameter-sensitive, frequently plateauing or decaying.

COALA matches or beats DPO and ORPO on AlpacaEval2 length-controlled (LC) win rate across all three backbones and datasets.

Table 1. AlpacaEval2 metrics (mean ± std over 3 runs) by alignment method for three models across three datasets.

In a double-blind study, 107 volunteers — 50 from a graduate deep-learning course and 57 from the technology sector — selected the most helpful answer (EduFeedback) or most positive review (IMDb). COALA achieves the highest win rate on both datasets, with 95% Wald intervals lying entirely above the SFT and ORPO baselines, indicating its gains are statistically significant.

Table 2. Human-evaluation win rates (mean ± 95% CI) from a double-blind study with 107 evaluators across two datasets.

On Llama-3.1-8B, COALA uses roughly 17.6% of DPO's TFLOPs — a 5.7× reduction — while attaining competitive or superior alignment quality.

Table 3. TFLOPs measurements per method on the EduFeedback dataset.

We release EduFeedback: 26,621 GPT-4o-generated student–tutor conversations spanning eleven fields of study such as science and philosophy. Using the novel Alternating Population Strategy we expand these into 65,606 preference triplets without any external reward or re-ranking model: for each prompt, the tutor's immediate response becomes the "chosen" answer and a later, topically related but less direct response from the same conversation becomes the "rejected" answer, exploiting the natural drift from direct answers to tangential follow-ups in objective-domain dialogs.

COALA trained on these pairs achieves a 39.1% win rate in human evaluation, suggesting the strategy captures genuine preference signal at scale. We note this is a domain-specific directness heuristic for objective, succinct dialogs rather than a general preference-labeling method, so only EduFeedback uses it; the three other datasets validate generalization.

This work was supported in part by the NSF CAREER Award (CCF-2236829), the National Institutes of Health (1R01AG08950901A1), the Office of Naval Research (N00014-24-1-2164), and DARPA (HR00112490441). Miria Feng was supported in part by the Stanford Graduate Fellowship. The authors thank Lucy Woof, Zachary Frangella, Kevin Nam, and Zhongwei Dang for valuable feedback and discussions, and all human volunteers at FCS Solutions for participating in the preference-alignment feedback survey.