Pallatom
全原子蛋白质生成扩散模型技术报告

🎯问题背景：为什么需要全原子生成

蛋白质结构预测与生成领域已经走过了几个阶段：从 AlphaFold 的残基层级精确预测，到 RFdiffusion 基于骨架的扩散设计，再到 Proteus、Chroma 等生成模型。这些方法大多以骨架或粗粒度表示作为生成目标，后续再侧链填充或全原子重建。

但真实蛋白质的功能不仅取决于骨架，还取决于侧链的精确构象、原子间的范德华力、氢键网络、配体结合口袋等。因此，直接生成全原子坐标具有天然优势：

1. 避免两阶段误差累积。先生成骨架再重建侧链，每一阶段都会引入误差；全原子生成一次性建模所有原子。

2. 更自然表达侧链多样性。不同氨基酸侧链构象（rotamer）直接由扩散模型学习，不需要额外的 rotamer 库。

3. 更接近物理能量。原子级坐标可以直接与 MD 力场、对接打分、几何约束对接，便于下游任务。

Pallatom 正是这条路径上的一个代表实现：它把扩散模型直接建模对象从 Cα 骨架推进到所有重原子坐标，并采用token 级 + 原子级双轨表示来兼顾计算效率与表达能力。

🧬Pallatom 整体架构

Pallatom 的核心设计可以用三个关键词概括：双轨表示、EDM 扩散、Haiku 实现。

双轨表示（dual-track）：模型同时维护两套特征。一套是 token 级（residue-level）表示，包括每个残基的单一表示（single）和成对表示（pair）；另一套是原子级（atom-level）表示，包括所有原子的坐标、元素、电荷等。两套表示在模型主干中通过 AtomEncoder 上采样、AtomDecoder 下采样进行信息交换，既能利用残基级别的全局上下文，又能保留原子级别的精细几何。

EDM（Elucidated Diffusion Model）：Pallatom 使用 Karras 等人提出的 EDM 框架作为扩散与采样基础。它用对数正态分布调度噪声水平，通过概率流 ODE（probability flow ODE）进行反向采样，并支持 gamma 噪声增强和自我条件（self-conditioning）机制。

Haiku 实现：模型本身基于 DeepMind 的 Haiku 框架构建（hk.transform + jax.jit）。这并不意外——Pallatom 的代码库大量复用了 AlphaFold 的模块结构，而 AlphaFold 内部原型常用 Haiku。

🔄生成管线：从噪声到 PDB

完整的生成流程从 pallatom.py 的 main() 开始，可划分为 7 个阶段：

def main(data_dir, model_name, savepath, L, cuda_devices, t_min, t_max, gamma, ...):
    os.environ['CUDA_VISIBLE_DEVICES'] = cuda_devices
    os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = "false"
    
    T = 50                              # 扩散步数
    model = PallatomRunner(...)         # 加载 Haiku 模型
    sampler = Sampler(...)              # 初始化 EDM 采样器
    
    for round_idx in range(rounds):
        batch = sampler.prepare_batch(L)         # 准备空 batch + 参考特征
        batch = SampleReference(batch, key1)       # 采样初始噪声
        batch, out_traj = sampler.Sample(batch, key2)  # T 步去噪
        
        # 转换为 atom37 并保存 PDB
        final_atoms = atom14_to_atom37(final_aa, final_atoms)
        save_all_pdb(savepath, final_aa, final_atoms, ...)

1. 环境初始化：设置可见 GPU、关闭 XLA 预分配，保证显存按需分配。

2. 模型加载：PallatomRunner 加载预训练的 Haiku 参数和配置。模型前向被包装为 jax.jit(hk.transform(_forward_fn).apply)，这意味着每次推理都是纯函数调用，自动图编译。

3. 采样器初始化：Sampler 类在 modules/sampling.py 中实现。它根据模型配置中的 r3_edmp 参数创建 EDM 采样器，设置扩散步数 T、噪声参数（sigma_data、sigma_mean、sigma_std）以及是否启用自我条件。

4. Batch 准备：prepare_batch(L) 生成一个空 batch，包含 token 索引、原子掩码、序列位置等占位符，并注入从 SDF 文件解析得到的参考特征（reference features）。

5. 初始噪声采样：SampleReference 在最大噪声水平处采样高斯噪声，得到初始的 all_atom_positions。这是扩散过程的起点。

6. T 步去噪：sampler.Sample(batch, key) 执行 T 步反向扩散，每步调用 SampleIter。输出包含每一步的中间轨迹 out_traj 和最终预测 px0。

