Reconciling In-Context and In-Weight Learning via Dual Representation Space Encoding
This addresses a fundamental issue in Transformer design for researchers and practitioners, though it is incremental as it builds on prior observations of ICL-IWL conflicts.
The paper tackles the conflict between in-context learning (ICL) and in-weight learning (IWL) in Transformers by proposing a modified architecture, CoQE, that encodes context and samples into separate dual representation spaces. It demonstrates improved ICL performance and reconciles ICL and IWL capabilities in synthetic few-shot classification and a pseudo-arithmetic task.
In-context learning (ICL) is a valuable capability exhibited by Transformers pretrained on diverse sequence tasks. However, previous studies have observed that ICL often conflicts with the model's inherent in-weight learning (IWL) ability. By examining the representation space learned by a toy model in synthetic experiments, we identify the shared encoding space for context and samples in Transformers as a potential source of this conflict. To address this, we modify the model architecture to separately encode the context and samples into two distinct spaces: a task representation space and a sample representation space. We model these two spaces under a simple yet principled framework, assuming a linear representational structure and treating them as a pair of dual spaces. Both theoretical analysis and empirical results demonstrate the effectiveness of our proposed architecture, CoQE, in the single-value answer setting. It not only enhances ICL performance through improved representation learning, but also successfully reconciles ICL and IWL capabilities across synthetic few-shot classification and a newly designed pseudo-arithmetic task. Code: https://github.com/McGuinnessChen/dual-representation-space-encoding