BAMBOO fidelity check: `xnn` vs the original `bytedance/bamboo`, block by block#

This notebook verifies that the xnn re-implementation of BAMBOO (xnn.hybrid.models.bamboo.BAMBOO, Gong et al. 2024, arXiv:2404.07181) is the same function as the original bytedance/bamboo model, block by block, by transplanting the upstream weights and comparing every intermediate quantity.

The xnn model was written from scratch on the shared xnn abstractions (InteratomicPotential, the xnn.transformer attention/radial primitives, the uniform ForceStressOutput autograd forces); no upstream code is vendored. We only import the upstream package here to compare against it.

The one intentional difference: forces. Upstream reports a charge-non-conservative force nn_forces + coul_forces (charges held fixed) and regularises the charge–position derivative qeq_force toward zero during training (Supplementary Theorem A.2). xnn instead returns the full conservative force -dE/dr uniformly via autograd, which equals upstream forces + qeq_force to machine precision. We check exactly that identity below.

0. Setup: `float64` and the upstream clone#

We run everything in double precision so the comparison is at machine precision, and clone the upstream repository next to this notebook (skipped if already present). Upstream imports torch_runstats.

import logging, warnings
logging.disable(logging.WARNING)
warnings.filterwarnings("ignore")

import os, sys, subprocess
import numpy as np
import torch
torch.set_default_dtype(torch.float64)

UPSTREAM = os.path.abspath("bamboo_upstream")
if not os.path.isdir(UPSTREAM):
    subprocess.run(["git", "clone", "--depth", "1",
                    "https://github.com/bytedance/bamboo", UPSTREAM], check=True)
sys.path.insert(0, UPSTREAM)

import torch.nn as nn
from models.bamboo_get import BambooGET            # the original model
print("upstream bamboo imported from", UPSTREAM)

upstream bamboo imported from /D3/sina/xnn/examples/fidelity_checks/bamboo_upstream

1. Build both models with identical hyper-parameters#

A small-but-real BAMBOO: feature width dim=64, num_rbf=32, 3 GET layers, 16 attention heads, cutoff 5 Å (the upstream defaults).

import xnn
from xnn.common.config import from_dict
from xnn.common.data import structure_to_graph
from xnn.common.models import ForceStressOutput, build_model

DIM, NRBF, NLAYERS, NHEADS, CUT = 64, 32, 3, 16, 5.0

up = BambooGET(
    device=torch.device("cpu"),
    coul_disp_params={"coul_damping_beta": 18.7, "coul_damping_r0": 2.2, "disp_cutoff": 10.0},
    nn_params={"dim": DIM, "num_rbf": NRBF, "rcut": CUT, "charge_ub": 2.0,
               "act_fn": nn.SiLU(), "charge_mlp_layers": 2, "energy_mlp_layers": 2},
    gnn_params={"n_layers": NLAYERS, "num_heads": NHEADS, "act_fn": nn.GELU()},
).eval()

cfg = from_dict({"model": {"name": "bamboo", "cutoff": CUT, "n_features": DIM,
    "n_rbf": NRBF, "n_interactions": NLAYERS, "extra": {"num_heads": NHEADS}}})
xm = build_model(cfg.model).eval()

print("xnn params  :", sum(p.numel() for p in xm.parameters()))
print("upstream params:", sum(p.numel() for p in up.parameters()))

xnn params  : 103620
upstream params: 306750

2. The transplant map (upstream → xnn)#

Every upstream tensor maps to exactly one xnn tensor. Because the read-out MLPs and the transformer projections keep matching shapes, most transplants are a plain load_state_dict.

upstream	xnn
`atom_embtab`	`atom_emb`
`dis_rbf` (`means`, `betas`)	`dis_rbf`
`rbf_proj`	`rbf_proj`
`energy_mlp`	`energy_mlp`
`charge_mlp`	`charge_mlp`
`pred_electronegativity_mlp`	`electronegativity_mlp`
`pred_electronegativity_hardness_mlp`	`hardness_mlp`
`first_attn` / `attns[i]` / `last_attn`: `qkv_proj`, `layer_norm`	`layers[k].attn.qkv_proj`, `.attn.layer_norm`
… `output_proj`, `vec_proj`	`layers[k].output_proj`, `.vec_proj`

def copy(dst, src):
    dst.load_state_dict(src.state_dict())

copy(xm.atom_emb, up.atom_embtab)
copy(xm.dis_rbf, up.dis_rbf)
copy(xm.rbf_proj, up.rbf_proj)
copy(xm.energy_mlp, up.energy_mlp)
copy(xm.charge_mlp, up.charge_mlp)
copy(xm.electronegativity_mlp, up.pred_electronegativity_mlp)
copy(xm.hardness_mlp, up.pred_electronegativity_hardness_mlp)