7. 后处理与输出：从轨迹中取出最后一帧，用 atom14_to_atom37 将原子 14 表示（每残基最多 14 原子）转换为标准 atom37 表示，然后保存为 PDB 文件。

⚙️单步去噪：SampleIter

单步去噪是 Pallatom 最核心的计算循环。它在 modules/sampling.py:208 的 SampleIter 中实现，每个时间步 t 执行以下流程：

class Sampler:
    def SampleIter(self, batch, key0, t):
        # 1. 当前噪声水平 t_hat
        t_hat = self.r3_sampler.get_noise_level_from_float(t)
        
        # 2. gamma 噪声增强
        xt_noisy, r3_t_hat, r3_sigma_next = self.r3_sampler.add_additional_noise(
            batch, t, key0
        )
        
        # 3. 模型预测干净结构 x0
        out = self.pallatom.model_runner.apply(
            self.pallatom.model_params, key0, batch, is_training=False
        )
        px0 = out['representations']['structure']
        
        # 4. 反向扩散 ODE 步
        batch["all_atom_positions"] = self.r3_sampler.reverse(
            xt_noisy, px0, r3_t_hat, r3_sigma_next, ...
        )
        
        # 5. 自我条件更新
        if self.use_selfcond:
            batch = self.selfcond_manger.update_selcond(out)
        
        return batch, {"px0": px0, ...}

噪声水平调整：get_noise_level_from_float(t) 使用 EDM 的对数正态调度，将归一化时间 t 转换为噪声标准差 sigma。然后 add_additional_noise 应用 gamma 增强：临时把噪声提高到 t_hat = sigma * (1 + gamma)，再注入适量高斯噪声，以改善采样动态。

模型前向：模型接收带噪声的 batch，输出预测的干净结构 x0。注意这是 Haiku 的 apply 调用，模型参数和 RNG 密钥从外部传入，整个调用被 jax.jit 编译。

反向 ODE 步：根据 EDM 的概率流 ODE，由当前噪声样本 xt_noisy 和预测 x0 计算 score，然后更新坐标。公式形式为：

score = (xt_noisy - x0) / t_hat
step = score * step_scale * (sigma_next - t_hat)
new_xt = xt_noisy + step

自我条件（self-conditioning）：Pallatom 支持将上一步预测的结构作为额外条件输入到下一步。这通过 selfcond_manger.update_selcond(out) 更新 batch 中的模板/条件特征实现。self-conditioning 已被证明能提升扩散模型的生成一致性和质量。

🏗️模型前向：EmbeddingsAndDenoise

模型前向传播的核心是 modules.py 中的 EmbeddingsAndDenoise 模块。它把输入 batch 通过编码器、主干块、原子解码器，最终输出预测结构。流程如下：

class EmbeddingsAndDenoise:
    def __call__(self, batch, is_training):
        # 1. 时间嵌入
        ti = TimeFourierEmbedding(self.config)(batch['time_embedding'])
        si += ti
        
        # 2. 原子特征编码：token 级 (a) + 原子级 (ql, cl, plm)
        a, ql, cl, plm, ref_plm, ref_cl = atom_feature_encoder(batch, rl, si, zij, ...)
        
        # 3. 8 个主干块迭代优化
        num_block = 8
        main_stack = layer_stack.layer_stack(num_block, with_state=True)(main_fn)
        (outputs, safe_key), (ql_acts, plm_acts, cl_acts, single_acts, si_acts, zij_acts) = main_stack(...)
        
        # 4. 原子解码器：将 token 特征解码回原子更新
        ql, plm, cl, rl_update, ql_brocast = atom_decoder(rl, ref_ql, ref_plm, cl, ...)
        
        # 5. 坐标更新 + EDM 预处理
        rl = sigma_data**2 / (sigma_data**2 + t_hat**2) * rl_in + \
             sigma_data * t_hat / jnp.sqrt(sigma_data**2 + t_hat**2) * rl_update
        
        # 6. pair 表示更新
        zij, affine_pair, dgram = self.affine2triangle(zij, cb_xyz, masks['pair'], ...)
        
        return representations

时间嵌入：扩散时间步通过 Fourier 特征编码为时间条件信号，注入到 token 和原子表示中，使模型知道当前噪声水平。

Atom Feature Encoder：这是双轨架构的关键。它接收原子坐标 rl 和模板特征，输出两套表示：token 级表示 a（用于残基间注意力），以及原子级表示 ql（query 原子）、cl（context 原子）、plm（pair 级别原子特征）。原子级表示比 token 级表示更精细，能表达侧链和局部几何。

