Графові нейронні мережі для молекулярної динаміки: навчання силових полів з квантових даних

Молекули як графи

Молекула природно є графом: атоми — це вузли з ознаками (тип елемента, заряд), а зв'язки — або просто сусідні атоми в межах радіуса обрізання — це ребра з геометричними ознаками (відстань, іноді порядок зв'язку). На відміну від веб-графа PageRank, цей граф живе вбудованим у 3D-простір, тому його ребра кодують неперервну геометрію, а не лише дискретну зв'язність:

Саме таке представлення графа робить потенціали GNN переносними за розміром: модель, навчена на 100-атомних молекулах, у принципі може виконувати інференс для системи з 10 000 атомів, оскільки вона завжди оперує лише локальними околами, а не всією системою одразу.

Передача повідомлень: базова операція

Кожен шар message-passing GNN повторює одні й ті самі два кроки: кожен вузол збирає «повідомлення» від сусідів, а потім оновлює свій прихований стан на основі агрегованого повідомлення:

Накладання L таких шарів дозволяє інформації поширюватися на L зв'язкових кроків від кожного атома — достатньо шарів, і кінцеве представлення кожного атома неявно кодує все його локальне хімічне оточення, від безпосередніх зв'язків до дальших стеричних та електростатичних ефектів.

Чому важлива еквіваріантність

Фізика молекули не змінюється, якщо повернути, перенести чи відбити всю систему — енергія інваріантна до цих симетрій, а сили (вектори) мають обертатися разом із молекулою (еквіваріантність). GNN, що ігнорує це, змушена вивчати обертальну інваріантність лише з даних, марнуючи ємність і погано узагальнюючись на невидимі орієнтації.

Передбачення енергозберігаючих сил

Щоб траєкторія молекулярної динаміки була фізично коректною, сили мають бути від'ємним градієнтом єдиної скалярної потенційної енергії — інакше енергія витікатиме або накопичуватиметься нефізично протягом довгих симуляцій. Стандартний трюк: передбачити скалярну повну енергію E, а сили отримати автоматичним диференціюванням, ніколи не передбачаючи векторів сили напряму:

Оскільки E гарантовано є дійсною скалярною функцією позицій, її диференціювання ніколи не може дати силове поле з «завихренням» — модель архітектурно не здатна порушити збереження енергії, чого не можна гарантувати при наївному прямому передбаченні сил.

Шар message passing на JavaScript

function messagePassingLayer(nodeFeatures, edges, mlpMsg, mlpUpdate) {
  const n = nodeFeatures.length;
  const messages = new Array(n).fill(null).map(() => zeros(mlpMsg.outDim));

  // Крок 1: кожне ребро генерує повідомлення, що сумується у вузлі-отримувачі
  for (const { i, j, rij } of edges) {
    const eij = radialBasisExpand(rij);
    const msg = mlpMsg.forward([...nodeFeatures[i], ...nodeFeatures[j], ...eij]);
    messages[i] = addVec(messages[i], msg);
  }

  // Крок 2: об'єднати агреговане повідомлення зі старим станом, щоб отримати новий
  return nodeFeatures.map((h, i) => mlpUpdate.forward([...h, ...messages[i]]));
}

function predictEnergyAndForces(positions, atomicNumbers, layers) {
  let h = atomicNumbers.map(embed);
  const edges = buildNeighborList(positions, 5.0); // радіус обрізання 5 Å
  for (const layer of layers) h = messagePassingLayer(h, edges, layer.msg, layer.update);

  const perAtomEnergy = h.map(hi => layers.readout.forward(hi)[0]);
  const totalEnergy = perAtomEnergy.reduce((a, b) => a + b, 0);
  const forces = autodiffGradient(totalEnergy, positions).map(g => scale(g, -1));
  return { energy: totalEnergy, forces };
}

Навчання на квантових еталонних даних

Ці моделі навчаються з учителем на наборах даних із трійок (геометрія, енергія, сили), обчислених методами DFT або coupled-cluster — таких наборах, як MD17, ANI-1 чи OC20. Функція втрат поєднує помилку енергії та сил, причому силам зазвичай надається значно більша вага, оскільки саме вони безпосередньо керують динамікою:

Застосування поза дизайном ліків і матеріалів

• Розробка ліків: швидкі, точні оцінки енергії зв'язування для віртуального скринінгу молекул-кандидатів.
• Матеріалознавство: передбачення властивостей електролітів батарей, каталізаторів і нових сплавів з майже DFT-точністю.
• Динаміка згортання білків: симуляції на великих часових масштабах, нездійсненні методами ab initio.
• Каталіз: картування шляхів реакцій і перехідних станів на вивчених поверхнях потенційної енергії.
• Хімія клімату та атмосфери: симуляція реактивних процесів у великих ансамблях молекул.