up_layers = [up.first_attn] + list(up.attns) + [up.last_attn]
for xl, ul in zip(xm.layers, up_layers):
    copy(xl.attn.qkv_proj, ul.qkv_proj)
    copy(xl.attn.layer_norm, ul.layer_norm)
    copy(xl.output_proj, ul.output_proj)
    if xl.vec_proj is not None:
        copy(xl.vec_proj, ul.vec_proj)
print("weights transplanted")

weights transplanted

3. A shared test geometry#

A 7-atom gas-phase cluster with a spread of the electrolyte elements (Li, C, N, O, F, H). We build the xnn graph and, from the same geometry, the upstream input dict. In xnn edge_index = [src, dst] with dst the centre (receiver); upstream calls these row (centre) and col (neighbour), so upstream’s edge_index is the xnn one flipped and its edge_cell_shift is exactly the xnn edge vector pos[dst] - pos[src].

rng = np.random.default_rng(3)
N = 7
pos = rng.uniform(0, 4.5, (N, 3))
Z = [8, 1, 1, 3, 9, 6, 7]
g = structure_to_graph({"pos": pos, "atomic_numbers": Z}, CUT)

posf = g.pos
src, dst = g.edge_index[0], g.edge_index[1]        # src=neighbour, dst=centre
edge_vec = (posf[dst] - posf[src]).detach()

# all ordered intra-cluster pairs for the Coulomb sum
rows, cols = zip(*[(i, j) for i in range(N) for j in range(N) if i != j])
row_all = torch.tensor(rows); col_all = torch.tensor(cols)

up_inputs = {
    "atom_types": g.atomic_numbers.clone(),
    "edge_index": torch.stack([dst, src], 0),         # [centre, neighbour]
    "edge_cell_shift": edge_vec.clone(),
    "all_edge_index": torch.stack([row_all, col_all], 0),
    "all_edge_cell_shift": (posf[row_all] - posf[col_all]).detach(),
    "mol_ids": torch.zeros(N, dtype=torch.long),
    "total_charge": torch.zeros(1),
    "pos": posf.clone(),
}
print("edges:", g.num_edges, " all-pairs:", row_all.numel())

edges: 42  all-pairs: 42

4. Block 1: atom embedding, radial basis, cutoff#

The very first tensors: the type embedding x^0, the exponential-normal radial basis dis_rbf(r), and the cosine cutoff envelope.

def maxabs(a, b): return float((a - b).abs().max())

# reproduce the upstream front-end tensors
coord_diff = up_inputs["edge_cell_shift"]
r = coord_diff.norm(dim=-1)
unit = coord_diff / r.unsqueeze(-1)

up_x0 = up.atom_embtab(up_inputs["atom_types"])
xm_x0 = xm.atom_emb(g.atomic_numbers)
print("atom embedding   :", maxabs(up_x0, xm_x0))
print("radial basis     :", maxabs(up.dis_rbf(r), xm.dis_rbf(r)))
print("cutoff envelope  :", maxabs(up.cutoff(r), xm.cutoff_fn(r)))

atom embedding   : 0.0
radial basis     : 0.0
cutoff envelope  : 0.0

5. Block 2: the edge feature and the equivariant edge vector#

edge_feat = act(W_rbf · rbf(r)) reshaped per head, and the equivariant edge vector e_ij = edge_feat ⊗ r_hat_ij (built from the unit edge direction, as in both codes’ front-ends). We reuse this single edge_feat/edge_vec pair to drive the GET layers of both models below.

H, DPH = NHEADS, DIM // NHEADS
weights_rbf = up.dis_rbf(r)

ef = up.rbf_proj(weights_rbf).reshape(-1, H, DPH)
ef_x = xm.rbf_proj(weights_rbf).reshape(-1, H, DPH)
print("edge feature :", maxabs(ef, ef_x))

edge_feat = ef                                   # identical, reuse for both
edge_vec = edge_feat.unsqueeze(-3) * unit.unsqueeze(-1).unsqueeze(-1)
print("edge vector shape:", tuple(edge_vec.shape))

edge feature : 0.0
edge vector shape: (42, 3, 16, 4)

6. Block 3: every GET layer, one at a time#

Each Graph Equivariant Transformer layer is driven with identical inputs and its scalar (and, except in the last layer, vector) update is compared. This is the heart of the model: multi-head QKV attention on edges, coupled to the scalar and vector channels.

row, col = up_inputs["edge_index"][0], up_inputs["edge_index"][1]  # centre, neighbour
radial = up.cutoff(r)
x0 = up_x0                                        # identical embedding (Block 1)

# --- first layer (no vector input) ---
u_df, u_dv = up.first_attn(x0, edge_feat, edge_vec, row, col, radial, N)
x_df, x_dv = xm.layers[0](x0, edge_feat, edge_vec, None, row, col, radial, N)
print(f"layer 0 (first): d_scalar={maxabs(u_df, x_df):.2e}  d_vector={maxabs(u_dv, x_dv):.2e}")

