Concept Tokens: Learning Behavioral Embeddings Through Concept Definitions

TL;DR

Concept Tokens add new concepts to frozen LLMs by learning a single token's embedding from multiple natural-language definitions, enabling concept-specific control without updating model weights. Across hallucination control, recasting, and qualitative tower studies, the learned embedding shows directional effects: negation reduces undesired behavior while assertion amplifies it, though precision trade-offs and abstentions arise. The Eiffel Tower experiment demonstrates effective activation of known representations, whereas the fictional Austral Tower reveals limits in storing new factual details, suggesting the token acts as a semantic attractor rather than a factual memory. Overall, Concept Tokens offer a lightweight, on-device mechanism to steer behavior via definitional supervision, with future work exploring design choices, multiple tokens, and mechanistic interpretability.

Abstract

We propose Concept Tokens, a lightweight method that adds a new special token to a pretrained LLM and learns only its embedding from multiple natural language definitions of a target concept, where occurrences of the concept are replaced by the new token. The LLM is kept frozen and the embedding is optimized with the standard language-modeling objective. We evaluate Concept Tokens in three settings. First, we study hallucinations in closed-book question answering on HotpotQA and find a directional effect: negating the hallucination token reduces hallucinated answers mainly by increasing abstentions, whereas asserting it increases hallucinations and lowers precision. Second, we induce recasting, a pedagogical feedback strategy for second language teaching, and observe the same directional effect. Moreover, compared to providing the full definitional corpus in-context, concept tokens better preserve compliance with other instructions (e.g., asking follow-up questions). Finally, we include a qualitative study with the Eiffel Tower and a fictional "Austral Tower" to illustrate what information the learned embeddings capture and where their limitations emerge. Overall, Concept Tokens provide a compact control signal learned from definitions that can steer behavior in frozen LLMs.

Figures (5)

• Figure 1: Overview of concept tokens. (1) We build a definitional corpus for a concept and instantiate it by replacing each mention with a new token $t_c$. (2) Keeping the LLM frozen, we optimize only the concept embedding $e_c$ by minimizing cross-entropy loss on the instantiated corpus. (3) The learned embedding captures a generalized representation that best satisfies the (potentially competing) constraints imposed by multiple definitions.
• Figure 2: Category proportions (Correct, Hallucination, No Answer) for three conditions: concept token negated, no instruction, and concept token asserted.
• Figure 3: Qualitative examples comparing concept tokens (CT) and the definitional-corpus in-context baseline (DIC) using questions from HotpotQA.
• Figure 4: Category proportions (Recasting or Explicit correction) for three conditions: concept token negated, no instruction, and concept token asserted.
• Figure 5: Qualitative examples comparing concept tokens (CT) and the definitional-corpus in-context baseline (DIC).