Hello Nick,
The formula is correct.
I can suggest an alternative approach as well. Given a point X on the manifold and the tangent space description,
Riemannian gradient = arg min_eta - <Euclidean gradient, eta> + 0.5 g(eta, eta), subject to eta in the tangent space at X,
where
eta is the variable,
< . , . > is the standard inner product, and
g(. , .) is the Riemannian metric imposed on the manifold.
For the Stiefel manifold, the Riemannian metric is also the standard inner product. Therefore, given the Euclidean Gradient of f(X) = Z, the above formula boils down to
Riemannian gradient = arg min_eta - trace(eta' Z) + 0.5 trace(eta' eta), subject to eta'X + X'eta = 0,
which is precisely the formula that you got.
Regards,
BM