def testPairwiseFlatInnerMatrix(self):
# Compare pairwise_flat_inner_projected against naive implementation.
what1 = initializers.random_matrix_batch(((2, 3, 4), None), 4, batch_size=3,
dtype=self.dtype)
what2 = initializers.random_matrix_batch(((2, 3, 4), None), 4, batch_size=4,
dtype=self.dtype)
where = initializers.random_matrix(((2, 3, 4), None), 3,
dtype=self.dtype)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired = batch_ops.pairwise_flat_inner(projected1, projected2)
actual = riemannian.pairwise_flat_inner_projected(projected1, projected2)
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
with self.assertRaises(ValueError):
# Second argument is not a projection on the tangent space.
riemannian.pairwise_flat_inner_projected(projected1, what2)
where2 = initializers.random_matrix(((2, 3, 4), None), 3,
dtype=self.dtype)
another_projected2 = riemannian.project(what2, where2)
with self.assertRaises(ValueError):
# The arguments are projections on different tangent spaces.
riemannian.pairwise_flat_inner_projected(projected1, another_projected2)
def testHessianVectorProduct(self):
w = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
A = initializers.random_matrix(([5] * 3, [5] * 3), dtype=self.dtype)
x = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
z = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
projected_vector = riemannian.project(z, x)
def func1(x):
return 0.5 * ops.flat_inner(x, w) ** 2
# Grad: <x, w> w
# Hessian: w w.T
# Hessian by vector: w <w, P_x z>
desired1 = riemannian.project(w * ops.flat_inner(projected_vector, w), x)
desired1 = ops.full(desired1)
self._TestSingleHessianByVector(func1, x, z, desired1)
def func2(x):
return ops.bilinear_form(A, x, x)
# Hessian of <x, Ax> is A + A.T
hessian_by_vector = ops.matmul(ops.transpose(A) + A, projected_vector)
desired2 = ops.full(riemannian.project(hessian_by_vector, x))
self._TestSingleHessianByVector(func1, x, z, desired1)
def func3(x):
# A function which is not invariant to different representations of the
# same tensor, i.e. it does not even have a Riemannian gradient or
# hessian.
return tf.add_n([tf.reduce_sum(c) for c in x.tt_cores]) ** 2
with self.assertRaises(tf.errors.InvalidArgumentError):
actual3 = ops.full(autodiff.hessian_vector_product(func3, x, z))
self.evaluate(actual3)
def testWeightedAddNProjectedBatch(self):
# Add several TT-batches from the same tangent space with coefs.
what1 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=3)
what2 = initializers.random_tensor_batch((2, 3, 4), 1, batch_size=3)
where = initializers.random_tensor((2, 3, 4), 3)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired_0 = ops.full(1.2 * projected1[0] + -2.0 * projected2[0])
desired_1 = ops.full(1.9 * projected1[1] + 2.0 * projected2[1])
desired_2 = ops.full(0.0 * projected1[2] + 1.0 * projected2[2])
desired = tf.stack((desired_0, desired_1, desired_2), axis=0)
actual = ops.full(riemannian.add_n_projected((projected1, projected2),
coef=[[1.2, 1.9, 0.0],
[-2.0, 2.0, 1.0]]))
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
def testCompareProjectSumAndProject(self):
# Compare results of project_sum and project.
tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4)
tangent_tens = initializers.random_tensor((2, 3, 4), 4)
project_sum = riemannian.project_sum(tens, tangent_tens, tf.eye(4))
project = riemannian.project(tens, tangent_tens)
with self.test_session() as sess:
res = sess.run((ops.full(project_sum), ops.full(project)))
project_sum_val, project_val = res
self.assertAllClose(project_sum_val, project_val)
def testCompareProjectSumAndProject(self):
# Compare results of project_sum and project.
tens = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=4,
dtype=self.dtype)
tangent_tens = initializers.random_tensor((2, 3, 4), 4,
dtype=self.dtype)
project_sum = riemannian.project_sum(tens, tangent_tens, np.eye(4))
project = riemannian.project(tens, tangent_tens)
res = self.evaluate((ops.full(project_sum), ops.full(project)))
project_sum_val, project_val = res
self.assertAllClose(project_sum_val, project_val)
def testAddNProjected(self):
# Add several TT-objects from the same tangent space.
what1 = initializers.random_tensor_batch((2, 3, 4), 4, batch_size=3)
what2 = initializers.random_tensor_batch((2, 3, 4), 3, batch_size=3)
where = initializers.random_tensor((2, 3, 4), 3)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired = ops.full(projected1 + projected2)
actual = ops.full(riemannian.add_n_projected((projected1, projected2)))
with self.test_session() as sess:
desired_val, actual_val = sess.run((desired, actual))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
with self.assertRaises(ValueError):
# Second argument is not a projection on the tangent space.
riemannian.add_n_projected((projected1, what2))
where2 = initializers.random_tensor((2, 3, 4), 3)
another_projected2 = riemannian.project(what2, where2)
with self.assertRaises(ValueError):
# The arguments are projections on different tangent spaces.
riemannian.add_n_projected((projected1, another_projected2))
def testProjectMatmul(self):
# Project a TT-matrix times TT-vector on a TT-vector.
tt_mat = initializers.random_matrix(((2, 3, 4), (2, 3, 4)))
tt_vec_what = initializers.random_matrix_batch(((2, 3, 4), None),
batch_size=3)
tt_vec_where = initializers.random_matrix(((2, 3, 4), None))
proj = riemannian.project_matmul(tt_vec_what, tt_vec_where, tt_mat)
matvec = ops.matmul(tt_mat, tt_vec_what)
proj_desired = riemannian.project(matvec, tt_vec_where)
with self.test_session() as sess:
actual_val, desired_val = sess.run((ops.full(proj), ops.full(proj_desired)))
self.assertAllClose(desired_val, actual_val, atol=1e-5, rtol=1e-5)
def testWeightedAddNProjected(self):
# Add several TT-objects from the same tangent space with coefs.
what1 = initializers.random_tensor((2, 3, 4), 4, dtype=self.dtype)
what2 = initializers.random_tensor((2, 3, 4), 1, dtype=self.dtype)
where = initializers.random_tensor((2, 3, 4), 3, dtype=self.dtype)
projected1 = riemannian.project(what1, where)
projected2 = riemannian.project(what2, where)
desired = ops.full(1.2 * projected1 + -2.0 * projected2)
actual = ops.full(riemannian.add_n_projected((projected1, projected2),
coef=[1.2, -2.0]))
desired_val, actual_val = self.evaluate((desired, actual))
self.assertAllClose(desired_val, actual_val)
with self.assertRaises(ValueError):
# Second argument is not a projection on the tangent space.
riemannian.add_n_projected((projected1, what2), coef=[1.2, -2.0])
where2 = initializers.random_tensor((2, 3, 4), 3, dtype=self.dtype)
another_projected2 = riemannian.project(what2, where2)
with self.assertRaises(ValueError):
# The arguments are projections on different tangent spaces.
riemannian.add_n_projected((projected1, another_projected2),
coef=[1.2, -2.0])
def testGradients(self):
w = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
A = initializers.random_matrix(([5] * 3, [5] * 3), dtype=self.dtype)
x = initializers.random_matrix(([5] * 3, None), dtype=self.dtype)
def func1(x):
return 0.5 * ops.flat_inner(x, w) ** 2
desired1 = ops.full(riemannian.project(w, x) * ops.flat_inner(x, w))
self._TestSingleGradient(func1, x, desired1)
def func2(x):
return ops.bilinear_form(A, x, x)
grad = ops.matmul(ops.transpose(A) + A, x)
desired2 = ops.full(riemannian.project(grad, x))
self._TestSingleGradient(func2, x, desired2)
def func3(x):
# A function which is not invariant to different representations of the
# same tensor, i.e. it does not even have a Riemannian gradient.
return tf.add_n([tf.reduce_sum(c) for c in x.tt_cores]) ** 2
with self.assertRaises(tf.errors.InvalidArgumentError):
actual3 = ops.full(autodiff.gradients(func3, x))
self.evaluate(actual3)
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 augmented_outer_func(deltas_outer):
def augmented_inner_func(deltas_inner):
x_projection = riemannian.deltas_to_tangent_space(
deltas_inner, x, left, right)
return func(x_projection)
function_value, cores_grad = value_and_grad(
augmented_inner_func, deltas_outer)
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)
vec_deltas = riemannian.tangent_space_to_deltas(vector_projected)
products = [
tf.reduce_sum(a * b) for a, b in zip(cores_grad, vec_deltas)
]
return tf.add_n(products)
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)
