Improving large language models with concept-aware fine-tuning
Authors: name: Michael K. Chen; name: Xikun Zhang; name: Jiaxing Huang; name: Dacheng Tao
Published: 2025-06-09 14:55:00

Abstract: Large language models (LLMs) have become the cornerstone of modern AI.
However, the existing paradigm of next-token prediction fundamentally limits
their ability to form coherent, high-level concepts, making it a critical
barrier to human-like understanding and reasoning. Take the phrase "ribonucleic
acid" as an example: an LLM will first decompose it into tokens, i.e.,
artificial text fragments ("rib", "on", ...), then learn each token
sequentially, rather than grasping the phrase as a unified, coherent semantic
entity. This fragmented representation hinders deeper conceptual understanding
and, ultimately, the development of truly intelligent systems. In response, we
introduce Concept-Aware Fine-Tuning (CAFT), a novel multi-token training method
that redefines how LLMs are fine-tuned. By enabling the learning of sequences
that span multiple tokens, this method fosters stronger concept-aware learning.
Our experiments demonstrate significant improvements compared to conventional
next-token finetuning methods across diverse tasks, including traditional
applications like text summarization and domain-specific ones like de novo
protein design. Multi-token prediction was previously only possible in the
prohibitively expensive pretraining phase; CAFT, to our knowledge, is the first
to bring the multi-token setting to the post-training phase, thus effectively
democratizing its benefits for the broader community of practitioners and
researchers. Finally, the unexpected effectiveness of our proposed method
suggests wider implications for the machine learning research community. All
code and data are available at https://github.com/michaelchen-lab/caft-llm