Main Trunk（8 个 block）：每个 block 包含 node 更新（token 级自注意力 + pair bias）、原子更新、pair 更新。使用 layer_stack 将 8 个 block 扫描执行，状态在 block 之间传递。这是典型的 AlphaFold-style 迭代精炼结构，但扩展到原子级。

Atom Decoder：将精炼后的 token 特征重新解码为原子级更新，包括坐标更新 rl_update。解码器采用 AtomAttentionDecoder，将 token 信息广播到每个原子。

EDM 预处理：模型输出的是去噪残差，而不是直接输出坐标。最终坐标通过 EDM 的预处理公式组合输入噪声坐标和预测残差，保证训练与采样时的缩放一致性。

Pair 更新：pair 表示 zij 在每次 block 中通过当前原子坐标计算的几何信息（如距离、角度）更新，再反过来增强 token 注意力。这种结构-特征循环是精确建模原子几何的关键。

📐参考特征工程：从 SDF 到原子坐标

为了让模型知道每种氨基酸的“理想几何”，Pallatom 在 batch 中注入参考特征。这些参考特征来自 RDKit 解析的 SDF 文件，包含 20 种标准氨基酸的 ideal 全原子坐标和元素信息。

def generate_sdf_amino_acid_specific_ref_features(restype3):
    # 读取 SDF 文件
    file_path = restype3_sdf_file_path[restype3]
    supplier = Chem.SDMolSupplier(file_path)
    molecule = supplier[0]
    specific_atom_indices = amino_acid_atom_indices[restype3]
    
    ref_pos, ref_element, ref_charge = [], [], []
    for atom_idx in specific_atom_indices:
        atom = molecule.GetConformer().GetAtomPosition(atom_idx)
        ref_pos.append([atom.x, atom.y, atom.z])
        ref_element.append(atom.GetAtomicNum())
        ref_charge.append(atom.GetFormalCharge())
    
    return {
        'ref_pos': np.array(ref_pos),
        'ref_element': np.array(ref_element),
        'ref_charge': np.array(ref_charge),
    }

在 prepare_batch 中，这些单残基参考特征被扩展到目标长度 L，并通过 dynamic_update_slice 插入到 batch 的对应位置：

def prepare_batch(self, L):
    batch = create_null_batch(token_size=self.fix_size, atom_size=self.atom_fix_size)
    ref_batch = genertate_entangled_ref_features(1, element_dim=128, name_chars_dim=...)
    
    batch["ref_pos"] = dynamic_update_slice(batch["ref_pos"], ref_batch['ref_pos'], ...)
    batch["ref_mask"] = dynamic_update_slice(batch["ref_mask"], ref_batch['ref_mask'], ...)
    batch["ref_element"] = dynamic_update_slice(
        batch["ref_element"], ref_batch['ref_element'], ...
    )

参考特征的作用：

1. 几何先验：为模型提供每种氨基酸的标准键长、键角和理想侧链模板，降低模型学习基本化学结构的压力。

2. 编码器对齐：原子编码器将参考特征与当前坐标对齐，帮助模型在扩散过程中保持局部几何合理性。

3. 元素与电荷信息：元素通过 one-hot 编码（128 维）嵌入，电荷作为标量特征，辅助模型理解原子化学属性。

🧹训练数据整理：筛选、去重与聚类

Pallatom 的训练数据流程在 db_scripts/pipeline.py 中实现，分为四步：质量过滤、高相似去重、多样性聚类、输出组织。

class DatabasePipeline:
    def run(self):
        # Step 1: 结构质量过滤
        step1_df = self.step1_filter_metrics()
        
        # Step 2: TM-score > 0.8 去重
        step2_df = easy_cluster(valid_paths, tm_score=0.8, ...)
        
        # Step 3: TM-score > 0.6 多样性聚类
        step3_df = easy_cluster(valid_paths, tm_score=0.6, ...)
        
        # Step 4: 输出组织
        self.step4_copy_output(step3_df)

Step 1：结构质量过滤：使用 DSSP 计算二级结构，并计算回转半径（gyration）、堆积密度（packing）、SSE 含量等。过滤条件包括长度、回转半径、堆积密度、二级结构含量等，剔除低质量结构。

Step 2：高相似去重（TM-score > 0.8）：调用 Foldseek 的 easy-cluster，对结构进行 TM-score > 0.8 的聚类。每个簇只保留一个代表结构，剔除近重复。

