Torchmetrics中R²分数计算的总平方和公式解析

在机器学习模型的回归任务评估中，R²分数（决定系数）是一个常用的性能指标。作为Lightning-AI生态中的重要组件，Torchmetrics库实现了这一指标的快速计算。本文将深入分析其实现中总平方和（TSS）的计算公式选择。

总平方和的数学本质

总平方和传统定义为观测值与其均值之差的平方和：

TSS = Σ(y_i - ȳ)²

其中ȳ表示观测值的算术平均数。这个公式直观反映了数据围绕均值的离散程度。

Torchmetrics的实现方式

Torchmetrics采用了等价的代数变形公式：

TSS = Σy_i² - (Σy_i)ȳ

这种形式通过三个基本统计量实现计算：

1. 观测值平方和（Σy_i²）
2. 观测值总和（Σy_i）
3. 样本数量（n）

公式等价性证明

通过展开传统定义中的平方项：

TSS = Σ(y_i - ȳ)² 
    = Σ(y_i² - 2y_iȳ + ȳ²)
    = Σy_i² - 2ȳΣy_i + nȳ²

代入ȳ=Σy_i/n的关系后：

TSS = Σy_i² - 2(Σy_i)²/n + (Σy_i)²/n
    = Σy_i² - (Σy_i)²/n
    = Σy_i² - (Σy_i)ȳ

实现选择的优势

这种计算方式具有两个显著优点：

1. 数值稳定性：避免直接计算大量小数的平方和
2. 计算效率：只需遍历数据一次即可获得所需统计量

实际应用启示

理解这种实现方式有助于：

• 在自定义指标时选择更优的计算公式
• 处理大规模数据时保证数值精度
• 深入理解统计量与模型评估指标的关系

Torchmetrics的这种实现体现了机器学习库设计中兼顾数学正确性与计算效率的典型思路，值得开发者学习借鉴。