Второй порядок дифференцирования

TensorFlow.js предоставляет мощный инструментарий для работы с автоматическим дифференцированием, включая вычисление градиентов второго порядка, что позволяет строить сложные модели оптимизации и решать задачи, связанные с гессианами, лапласианами и другими производными второго порядка.

Автоматическое дифференцирование

TensorFlow.js опирается на концепцию тензорных вычислений и графа операций. Каждый тензор в библиотеке поддерживает отслеживание операций с помощью механизма tf.GradientTape, который позволяет вычислять производные. Для вычисления производных второго порядка используется вложенное использование GradientTape.

Пример базовой структуры:

import * as tf from '@tensorflow/tfjs';

const x = tf.variable(tf.scalar(2.0));

tf.tidy(() => {
  tf.grads((x) => x.square().mul(x))(x);
});

Здесь используется tf.grads для вычисления градиента функции, а для второго порядка градиентов необходимо создать вложенный GradientTape.

Вычисление градиента второго порядка

Вычисление производной второго порядка требует сохранения операций первого градиента внутри внутреннего контекста GradientTape.

tf.tidy(() => {
  const x = tf.variable(tf.scalar(3.0));

  const secondOrderGrad = tf.tidy(() => {
    return tf.grad((x) => {
      return tf.tidy(() => {
        const y = tf.square(x); // y = x^2
        return y;
      });
    })(x);
  });

  console.log(secondOrderGrad.dataSync());
});

Для вычисления второй производной, лучше использовать вложенные тензорные ленивые вычисления через tf.GradientTape:

tf.tidy(() => {
  const x = tf.variable(tf.scalar(3.0));

  const secondOrderGrad = tf.tidy(() => {
    return tf.grad((x) => {
      return tf.grad((x) => x.square())(x);
    })(x);
  });

  console.log(secondOrderGrad.dataSync()); // ожидаемый результат: 2
});

Ключевой момент: внутренняя градиентная функция возвращает первую производную, а внешняя градиентная функция применяет дифференцирование к ней, формируя производную второго порядка.

Использование `GradientTape` для второго порядка

Более гибкий подход заключается в явном использовании tf.GradientTape. Для второго порядка необходимо создать вложенные контексты с опцией persistent: true:

tf.tidy(() => {
  const x = tf.variable(tf.scalar(4.0));

  tf.tidy(() => {
    const firstOrderGrad = tf.tidy(() => {
      return tf.variable(tf.scalar(0));
    });

    const secondOrderGrad = tf.tidy(() => {
      return tf.variable(tf.scalar(0));
    });

    tf.engine().startScope();

    const tape1 = tf.GradientTape({ persistent: true });
    tape1.watch(x);

    const tape2 = tf.GradientTape();
    tape2.watch(x);

    const y = x.pow(tf.scalar(3)); // y = x^3

    const dy_dx = tape2.gradient(y, x); // первая производная: dy/dx = 3x^2

    const d2y_dx2 = tape1.gradient(dy_dx, x); // вторая производная: d^2y/dx^2 = 6x

    console.log(dy_dx.dataSync());   // [48] при x=4
    console.log(d2y_dx2.dataSync()); // [24] при x=4

    tf.engine().endScope();
  });
});

Особенности:

persistent: true позволяет использовать один GradientTape несколько раз.
Внутренний GradientTape вычисляет первую производную, внешний — вторую.
После вычислений тензоры можно безопасно очищать через tf.tidy() для экономии памяти.

Гессиан и векторные функции

Для многомерных функций (f: ^n ) вычисление матрицы Гессиана требует использования комбинации tf.grads или tf.valueAndGrads:

const f = x => x[0].square().add(x[1].mul(x[1]).mul(tf.scalar(2)));

const x = [tf.variable(tf.scalar(1)), tf.variable(tf.scalar(2))];

const hessian = x.map((xi, i) =>
  x.map((xj, j) =>
    tf.grad(x => tf.grad(f)(x)[i])(x)[j]
  )
);

hessian.forEach(row => row.forEach(t => t.print()));

Выводы по работе с Гессианом:

Гессиан — это матрица второй производной.
Вычисление требует двух уровней дифференцирования.
Для больших размерностей рекомендуется использовать оптимизации памяти и tf.tidy().

Практическое применение второго порядка

Вторые производные необходимы для:

Методов оптимизации второго порядка (например, Newton-Raphson).
Анализа кривизны функции потерь.
Вычисления лапласианов в физических моделях.
Обучения нейронных сетей с продвинутыми метриками регуляризации.