Step 3：多样性聚类（TM-score > 0.6）：在已去重的基础上，再次用 TM-score > 0.6 聚类，为每个训练样本分配 cluster ID，保证训练集覆盖结构空间。

Step 4：输出组织：将最终筛选出的结构复制到输出目录，命名为 AFDB_{cluster}_{count}.pdb，便于后续训练加载。

这套数据流程与 AlphaFold 训练集的整理思路一致，但更强调全原子结构的筛选与去重。TM-score 阈值的选择保证了训练集既不过度冗余，又保留足够多样性。

🚀使用方式：如何运行 Pallatom

从代码结构看，Pallatom 的主要入口是 pallatom.py 的 main() 函数。典型调用参数包括：

python pallatom.py \
    --data_dir /path/to/data \
    --model_name pallatom \
    --savepath /path/to/output \
    --L 128 \                  # 目标序列长度
    --cuda_devices 0 \
    --t_min 0.0 \              # 最小噪声水平
    --t_max 1.0 \              # 最大噪声水平
    --gamma 0.1 \              # gamma 噪声增强
    --step_scale 1.0 \         # ODE 步长缩放
    --T 50 \                   # 扩散步数
    --rounds 10                # 生成样本数

运行流程：

1. 准备模型权重：下载或训练 Pallatom 的 Haiku 检查点，放入 data_dir。

2. 准备 SDF 参考文件：确保 restype3_sdf_file_path 字典指向 20 种氨基酸的 SDF 文件。

3. 执行生成：main 函数会为每个 round 创建一个 batch，采样噪声，执行 T 步去噪，保存 PDB。

4. 评估生成结构：生成的 PDB 可以用传统结构评估工具（如 RMSD、TM-score、GDT、Clash 检测）或下游任务（如对接、MD 稳定性）进一步评估。

对于训练新模型，需要先运行 db_scripts/pipeline.py 整理数据，然后按照模型配置进行训练。模型配置（如 block 数、通道数、噪声参数）在 cfg.model.pallatom 中定义。

🎯基于靶点的蛋白质设计工作流示例

Pallatom 本身是一个无条件的全原子生成模型，但可以通过修改 batch 输入把它嵌入到一条基于靶点（target-based）的蛋白质设计管线中。下面是可复用的工作逻辑与示例代码。

1. 定义靶点与活性位点

假设目标靶点是一个酶，需要设计一个能与其活性口袋结合的小蛋白 binder。首先用靶点结构（PDB）解析出活性位点的几何约束：关键残基的 Cα 坐标、结合口袋的轮廓、保守氢键供体/受体位置。

from Bio.PDB import PDBParser
import numpy as np

parser = PDBParser(QUIET=True)
structure = parser.get_structure('target', 'target.pdb')

# 提取活性位点残基的 Cα 坐标（例如残基 42, 85, 127, 198）
binding_site_residues = [42, 85, 127, 198]
ca_coords = []
for residue in structure[0]['A']:
    if residue.get_id()[1] in binding_site_residues:
        ca_coords.append(residue['CA'].get_coord())

binding_site_center = np.mean(ca_coords, axis=0)
print(f'Binding site center: {binding_site_center}')

2. 生成初始 scaffold

用 Pallatom 生成多条候选骨架。此时可以设计一个比靶点口袋略大的长度，例如 80-120 残基，作为后续 docking 的候选 pool。

import subprocess

N_DESIGNS = 100
L = 100

for i in range(N_DESIGNS):
    subprocess.run([
        'python', 'pallatom.py',
        '--data_dir', 'data/',
        '--model_name', 'pallatom',
        '--savepath', f'output/design_{i:03d}/',
        '--L', str(L),
        '--T', '50',
        '--rounds', '1',
        '--cuda_devices', '0',
    ])

3. 几何约束筛选：口袋匹配

对生成的每个 PDB 计算其表面形状与靶点口袋的互补性。可用 shape complementarity、可接触表面积（SASA）、或简单距离约束初筛。

def pocket_matching_score(pdb_path, binding_site_center, radius=12.0):
    """返回设计结构与靶点口袋中心附近原子的数量。"""
    from Bio.PDB import PDBParser
    import numpy as np
    
    structure = PDBParser(QUIET=True).get_structure('d', pdb_path)
    atoms = [atom.get_coord() for atom in structure.get_atoms()]
    atoms = np.array(atoms)
    
    distances = np.linalg.norm(atoms - binding_site_center, axis=1)
    n_nearby = np.sum(distances < radius)
    return n_nearby