node_feat_u, node_vec_u = x0 + u_df, u_dv
node_feat_x, node_vec_x = x0 + x_df, x_dv

# --- middle layers ---
for i, (ul, xl) in enumerate(zip(up.attns, xm.layers[1:-1]), start=1):
    u_df, u_dv = ul(node_feat_u, edge_feat, node_vec_u, edge_vec, row, col, radial, N)
    x_df, x_dv = xl(node_feat_x, edge_feat, edge_vec, node_vec_x, row, col, radial, N)
    print(f"layer {i} (mid) : d_scalar={maxabs(u_df, x_df):.2e}  d_vector={maxabs(u_dv, x_dv):.2e}")
    node_feat_u = node_feat_u + u_df; node_vec_u = node_vec_u + u_dv
    node_feat_x = node_feat_x + x_df; node_vec_x = node_vec_x + x_dv

# --- last layer (no vector output) ---
u_df = up.last_attn(node_feat_u, edge_feat, node_vec_u, row, col, radial, N)
x_df, _ = xm.layers[-1](node_feat_x, edge_feat, edge_vec, node_vec_x, row, col, radial, N)
print(f"layer {NLAYERS-1} (last): d_scalar={maxabs(u_df, x_df):.2e}")
print("final node features:", maxabs(node_feat_u + u_df, node_feat_x + x_df))

layer 0 (first): d_scalar=0.00e+00  d_vector=0.00e+00
layer 1 (mid) : d_scalar=0.00e+00  d_vector=0.00e+00
layer 2 (last): d_scalar=0.00e+00
final node features: 0.0

7. Block 4: charges, electronegativity/hardness, and the energy pieces#

The full upstream energy_nn gives the NN energy, the (conserved) partial charges and the electronegativity energy; we compare against the xnn forward outputs. Note xnn spreads the electrostatic energy over node_energy so that energy == sum(node_energy), so we compare the aggregated component energies.

nn_energy_u, charge_u, eneg_u = up.energy_nn({
    **up_inputs, "edge_cell_shift": up_inputs["edge_cell_shift"].clone()})

xout = xm(g)
# xnn aggregates electrostatics into node_energy, so compare component sums:
# upstream (coulomb + electronegativity) == xnn energy_elec.
print("partial charges     :", maxabs(charge_u, xout["charges"]))
print("NN energy           :", abs(float(nn_energy_u[0]) - float(xout["energy_nn"][0])))
up_predict = up.predict({**up_inputs, "edge_cell_shift": up_inputs["edge_cell_shift"].clone()})
elec_up = float(up_predict["energy"][0]) - float(nn_energy_u[0])   # coul + electronegativity
print("electrostatic energy:", abs(elec_up - float(xout["energy_elec"][0])))

partial charges     : 3.608224830031759e-16
NN energy           : 0.0
electrostatic energy: 1.7763568394002505e-14

8. Block 5: the full model (energy, forces, charges, dipole)#

Finally the whole model through the uniform ForceStressOutput. Energy, charges and dipole match directly; the xnn conservative force equals the upstream forces + qeq_force (see the note at the top).

fs = ForceStressOutput(xm, compute_forces=True)
xfull = fs(g)
up_out = up.predict({**up_inputs, "edge_cell_shift": up_inputs["edge_cell_shift"].clone()})

e_rel = abs(float(xfull["energy"][0]) - float(up_out["energy"][0])) / abs(float(up_out["energy"][0]))
full_force = up_out["forces"] + up_out["qeq_force"]
print(f"total energy (rel)        : {e_rel:.2e}")
print(f"partial charges           : {maxabs(up_out['charge'], xfull['charges']):.2e}")
print(f"dipole                    : {maxabs(up_out['dipole'][0], xfull['dipole'][0]):.2e}")
print(f"conservative force        : {maxabs(full_force, xfull['forces']):.2e}")
print(f"  (vs upstream nn+coul only: {maxabs(up_out['forces'], xfull['forces']):.2e} — the qeq residual)")

total energy (rel)        : 6.53e-16
partial charges           : 3.61e-16
dipole                    : 3.55e-15
conservative force        : 1.07e-14
  (vs upstream nn+coul only: 4.69e+00 — the qeq residual)

Summary: every block matches#

block	quantity	agreement
1	atom embedding, radial basis, cutoff	machine precision
2	edge feature, equivariant edge vector	machine precision
3	every GET layer’s scalar & vector update	machine precision
4	partial charges, component energies	machine precision
5	total energy, charges, dipole	machine precision
5	conservative force	`= upstream forces + qeq_force` (machine precision)

The xnn BAMBOO is the same function as bytedance/bamboo, block by block. The only difference is deliberate and documented: xnn returns the full conservative -dE/dr (equal to upstream forces + qeq_force), whereas upstream splits off the charge-equilibrium residual qeq_force as a training penalty.

BAMBOO fidelity check: xnn vs the original bytedance/bamboo, block by block#

0. Setup: float64 and the upstream clone#