def gradients(func, x, name='t3f_gradients', runtime_check=True):
"""Riemannian autodiff: returns gradient projected on tangent space of TT.
Computes projection of the gradient df/dx onto the tangent space of TT tensor
at point x.
Warning: this is experimental feature and it may not work for some function,
e.g. ones that include QR or SVD decomposition (t3f.project, t3f.round) or
for functions that work with TT-cores directly (in contrast to working with
TT-object only via t3f functions). In this cases this function can silently
return wrong results!
Example:
# Scalar product with some predefined tensor squared 0.5 * <x, t>**2.
# It's gradient is <x, t> t and it's Riemannian gradient is
# t3f.project(<x, t> * t, x)
f = lambda x: 0.5 * t3f.flat_inner(x, t)**2
projected_grad = t3f.gradients(f, x) # t3f.project(t3f.flat_inner(x, t) * t, x)
Args:
func: function that takes TensorTrain object as input and outputs a number.
x: point at which to compute the gradient and on which tangent space to
project the gradient.
name: string, name of the Op.
runtime_check: [True] whether to do a sanity check that the passed
function is invariant to different TT representations (otherwise
the Rieamnnian gradient doesn't even exist). It makes things slower,
but helps catching bugs, so turn it off during production deployment.
Returns:
`TensorTrain`, projection of the gradient df/dx onto the tangent space at
point x.
See also:
t3f.hessian_vector_product
"""
with tf.name_scope(name, values=x.tt_cores):
left = decompositions.orthogonalize_tt_cores(x)
right = decompositions.orthogonalize_tt_cores(left,
left_to_right=False)
deltas = [right.tt_cores[0]]
deltas += [tf.zeros_like(cc) for cc in right.tt_cores[1:]]
x_projection = riemannian.deltas_to_tangent_space(
deltas, x, left, right)
function_value = func(x_projection)
if runtime_check:
assert_op = _is_invariant_to_input_transforms(
function_value, func(x))
else:
assert_op = tf.no_op()
with tf.control_dependencies([assert_op]):
cores_grad = tf.gradients(function_value, deltas)
deltas = _enforce_gauge_conditions(cores_grad, left)
return riemannian.deltas_to_tangent_space(deltas, x, left, right)
def testToAndFromDeltas(self):
# Test converting to and from deltas representation of the tangent space
# element.
what = initializers.random_tensor((2, 3, 4), 4, dtype=self.dtype)
where = initializers.random_tensor((2, 3, 4), 3, dtype=self.dtype)
projected = riemannian.project(what, where)
deltas = riemannian.tangent_space_to_deltas(projected)
reconstructed_projected = riemannian.deltas_to_tangent_space(deltas, where)
# Tangent space element norm can be computed from deltas norm.
projected_normsq_desired = ops.frobenius_norm_squared(projected)
projected_normsq_actual = tf.add_n([tf.reduce_sum(c * c) for c in deltas])
desired_val, actual_val = self.evaluate((ops.full(projected),
ops.full(reconstructed_projected)))
self.assertAllClose(desired_val, actual_val)
desired_val, actual_val = self.evaluate((projected_normsq_desired,
projected_normsq_actual))
self.assertAllClose(desired_val, actual_val)
def testToAndFromDeltasBatch(self):
# Test converting to and from deltas representation of the tangent space
# element in the batch case.
what = initializers.random_matrix_batch(((2, 3, 4), (3, 3, 3)), 4,
batch_size=3, dtype=self.dtype)
where = initializers.random_matrix(((2, 3, 4), (3, 3, 3)), 3,
dtype=self.dtype)
projected = riemannian.project(what, where)
deltas = riemannian.tangent_space_to_deltas(projected)
reconstructed_projected = riemannian.deltas_to_tangent_space(deltas, where)
# Tangent space element norm can be computed from deltas norm.
projected_normsq_desired = ops.frobenius_norm_squared(projected)
d_normssq = [tf.reduce_sum(tf.reshape(c, (3, -1)) ** 2, 1) for c in deltas]
projected_normsq_actual = tf.add_n(d_normssq)
desired_val, actual_val = self.evaluate((ops.full(projected),
ops.full(reconstructed_projected)))
self.assertAllClose(desired_val, actual_val)
desired_val, actual_val = self.evaluate((projected_normsq_desired,
projected_normsq_actual))
self.assertAllClose(desired_val, actual_val)
def hessian_vector_product(func,
x,
vector,
name='t3f_hessian_vector_product',
runtime_check=True):
"""P_x [d^2f/dx^2] P_x vector, i.e. Riemannian hessian by vector product.
Computes
P_x [d^2f/dx^2] P_x vector
where P_x is projection onto the tangent space of TT at point x and
d^2f/dx^2 is the Hessian of the function.
Note that the true Riemannian hessian also includes the manifold curvature
term which is ignored here.
Warning: this is experimental feature and it may not work for some function,
e.g. ones that include QR or SVD decomposition (t3f.project, t3f.round) or
for functions that work with TT-cores directly (in contrast to working with
TT-object only via t3f functions). In this cases this function can silently
return wrong results!
Example:
# Quadratic form with matrix A: <x, A x>.
# It's gradient is (A + A.T) x, it's Hessian is (A + A.T)
# It's Riemannian Hessian by vector product is
# proj_vec = t3f.project(vector, x)
# t3f.project(t3f.matmul(A + t3f.transpose(A), proj_vec), x)
f = lambda x: t3f.bilinear_form(A, x, x)
res = t3f.hessian_vector_product(f, x, vector)
Args:
func: function that takes TensorTrain object as input and outputs a number.
x: point at which to compute the Hessian and on which tangent space to
project the gradient.
vector: `TensorTrain` object which to multiply be the Hessian.
name: string, name of the Op.
runtime_check: [True] whether to do a sanity check that the passed
function is invariant to different TT representations (otherwise
the Rieamnnian gradient doesn't even exist). It makes things slower,
but helps catching bugs, so turn it off during production deployment.
Returns:
`TensorTrain`, result of the Riemannian hessian by vector product.
See also:
t3f.gradients
"""
all_cores = list(x.tt_cores) + list(vector.tt_cores)
with tf.name_scope(name, values=all_cores):
left = decompositions.orthogonalize_tt_cores(x)
right = decompositions.orthogonalize_tt_cores(left,
left_to_right=False)
deltas = [right.tt_cores[0]]
deltas += [tf.zeros_like(cc) for cc in right.tt_cores[1:]]
x_projection = riemannian.deltas_to_tangent_space(
deltas, x, left, right)
function_value = func(x_projection)
if runtime_check:
assert_op = _is_invariant_to_input_transforms(
function_value, func(x))
else:
assert_op = tf.no_op()
with tf.control_dependencies([assert_op]):
vector_projected = riemannian.project(vector, x)
cores_grad = tf.gradients(function_value, deltas)
vec_deltas = riemannian.tangent_space_to_deltas(vector_projected)
products = [
tf.reduce_sum(a * b) for a, b in zip(cores_grad, vec_deltas)
]
grad_times_vec = tf.add_n(products)
second_cores_grad = tf.gradients(grad_times_vec, deltas)
final_deltas = _enforce_gauge_conditions(second_cores_grad, left)
return riemannian.deltas_to_tangent_space(final_deltas, x, left, right)
def augmented_inner_func(deltas_inner):
x_projection = riemannian.deltas_to_tangent_space(
deltas_inner, x, left, right)
return func(x_projection)