医疗诊断模型的训练

题意

给定一个简单的"线性映射 + 线性分类"模型，需要完成三步计算：

前向传播：输入矩阵 $X$ （ $L \times D$ ）经过映射矩阵 $W_{mlp}$ （ $D \times D$ ）得到隐层 $H = X \cdot W_{mlp}$ ，再经过分类矩阵 $W_{cls}$ （ $D \times K$ ）得到打分 $S = H \cdot W_{cls}$ ，对 $L$ 条记录取平均得到预测向量 $\hat{y}$ （ $K$ 维）。
计算 MSE 损失： $\text{MSE} = \frac{1}{K}\sum_{j=1}^{K}(\hat{y}_j - y_j)^2$ 。
SGD 更新：用学习率 $\eta$ 对 $W_{mlp}$ 和 $W_{cls}$ 各做一步梯度下降。

输出预测值、损失值、更新后的两个权重矩阵，保留 2 位小数。

思路

前向传播就是简单的矩阵乘法加平均，关键在于反向传播的梯度推导。

设 $\hat{y} = \frac{1}{L}\sum_{i=1}^{L} S_i$ ，其中 $S = H \cdot W_{cls}$ ， $H = X \cdot W_{mlp}$ 。

MSE 对 $\hat{y}$ 的梯度：

$ $\frac{\partial \mathcal{L}}{\partial \hat{y}_j} = \frac{2(\hat{y}_j - y_j)}{K}$ $

由于 $\hat{y}$ 是对 $S$ 按行取平均，梯度均匀分配到每一行：

$ $\frac{\partial \mathcal{L}}{\partial S_{ij}} = \frac{1}{L} \cdot \frac{\partial \mathcal{L}}{\partial \hat{y}_j}$ $

然后按矩阵乘法的链式法则：

$ $\frac{\partial \mathcal{L}}{\partial W_{cls}} = H^T \cdot \frac{\partial \mathcal{L}}{\partial S}$ $

$ $\frac{\partial \mathcal{L}}{\partial H} = \frac{\partial \mathcal{L}}{\partial S} \cdot W_{cls}^T$ $

$ $\frac{\partial \mathcal{L}}{\partial W_{mlp}} = X^T \cdot \frac{\partial \mathcal{L}}{\partial H}$ $

最后用 SGD 更新： $W \leftarrow W - \eta \cdot \nabla W$ 。

整个过程就是手写一遍神经网络的前向和反向传播，没有激活函数，纯线性运算。

时间复杂度 $O(L \cdot D \cdot (D + K))$ ，空间复杂度 $O(L \cdot D + D \cdot K)$ 。

代码

import sys

def solve():
    data = sys.stdin.read().split('\n')
    line0 = data[0].split(',')
    L, D, K = int(line0[0]), int(line0[1]), int(line0[2])
    eta = float(line0[3])

    y = list(map(float, data[1].split(',')))

    X_flat = list(map(float, data[2].split(',')))
    X = [X_flat[i * D:(i + 1) * D] for i in range(L)]

    W_mlp_flat = list(map(float, data[3].split(',')))
    W_mlp = [W_mlp_flat[i * D:(i + 1) * D] for i in range(D)]

    W_cls_flat = list(map(float, data[4].split(',')))
    W_cls = [W_cls_flat[i * K:(i + 1) * K] for i in range(D)]

    # 前向传播：H = X @ W_mlp (L x D)
    H = [[sum(X[i][k] * W_mlp[k][j] for k in range(D)) for j in range(D)] for i in range(L)]

    # S = H @ W_cls (L x K)
    S = [[sum(H[i][k] * W_cls[k][j] for k in range(D)) for j in range(K)] for i in range(L)]

    # y_pred = mean(S, axis=0)
    y_pred = [sum(S[i][j] for i in range(L)) / L for j in range(K)]

    # MSE
    mse = sum((y_pred[j] - y[j]) ** 2 for j in range(K)) / K

    # 反向传播
    d_ypred = [2.0 * (y_pred[j] - y[j]) / K for j in range(K)]
    d_S = [[d_ypred[j] / L for j in range(K)] for _ in range(L)]

    d_W_cls = [[sum(H[l][i] * d_S[l][j] for l in range(L)) for j in range(K)] for i in range(D)]
    d_H = [[sum(d_S[i][k] * W_cls[j][k] for k in range(K)) for j in range(D)] for i in range(L)]
    d_W_mlp = [[sum(X[l][i] * d_H[l][j] for l in range(L)) for j in range(D)] for i in range(D)]

    # SGD 更新
    for i in range(D):
        for j in range(D):
            W_mlp[i][j] -= eta * d_W_mlp[i][j]
    for i in range(D):
        for j in range(K):
            W_cls[i][j] -= eta * d_W_cls[i][j]

    # 输出
    print(','.join(f'{v:.2f}' for v in y_pred))
    print(f'{mse:.2f}')
    print(','.join(f'{W_mlp[i][j]:.2f}' for i in range(D) for j in range(D)))
    print(','.join(f'{W_cls[i][j]:.2f}' for i in range(D) for j in range(K)))

solve()

复杂度分析

时间复杂度： $O(L \cdot D \cdot (D + K))$ ，主要是矩阵乘法的开销。
空间复杂度： $O(L \cdot D + D \cdot K)$ ，存储中间矩阵 $H$ 、 $S$ 及梯度。