# 筛选出与口袋有显著接触的设计
scores = [(p, pocket_matching_score(p, binding_site_center)) for p in generated_pdbs]
top_candidates = [p for p, s in sorted(scores, key=lambda x: x[1], reverse=True)[:20]]

4. 结构合理性过滤

用 DSSP 或 Foldseek 验证生成的结构是否有合理的二级结构，排除明显坍缩或 clashing 的设计。Pallatom 已经用 SDF 参考特征约束理想几何，但仍需后处理过滤。

from Bio.PDB.DSSP import DSSP

def is_valid_structure(pdb_path):
    try:
        structure = PDBParser(QUIET=True).get_structure('d', pdb_path)
        model = structure[0]
        dssp = DSSP(model, pdb_path)
        # 至少 30% 残基参与二级结构
        sse_ratio = sum(dssp[key][2] in ('H', 'B', 'E', 'G', 'I', 'T')
                        for key in dssp.keys()) / len(dssp.keys())
        return sse_ratio > 0.3
    except Exception:
        return False

valid_candidates = [p for p in top_candidates if is_valid_structure(p)]

5. 对接与亲和力评估

将候选设计与靶点进行对接。可以使用 DiffDock、AutoDock Vina、或 HADDOCK 等工具。如果目标是设计 binder，对接打分是核心选择指标。

# 示例：用 Vina 或 DiffDock 批量对接
for design in valid_candidates:
    subprocess.run([
        'diffdock', '--protein', 'target.pdb',
        '--ligand', design,  # 这里把候选蛋白当作 ligand 处理
        '--out', f'docking_results/{Path(design).stem}/'
    ])

6. 迭代精化：条件生成

更高级的做法是改造 Pallatom 的 batch，使其包含靶点条件信息。例如，可以在 prepare_batch 中注入靶点口袋的参考坐标作为 soft constraint，或在 self-conditioning 时把靶点模板特征传入。这样 Pallatom 从无条件生成变成条件生成，显著减少无效样本。

class TargetConditionedSampler(Sampler):
    def prepare_batch(self, L, target_template=None):
        batch = super().prepare_batch(L)
        if target_template is not None:
            # 将靶点模板坐标作为额外条件注入
            batch['target_template_pos'] = target_template['coords']
            batch['target_template_mask'] = target_template['mask']
        return batch

7. 完整管线总结

基于靶点的 Pallatom 设计流程可以总结为：定义靶点约束 → 批量无条件生成 → 几何初筛 → 结构合理性过滤 → 对接/亲和力评估 → 条件生成精化 → 实验验证。这个流程充分利用了 Pallatom 全原子生成的优势，同时通过下游任务过滤保证设计的可用性。

💡技术启示与复用价值

Pallatom 的代码架构和实现方式中有几个值得研究者借鉴的设计点：

1. 双轨表示的通用性：token 级 + 原子级双轨表示不仅适用于蛋白质，也可以推广到任何具有层级结构的全原子生成问题（如 RNA、DNA、小分子-蛋白复合物）。Encoder-Decoder 的上采样/下采样模式是一种可复用的架构模板。

2. EDM 框架的稳定性：EDM 的对数正态噪声调度、gamma 增强、概率流 ODE 都是已被验证的稳定扩散技术。在三维结构生成中，这些方法比朴素 DDPM 更鲁棒。

3. 自我条件机制：self-conditioning 不仅提升生成质量，还能让模型在采样时利用前一步预测作为额外上下文。这种迭代精化思想可以扩展到其他条件生成任务。

4. SDF 参考特征工程：用 RDKit 解析化学标准文件构建几何先验，是一种低成本但高收益的方法。对于任何需要学习化学几何的生成模型，这都是值得借鉴的模板。

5. Foldseek 数据去重：结构相似性去重比序列相似性去重更能保证训练集的多样性。TM-score 阈值的选择可以根据任务需求调整（如更严格的去重用 0.9，多样性保留用 0.5）。

6. Haiku + JAX 的工程实践：Pallatom 展示了如何在 Haiku 中构建复杂的多层级模型。虽然 Flax 现在是 Google 推荐的新框架，但理解 Haiku 的 transform/apply 模式对于阅读 DeepMind 的 AlphaFold 系列代码仍然重要。

7. atom14/atom37 转换：两种原子表示的转换是结构预测/生成中的常见需求。Pallatom 的转换逻辑可以直接复用或迁移到其他项目。

📄项目信息