Embedding-Aligned Language Models

Abstract

We propose a novel approach for training large language models (LLMs) to adhere to objectives defined within a latent embedding space. Our method leverages reinforcement learning (RL), treating a pre-trained LLM as an environment. Our embedding-aligned guided language (EAGLE) agent is trained to iteratively steer the LLM's generation towards optimal regions of the latent embedding space, w.r.t. some predefined criterion. We demonstrate the effectiveness of the EAGLE agent using the MovieLens 25M and Amazon Review datasets to surface content gaps that satisfy latent user demand. We also demonstrate the benefit of using an optimal design of a state-dependent action set to improve EAGLE's efficiency. Our work paves the way for controlled and grounded text generation using LLMs, ensuring consistency with domain-specific knowledge and data representations.

Figures (2)

• Figure 1: An illustration comparing the creation of descriptions of novel entities using ELM tennenholtz2024demystifying vs. EAGLE (ours). The latent embedding space ${\mathcal{Z}}$ is illustrated as a complex surface on the right. Red points on the surface are illustrated as latent embeddings of existing entities. Black points are used to illustrate hypothetical (i.e., non-existing) entities. In ELM, a utility is maximized in latent embedding space to identify an optimal point. This hypothetical embedding is then decoded back to ambient space ${\mathcal{X}}$ to form a description of the hypothetical entity. Conversely, the EAGLE agent utilizes a highly capable pre-trained LLM as an environment to search for novel entities in ambient space ${\mathcal{X}}$. EAGLE does not use a decoder, but rather only requires an encoder $E_D: {\mathcal{X}} \mapsto {\mathcal{Z}}$. More specifically, the EAGLE agent uses action prompts to change an existing entity using the environment LLM. Every new changed entity is then mapped back to the latent embedding space using the encoder $E_D$. The utility function is optimized by the EAGLE agent through a reward signal to stir the changes to areas of high utility in the latent embedding space. A final description is then returned after $H$ steps.
• Figure 2: An illustration of different forms of coverage and bias of an action space. As actions are textual prompts, their corresponding embedding directions may either provide good coverage, or partial coverage of the underlying embedding manifold (e.g., there may be directions that are not covered by the generated actions). Moreover, actions may be uniformly biased toward specific directions. We accommodate for this using a G-optimal design which reduces this bias through a set of exploratory actions in embedding space.

Theorems & Definitions (1)

• Definition 1: G-optimal design