Reproduce “Aligning Transformers with Continuous Feedback via Energy Rank Alignment” on OpenResearch

About this paper

Shriram Chennakesavalu, Frank Hu, Sebastian Ibarraran, Grant M. Rotskoff

Researchers at Stanford University introduced Energy Rank Alignment (ERA), a gradient-based algorithm that directly aligns autoregressive transformers to explicit, continuous reward functions for tasks in molecular and protein design. ERA demonstrates superior sample efficiency and diversity compared to existing methods in molecular docking, while also robustly guiding protein language models for directed evolution.

Key insights

• ERA's theoretical foundation demonstrates that its minimizer asymptotically converges to an ideal Gibbs-Boltzmann distribution, where the reward function acts as an energy, offering interpretable control over generated output entropy and diversity.
• The algorithm provides independent, tunable control over the reward's influence (via inverse temperature "beta") and regularization towards the reference policy (via "gamma"), allowing fine-grained balancing of exploitation and exploration.
• Unlike DPO, ERA's direct utilization of explicit reward functions contributes to greater stability and diversity of generated samples, mitigating issues like mode collapse particularly in scenarios with continuous feedback.

Summary by alphaXiv.