Language models adapt better with matched supervision for structured tasks
Representation-Aligned Auxiliary Supervision for Language Model Adaptation
Computation and LanguageMachine Learning
Summary
Adapting language models to structured tasks like chess can be tricky because they may understand different ways of showing the same information unevenly. The authors studied this issue in chess, where states can be shown in different formats. They found that training with extra guidance in formats compatible with the model helps it learn and generalize better. This extra training also helps in producing fact-based descriptions, not just predicting moves. Their work shows that matching the way information is represented during training is important for improving how language models handle structured problems.
What this means in practice
- •For game ai developers: Improve chess AI by training language models with auxiliary tasks in compatible data formats to predict optimal moves more accurately.
- •For conversational ai builders: Enhance factually grounded dialogue by integrating auxiliary supervision aligned with the model’s preferred data representation.
Authors
Kyuyoung Kim, Peiyao Sheng, Ashwin Hebbar, Peiyang Xu, Yunfei Xie, Kevin Wang, Rui Xin, Chen Wei, Zhangyang Wang, Jinwoo Shin, Pramod Viswanath, Sewoong Oh
Abstract
Language models exhibit strong reasoning capabilities, yet adapting them to structured domains remains challenging and can yield inconsistent outcomes. We identify representation compatibility, the extent to which a model effectively processes a representation for a structured task, as a key factor in adaptation. We study this in chess, which provides a controlled testbed with precise semantics, computable optimal actions, and multiple state representations, including a symbolic encoding (FEN) and a spatial format (ASCII). We find that models often process semantically equivalent inputs substantially differently, affecting both learning and generalization. Building on this observation, we propose representation-aligned auxiliary supervision, which uses environment-derived tasks expressed in compatible representations to improve adaptation to structured domains. Across models and representations, auxiliary supervision consistently improves optimal-move prediction relative to target-only training under identical target data. Tasks that expose environment dynamics provide larger and most consistent gains than surface-level or static supervision, while remaining competitive with substantially increasing the amount of target-task data. Moreover, ASCII-trained models transfer more effectively to FEN than FEN-trained models do to ASCII, even surpassing the FEN target-only baseline on FEN evaluation. The gains also extend beyond optimal-move prediction to open-ended, factually grounded commentary generation. Overall, our results show that auxiliary supervision in model-compatible representations can enable effective adaptation in structured domains.