深度学习基础回顾：神经网络、激活函数、损失函数、优化器

说实话，很多同学一上来就追着分布式训练跑，结果连单卡模型都调不稳。我记得当年做分子生成项目时，模型死活不收敛，折腾了两天，最后发现是激活函数选错了——这玩意儿看着不起眼，坑起来真要命。

所以这一章，咱们把老本行踏实捋一遍。重点不是堆公式，是我踩过的雷和实战中用得顺手的招式。

神经网络：从单层到深度

神经网络本质就是个多层复合函数。输入是分子描述符、原子坐标、键长这些特征，经过一层层线性变换+非线性激活，最后输出分子的属性或生成新的分子。

一个全连接层的基本操作：

z = W * x + b
a = activation(z)

我自己的习惯是，写代码前先把维度算清楚，画个表格贴显示器上。尤其是分布式场景里，各卡的输入张量切分方式一乱，整个loss直接飞掉。

层类型	输入形状	输出形状	参数量
全连接	(batch, 128)	(batch, 256)	128×256 + 256
卷积1D	(batch, 50, 32)	(batch, 48, 64)	kernel_size×32×64 + 64
图卷积	节点特征 (N, F)	(N, F')	邻接矩阵相关

你想想看，如果batch size是1024，张量在GPU间切分时，每个rank只拿到256条样本，但是全连接层的权重矩阵在所有卡上是一样的。这就是分布式训练的基础——数据并行。我个人建议先用单卡把小模型跑通，再拆成多卡，否则问题定位会非常痛苦。

激活函数：非线性是关键

没有激活函数，多少层都等于一层线性变换。分子生成里常用的就那几个，但选错一个，梯度就炸了。

• ReLU：最简单，梯度1或0。缺点是死神经元——我一台卡上曾经80%的ReLU单元全死了，loss死活降不下去，排查半天才发现初始化有问题。
• Leaky ReLU / PReLU：给负数一点小斜率，避免死亡。我个人习惯在分子编码器里用Leaky ReLU，斜率0.01就够。
• GELU：Transformer和扩散模型标配。分子生成里的注意力模块，我推荐用GELU，效果确实比ReLU稳一截。
• Sigmoid / Tanh：输出范围有限，常用于门控或分子二分类任务。但小心梯度饱和——我曾在一个深度生成网络上用了Tanh，结果深层梯度几乎为零，白训三天。

# 我在分子生成项目中常用的激活函数配置
import torch.nn as nn

class MolEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(256, 512)
        self.act = nn.GELU()
        self.fc2 = nn.Linear(512, 256)

    def forward(self, x):
        x = self.act(self.fc1(x))
        x = self.fc2(x)
        return x

损失函数：分子生成场景怎么选

损失函数直接决定了模型学什么方向。分子生成不是单纯的分类或回归，往往要组合多个loss。

任务类型	常用损失	原因
分子性质预测	MSE / MAE	数值回归，简单直接
分子图生成	交叉熵 + 重建损失	离散原子类型、键型
分子构象生成	扩散噪声损失 (L2 on epsilon)	符合高斯扩散假设
分子对接评分	对比损失 (InfoNCE)	区分正负分子对

我记得一个很典型的坑：做分子变分自编码器时，KL散度项权重调太大，模型直接坍缩到先验分布，生成出来的分子全是甲烷。你想想看，调了两周，结果每张卡都给同一坨东西。后来我把KL权重从1.0降到0.001，配合beta-VAE的渐进式增加，才算稳住。

优化器：SGD、Adam、LAMB怎么选

优化器这东西，说白了就是下山的路怎么走。单卡时我无脑AdamW，但多卡分布式下，AdamW的通信量其实不小——它要同步一阶动量和二阶动量。

• SGD + Momentum：经典，但收敛慢。现在大规模分子预训练很少用了。
• Adam / AdamW：自适应学习率，效果好，我90%的项目用这个。不过注意，weight decay在AdamW里要同时作用在梯度上，别用成Adam。
• LAMB：专门为大规模分布式训练设计的，能支持几万的batch size。我用它训过分子构象模型，单卡batch 32，32卡batch 1024，收敛速度和单卡几乎一致。夸张吧？
• Adafactor：内存省一半，适合分子大模型。如果你的显存快满了，可以试试。

# 我常用的分布式优化器配置示例
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2)
scheduler = CosineAnnealingLR(optimizer, T_max=100)

# 分布式场景下，最好用DistributedSampler + 按global_batch调整lr
lr = base_lr * (global_batch_size / 256)

本章知识图谱

小结

深度学习基础这部分，说白了就是搭积木。神经网络是骨架，激活函数给非线性，损失函数定方向，优化器让模型真正动起来。我这两年带项目发现，70%的训练崩溃都出在这四个基础模块的搭配上——选对了就顺，选错了就卡。

下一阶段咱们就要把这些基础模块扔到多卡环境里跑，到时候我会给你讲分布式数据并行、梯度同步、batch size和lr的联动策略。嗯，先把这章消化了，后面上多卡才不慌。