TypeProbe从预训练代码模型隐藏状态恢复类型表示

背景

预训练代码模型（如 CodeBERT、GraphCodeBERT 等）在代码理解与生成任务上表现优异，但其内部如何编码代码的语义信息——尤其是类型信息——仍缺乏深入理解。类型系统是编程语言的核心组成部分，直接影响程序的正确性、可读性和编译检查。尽管已有工作尝试通过可解释性方法分析代码模型的表示，但大多数研究聚焦于代码语法结构或变量名等表面特征，尚未系统性地探讨模型是否学习到了形式化的类型语义，以及这种表示是否具备跨语言泛化能力。

核心内容

该研究（TypeProbe）基于一组 Java 和 Python 代码示例的平行数据集，对预训练代码模型的残差流（residual stream）进行探测（probing），以恢复其内部产生的类型表示。实验设计如下：

• 使用多种预训练代码模型（如 CodeBERT、GraphCodeBERT 等），提取模型在处理代码片段时各层的隐藏状态。
• 构建一个平行数据集，其中相同逻辑的代码以 Java（强类型）和 Python（动态类型）分别实现，确保类型信息在不同语言中对应。
• 训练线性探测分类器（linear probe），从隐藏状态中解码出变量或表达式的类型标签，观察模型是否在内部形成了可区分的类型编码。

主要发现包括：

1. 跨语言类型表示在未类型化代码中涌现：即使模型处理的 Python 代码本身不包含显式类型注解，其隐藏状态中仍能提取出与 Java 类型系统对应的表示空间。这表明预训练阶段接触的多语言代码已使模型抽象出一种与具体语法无关的类型表征。
2. 线性可解码的类型推断：通过在一个语言（如 Java）上训练 probe，可以直接预测另一个语言（如 Python）中函数调用的参数类型和结果类型。这证明隐藏状态中编码了函数应用推导出的类型信息，且该表示是线性可分的。
3. 对词汇和句法变化的鲁棒性：当对代码进行词汇扰动（如重命名变量、替换标识符）或跨语言句法变换（如 Java 的显式类型声明 vs Python 的动态类型）时，探测到的类型表示结构部分保持稳定，说明模型学习的类型特征具有一定的抽象性和鲁棒性。

论文明确指出，据其了解，先前关于代码模型可解释性的工作没有直接关注形式化的类型语义或跨语言类型表示。研究者已公开其代码和数据集。

关键要点

• 使用线性探测方法从预训练代码模型的残差流中恢复类型表示，不需要模型微调。
• 发现强类型语言（Java）的类型表示可以从动态类型语言（Python）的隐藏状态中解码出来，反之亦然，表明跨语言类型表征的存在。
• 类型信息不仅存在于显式类型声明的代码中，也隐含在未类型化代码的表示里。
• 函数应用的结果类型可以通过另一种语言上训练的 probe 线性推断，说明模型内部编码了类型推导的逻辑。
• 类型表示对词汇变化和跨语言句法差异具有部分鲁棒性，暗示其超越表层特征。
• 该工作是首个针对预训练代码模型中形式类型语义和跨语言类型表示的可解释性研究。

意义与影响

这项工作的价值体现在以下几个方面：

• 深化可解释性研究：开辟了一个新视角——从类型系统出发探究代码模型的内在表示，而非仅停留在语法或变量用途。这有助于理解模型如何抽象出编程语言的核心语义要素。
• 指导模型设计与评估：如果预训练代码模型已经捕捉到跨语言的类型信息，那么未来在训练此类模型时可以更明确地融入类型信号，提升其在类型敏感任务（如类型推断、代码补全）上的表现。
• 跨语言泛化的理论基础：发现类型表示在不同语言间共享，为解释代码模型为何能实现跨语言迁移（如从 Java 到 Python）提供了证据，也启发利用类型结构作为跨语言桥梁。
• 对动态类型语言处理的影响：即使 Python 等语言没有显式类型，模型仍隐含了类型概念，这意味着基于静态类型语言的 probe 可以用于分析动态语言模型中的类型知识，拓展了可解释性工具的应用范围。
• 促进类型感知的代码理解工具：基于这项工作，可以构建不依赖编译器或类型推演引擎的轻量级类型恢复工具，用于代码片段理解、漏洞检测等下游任务。

总体而言，TypeProbe 不仅揭示了一种被忽视的内部表示结构，也为未来设计更透明、更可靠的代码 AI 系统提供了方法论和实证基础。