Сообщение Мэтью Харриган. Ознакомьтесь с кодом и сообщите о проблемах на GitHub.
Глубокое обучение произвело революцию в обработке изображений и речи, позволяя превращать края в кошек. В нашей лаборатории мы применяем эти методы для открытия лекарств из малых молекул.
Побочным продуктом революции в глубоком обучении стало создание нескольких высококачественных фреймворков машинного обучения с открытым исходным кодом, которые могут вычислять градиенты произвольных операций. Tensorflow от Google может быть самым известным. Мои исследования были сосредоточены на понимании результатов моделирования большой молекулярной динамики белков и других биомолекул. Вы можете легко представить себе, что предсказание энергий связывания малых молекул является проблемой обучения; можем ли мы использовать некоторые достижения глубокого обучения в области молекулярной динамики для вещей в дополнение к художественным изображениям белков?
Распространенной операцией в биофизике является вычисление сходства двух поз (конформаций) белков с помощью метрики расстояния RMSD. Эта метрика известна своим уважением к трансляционной и вращательной инвариантности. Грубо говоря, он перекрывает две белковые структуры и сообщает среднее расстояние между атомом и его партнером в другой структуре.
RMSD и вращения
Удовлетворение трансляционной симметрии легко: вы просто центрируете свои белки в начале координат, прежде чем проводить сравнение.
# traj = np.array(...) [shape (n_frames, n_atoms, 3)]
traj -= np.mean(traj, axis=1, keepdims=True)Удовлетворить вращательную симметрию сложнее. Вам нужно найти оптимальное вращение между каждой парой конформаций (оптимальный = минимизирует RMSD). Еще в 1976 году Kabsch выяснил, что можно выполнить SVD корреляционной матрицы 3x3 (xyz), чтобы найти оптимальную матрицу вращения.
# x1 = np.array(...) [shape (n_atoms, 3)]
correlation_matrix = np.dot(x1.T, x2)
V, S, W_tr = np.linalg.svd(correlation_matrix)
rotation = np.dot(V, W_tr)
x1_rot = np.dot(x1, rotation)Это не идеально, потому что SVD может дать вам «ротоинверсию», иначе говоря, неправильное вращение, иначе говоря, вращение с последующей инверсией. Мы должны явно проверить и исправить этот случай:
correlation_matrix = np.dot(x1.T, x2)
V, S, W_tr = np.linalg.svd(correlation_matrix)
is_reflection = (np.linalg.det(V) * np.linalg.det(W_tr)) < 0.0
if is_reflection:
V[:, -1] = -V[:, -1]
rotation = np.dot(V, W_tr)
x1_rot = np.dot(x1, rotation)Кватернионы спешат на помощь
В 1987 году компания Horn выяснила, что можно построить ключевую матрицу 4х4 из комбинаций элементов корреляционной матрицы. Он получил эту матрицу из кватернионной математики (хотя ключевой матрицей является нормальная матрица). Главное собственное значение этой матрицы можно использовать для коррекции вращения наивной квадратичной разности между координатами атомов.
correlation_matrix = np.dot(x1.T, x2)
F = key_matrix(correlation_matrix)
vals, vecs = np.linalg.eigh(F)
max_val = vals[-1] # numpy sorts ascending
sd = np.sum(traj1 ** 2 + traj2 ** 2) - 2 * max_val
msd = sd / n_atoms
rmsd = np.sqrt(msd)Важно отметить, что вам не нужно явно создавать матрицу вращения, чтобы найти значение RMSD. Если вам нужно вращение, вы можете восстановить его по ведущему собственному вектору ключевой матрицы.
Tensorflow может это сделать
Мы сформулировали задачу в виде векторных операций и одной самосопряженной задачи на собственные значения. Все эти операции реализованы в Tensorflow!
def key_matrix(r): # @formatter:off return [ [r[0][0] + r[1][1] + r[2][2], r[1][2] - r[2][1], r[2][0] - r[0][2], r[0][1] - r[1][0]], [r[1][2] - r[2][1], r[0][0] - r[1][1] - r[2][2], r[0][1] + r[1][0], r[0][2] + r[2][0]], [r[2][0] - r[0][2], r[0][1] + r[1][0], -r[0][0] + r[1][1] - r[2][2], r[1][2] + r[2][1]], [r[0][1] - r[1][0], r[0][2] + r[2][0], r[1][2] + r[2][1], -r[0][0] - r[1][1] + r[2][2]], ]def squared_deviation(frame, target): R = tf.matmul(frame, target, transpose_a=True) R_parts = [tf.unstack(t) for t in tf.unstack(R)] F_parts = key_matrix(R_parts) F = tf.stack(F_parts, name='F') vals, vecs = tf.self_adjoint_eig(F, name='eig') lmax = tf.unstack(vals)[-1] # tensorflow sorts ascending sd = tf.reduce_sum(frame ** 2 + target ** 2) - 2 * lmax
Преимущество в том, что теперь мы получаем деривативы бесплатно, поэтому можем делать интересные вещи. В качестве игрушечного примера это показывает нахождение «консенсусной» структуры, которая минимизирует среднее RMSD для каждого кадра в траектории молекулярной динамики.
Это обычный код тензорного потока, используемый для выполнения оптимизации.
target = tf.Variable(tf.truncated_normal((1, n_atoms, 3), stddev=0.3), name='target') msd, rot = pairwise_msd(traj, target) loss = tf.reduce_mean(msd, axis=0)optimizer = tf.train.AdamOptimizer(1e-3) train = optimizer.minimize(loss)sess = tf.Session() sess.run(tf.global_variables_initializer()) for step in range(2500): sess.run(train)
Tensorflow на самом деле не очень хорошо с этим справляется
Выполнение собственного разложения для каждой точки данных становится дорогостоящим, особенно потому, что я хочу иметь возможность выполнять попарные (R) вычисления MSD между большой траекторией и значительным количеством целевых структур. MDTraj может очень быстро выполнять огромное количество вычислений RMSD. Он использует лучшую стратегию для поиска начального собственного значения ключевой матрицы 4x4 сверху. Метод Theobald QCP 2005 года явно выписывает характеристический полином для ключевой матрицы. Мы используем тот факт, что существует граница для идентичных структур ((R) MSD = 0), чтобы выбрать начальную точку для итеративного метода Ньютона нахождения ведущего собственного значения. Если мы начнем с этого момента, мы гарантируем, что первый корень характеристического многочлена будет наибольшим собственным значением. Так что давайте запрограммируем это в Tensorflow! Не так быстро (буквально): вы действительно не можете выполнять итерацию в Tensorflow, и кто знает, насколько производительной она была бы, если бы вы могли.
Пользовательская парная операция MSD
Вместо этого я реализовал настраиваемый Tensorflow« op ». Поначалу меня пугала необходимость создавать и отслеживать настраиваемую установку Tensorflow. К счастью, Tensorflow с радостью загрузит разделяемые библиотеки для регистрации Ops во время выполнения. Более того, выпускник Pande Group Имран Хак реализовал быструю (R) реализацию вычисления MSD на C, которую я мог обернуть.
Я реализовал Op, который выполняет попарные вычисления MSD, где двойной цикл распараллеливается с OpenMP. Помимо ускорения в 10 000 раз благодаря собственной реализации метода Horn с тензорным потоком, мы немного быстрее, чем MDTraj, даже несмотря на то, что он использует ту же реализацию под капотом. Для MDTraj цикл по траектории выполняется с помощью OpenMP на C, но итерация по целям должна выполняться на Python с соответствующими накладными расходами.
Я запустил тест, который выполняет попарное вычисление RMSD среди траекторий fs пептида. В частности, от 2800 (шаг = 100) кадров до 28 целей (шаг = 100 * 100).
Реализация ……. Время / мс
TF Native Ops …….… .. 22 843
MDTraj …………… ..…. 33.3
Пользовательские операции TF …… ..…. 0.9
TF Custom Op (w / rot) .. 1.6
А как насчет градиентов?
Причина, по которой мы хотели использовать Tensorflow в первую очередь, заключалась в том, чтобы делать забавные вещи с помощью автоматической дифференциации. Бесплатного обеда нет, и Tensorflow не будет автоматически различать наши пользовательские задания. Coutsias et. др. указал, что производная MSD - это просто разница между координатами в наложенной паре структур. Мы можем это закодировать.
Первая проблема заключается в том, что теперь нам явно нужна матрица вращения, чтобы мы могли использовать ее для вычисления градиентов. Помните, что Теобальд придумал умный метод нахождения ведущего собственного значения, но он дает нам только значение RMSD, а не фактическое вращение (которое требует собственного вектора). К счастью, в 2010 году он расширил метод, чтобы использовать старшее собственное значение, чтобы быстро найти главный собственный вектор.
Я изменил операцию pairwise_msd, чтобы возвращать (n_frames, n_targets) попарную матрицу MSD и матрицы вращения (n_frames, n_targets, 3, 3). Пользователи никогда не должны использовать матрицы вращения для дальнейших вычислений, потому что я не реализовал производные для этого вывода. Вместо этого я использую этот результат в вычислении градиента для MSD. Если кто-то знает, как это сделать лучше, дайте мне знать.
В таблице тестов эта версия Op является вариантом «с гнилью» и работает медленнее (потому что она должна выполнять больше работы).
Слишком много подробностей о коде градиента
Большая часть этого кода просто придает тензорам правильную форму. Нам нужно применить наши матрицы вращения n_frames * n_targets индивидуально к каждой конформации, и нам нужно смешать градиент grad из предыдущей операции в графе вычислений, поэтому мы увеличиваем все до матрицы ранга 4 и явно tile преобразования, которые нужно повернуть, потому что matmul не выполняет широковещательную передачу.
rots = op.outputs[1] N1 = int(confs1.get_shape()[0]) N2 = int(confs2.get_shape()[0])# expand from (n_frames, n_targets) to (n_frame, n_targets, 1, 1) grad = tf.expand_dims(tf.expand_dims(grad, axis=-1), axis=-1)# expand from (n_frames OR n_targets, n_atoms, 3) # to (n_frames OR 1, 1 OR n_targets, n_atoms, 3) expand_confs1 = tf.expand_dims(confs1, axis=1) expand_confs2 = tf.expand_dims(confs2, axis=0)# Explicitly tile conformations for matmul big_confs1 = tf.tile(expand_confs1, [1, N2, 1, 1]) big_confs2 = tf.tile(expand_confs2, [N1, 1, 1, 1])# This is the gradient! dxy = expand_confs1 - tf.matmul(big_confs2, rots, transpose_b=True) dyx = expand_confs2 - tf.matmul(big_confs1, rots, transpose_b=False)
Фактическая форма градиента имеет несколько факторов, которые мы должны включить:
n_atom = float(int(confs1.get_shape()[1]))
dxy = 2 * dxy / n_atom
dyx = 2 * dyx / n_atomНаконец, мы суммируем по оси, которая имеет другие конформации, чтобы убедиться, что наши тензоры градиента совпадают по форме с их переменными.
dr_dc1 = tf.reduce_sum(grad * dxy, axis=1)
dr_dc2 = tf.reduce_sum(grad * dyx, axis=0)Вы забыли о трансляционной симметрии после того, как сосредоточились на вращении? Я так делал изначально! Протестируйте свой код на различных входных данных, включая траектории, которые предварительно не центрированы :). Давайте воспользуемся автоматическим дифференцированием Tensorflow для этой части.
В частности, мы настраиваем операцию «вперед» и вызываем для нее tf.gradients. Мы передаем наши градиенты w.r.t. вращение как аргумент grad_ys. Ура цепное правило!
centered1 = confs1 - tf.reduce_mean(confs1, axis=1, keep_dims=True)
centered2 = confs2 - tf.reduce_mean(confs2, axis=1, keep_dims=True)
dc_dx1 = tf.gradients(centered1, [confs1], grad_ys=dr_dc1)[0]
dc_dx2 = tf.gradients(centered2, [confs2], grad_ys=dr_dc2)[0]Кластеризация RMSD на основе KMeans
В качестве примера того, что мы можем сделать с нашей быстрой парной операцией MSD с градиентами, давайте найдем «оптимальные» центры кластеров (центроиды). Для траектории конформаций найдите центры, которые минимизируют расстояние между каждой точкой и ближайшим к ней центроидом. Чтобы предотвратить обнаружение одного и того же центроида дважды, мы добавляем штраф, чтобы центроиды разделялись. Будьте осторожны, чтобы убедиться, что этот штраф в какой-то момент насыщается, иначе ваша оптимизация просто сделает действительно разные центроиды без учета межкластерных расстояний.
# Out inputs n_clusters = 2 target = tf.Variable(tf.truncated_normal((n_clusters, traj.xyz.shape[1], 3), stddev=0.3))# Set up the compute graph msd, rot = rmsd_op.pairwise_msd(traj.xyz, target) nearest_cluster = msd * tf.nn.softmax(-msd) cluster_dist = tf.reduce_mean(nearest_cluster, axis=(0, 1)) cluster_diff, _ = rmsd_op.pairwise_msd(target, target) cluster_diff = cluster_diff[0, 1] loss = cluster_dist - tf.tanh(cluster_diff*10)# Train it in the normal way optimizer = tf.train.AdamOptimizer(5e-3) train = optimizer.minimize(loss)sess = tf.Session() sess.run(tf.global_variables_initializer()) for step in range(1000): sess.run(train)
Теперь вы можете делать tICA или делать МСМ в этом красивом месте.
Доступность кода
Весь код доступен на Github. Обязательно ознакомьтесь с инструкциями по установке в README, поскольку для пользовательской операции требуется работающий компилятор C ++. Пример консенсуса, пример кластеризации и сценарий профилирования находятся в папке examples и требуют набора данных fs peptide.
Собственная реализация тензорного потока находится в rmsd.py. Код низкого уровня для пользовательской операции находится во вложенной папке rmsd /, а именно в rmsd.cpp. Наконец, rmsd_op.py содержит удобную функцию для загрузки общего объекта, который регистрирует Op. Он также реализует градиенты (на Python).
