Skip to content

Commit 38242dd

Browse files
A. Unique TensorFlowerVijay Vasudevan
A. Unique TensorFlower
authored and
Vijay Vasudevan
committed
Add new matrix_solve_ls op for solving linear least-squares problems.
Change: 113064195
1 parent 7d4a063 commit 38242dd

12 files changed

Lines changed: 814 additions & 31 deletions

File tree

‎google/protobuf‎

‎tensorflow/core/kernels/matrix_solve_ls_op.cc‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
/* Copyright 2015 Google Inc. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
// See docs in ../ops/linalg_ops.cc.
17+
#include <cmath>
18+
19+
#include "third_party/eigen3/Eigen/Cholesky"
20+
#include "third_party/eigen3/Eigen/Core"
21+
#include "third_party/eigen3/Eigen/QR"
22+
#include "tensorflow/core/framework/kernel_def_builder.h"
23+
#include "tensorflow/core/framework/op_kernel.h"
24+
#include "tensorflow/core/kernels/binary_linalg_ops_common.h"
25+
#include "tensorflow/core/lib/core/errors.h"
26+
#include "tensorflow/core/platform/logging.h"
27+
#include "tensorflow/core/platform/port.h"
28+
#include "tensorflow/core/public/tensor_shape.h"
29+
30+
namespace tensorflow {
31+
32+
template <class Scalar, bool SupportsBatchOperationT>
33+
class MatrixSolveLsOp
34+
: public BinaryLinearAlgebraOp<Scalar, SupportsBatchOperationT> {
35+
public:
36+
explicit MatrixSolveLsOp(OpKernelConstruction* context)
37+
: BinaryLinearAlgebraOp<Scalar, SupportsBatchOperationT>(context) {
38+
OP_REQUIRES_OK(context, context->GetAttr("fast", &fast_));
39+
}
40+
41+
~MatrixSolveLsOp() override {}
42+
43+
TensorShape GetOutputMatrixShape(
44+
const TensorShape& input_matrix_shape,
45+
const TensorShape& rhs_matrix_shape) override {
46+
CHECK_EQ(input_matrix_shape.dims(), rhs_matrix_shape.dims());
47+
TensorShape output_matrix_shape = rhs_matrix_shape;
48+
output_matrix_shape.set_dim(
49+
output_matrix_shape.dims() - 2,
50+
input_matrix_shape.dim_size(output_matrix_shape.dims() - 1));
51+
return output_matrix_shape;
52+
}
53+
54+
int64 GetCostPerUnit(const TensorShape& input_matrix_shape,
55+
const TensorShape& rhs_matrix_shape) override {
56+
const int64 rows = input_matrix_shape.dim_size(0);
57+
const int64 rhss = rhs_matrix_shape.dim_size(1);
58+
if (rows > (1LL << 20)) {
59+
// A big number to cap the cost in case overflow.
60+
return kint32max;
61+
} else {
62+
return 2 * rows * rows * (rows + rhss);
63+
}
64+
}
65+
66+
using typename BinaryLinearAlgebraOp<Scalar, SupportsBatchOperationT>::Matrix;
67+
using typename BinaryLinearAlgebraOp<Scalar,
68+
SupportsBatchOperationT>::MatrixMap;
69+
using typename BinaryLinearAlgebraOp<Scalar,
70+
SupportsBatchOperationT>::ConstMatrixMap;
71+
72+
void ComputeMatrix(OpKernelContext* context, const ConstMatrixMap& matrix,
73+
const ConstMatrixMap& rhs, MatrixMap* output) override {
74+
const int64 rows = matrix.rows();
75+
const int64 cols = matrix.cols();
76+
OP_REQUIRES(
77+
context, rows == rhs.rows(),
78+
errors::InvalidArgument("Input matrix and rhs are incompatible."));
79+
const auto& l2_regularizer_in = context->input(2);
80+
OP_REQUIRES(
81+
context, TensorShapeUtils::IsScalar(l2_regularizer_in.shape()),
82+
errors::InvalidArgument("l2_regularizer must be scalar, got shape ",
83+
l2_regularizer_in.shape().DebugString()));
84+
const double l2_regularizer = l2_regularizer_in.scalar<double>()();
85+
86+
OP_REQUIRES(context, l2_regularizer >= 0,
87+
errors::InvalidArgument("l2_regularizer must be >= 0."));
88+
if (rows == 0 || cols == 0) {
89+
// The result is the empty matrix.
90+
return;
91+
}
92+
if (fast_) {
93+
// The fast branch assumes that matrix is not rank deficient and
94+
// not too ill-conditioned. Specifically, the reciprobal condition number
95+
// should be greater than the square root of the machine precision, i.e.
96+
// 1 / cond(matrix) > sqrt(std::numeric_limits<Scalar>::epsilon()).
97+
// This branch solves over- or underdetermined least-squares problems
98+
// via the normal equations and Cholesky decomposition.
99+
if (matrix.rows() >= matrix.cols()) {
100+
// Overdetermined case (rows >= cols): Solves the ordinary (possibly
101+
// regularized) least-squares problem
102+
// min || A * X - RHS ||_F^2 + l2_regularizer ||X||_F^2
103+
// by solving the normal equations
104+
// (A^T * A + l2_regularizer * I) X = A^T RHS
105+
// using Cholesky decomposition.
106+
Matrix gramian(cols, cols);
107+
gramian.template triangularView<Eigen::Lower>() =
108+
matrix.transpose() * matrix;
109+
if (l2_regularizer > 0) {
110+
gramian +=
111+
(Scalar(l2_regularizer) * Matrix::Ones(cols, 1)).asDiagonal();
112+
}
113+
const Eigen::LLT<Matrix, Eigen::Lower> llt(gramian);
114+
OP_REQUIRES(
115+
context, llt.info() == Eigen::Success,
116+
errors::InvalidArgument("Input matrix was rank deficient or "
117+
"ill-conditioned. Try setting fast=False "
118+
"or provide a larger l2_regularizer > 0."));
119+
*output = llt.solve(matrix.transpose() * rhs);
120+
} else {
121+
// Underdetermined case (rows < cols): Solves the minimum-norm problem
122+
// min ||X||_F^2 s.t. A*X = RHS
123+
// by solving the normal equations of the second kind
124+
// (A * A^T + l2_regularizer * I) Z = RHS, X = A^T * Z
125+
// using Cholesky decomposition.
126+
Matrix gramian(rows, rows);
127+
gramian.template triangularView<Eigen::Lower>() =
128+
matrix * matrix.transpose();
129+
if (l2_regularizer > 0) {
130+
gramian +=
131+
(Scalar(l2_regularizer) * Matrix::Ones(rows, 1)).asDiagonal();
132+
}
133+
const Eigen::LLT<Matrix, Eigen::Lower> llt(gramian);
134+
OP_REQUIRES(
135+
context, llt.info() == Eigen::Success,
136+
errors::InvalidArgument("Input matrix was rank deficient or "
137+
"ill-conditioned. Try setting fast=False "
138+
"or provide an l2_regularizer > 0."));
139+
*output = matrix.transpose() * llt.solve(rhs);
140+
}
141+
} else {
142+
// Use a rank revealing factorization (QR with column pivoting).
143+
//
144+
// NOTICE: Currently, Eigen's implementation of column pivoted Householder
145+
// QR has a few deficiencies:
146+
// 1. It does not implement the post-processing step to compute a
147+
// complete orthogonal factorization. This means that it does not
148+
// return a minimum-norm solution for underdetermined and
149+
// rank-deficient matrices. We could use the Eigen SVD instead, but
150+
// the currently available JacobiSVD is so slow that is it is
151+
// essentially useless (~100x slower than QR).
152+
// 2. The implementation is not blocked, so for matrics that do not fit
153+
// in cache, it is significantly slower than the equivalent blocked
154+
// LAPACK routine xGEQP3 (e.g. Eigen is ~3x slower for 4k x 4k
155+
// matrices). See http://www.netlib.org/lapack/lawnspdf/lawn114.pdf
156+
// 3. The implementation uses the numerically unstable norm downdating
157+
// formula from the original 1965 Businger & Golub paper. This can
158+
// lead to incorrect rank determination for graded matrices. I
159+
// (rmlarsen@) have a patch to bring this up to date by implementing
160+
// the robust formula from
161+
// http://www.netlib.org/lapack/lawnspdf/lawn176.pdf
162+
//
163+
// TODO(rmlarsen): a) Contribute 1. and 2. to Eigen.
164+
// b) Evaluate new divide-and-conquer SVD in Eigen when
165+
// it becomes available & robust.
166+
*output = matrix.colPivHouseholderQr().solve(rhs);
167+
}
168+
}
169+
170+
private:
171+
bool fast_;
172+
};
173+
174+
REGISTER_BINARY_LINALG_OP("MatrixSolveLs", (MatrixSolveLsOp<float, false>),
175+
float);
176+
REGISTER_BINARY_LINALG_OP("MatrixSolveLs", (MatrixSolveLsOp<double, false>),
177+
double);
178+
REGISTER_BINARY_LINALG_OP("BatchMatrixSolveLs", (MatrixSolveLsOp<float, true>),
179+
float);
180+
REGISTER_BINARY_LINALG_OP("BatchMatrixSolveLs", (MatrixSolveLsOp<double, true>),
181+
double);
182+
183+
} // namespace tensorflow

‎tensorflow/core/ops/linalg_ops.cc‎

Lines changed: 93 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@ Calculates the determinant of a square matrix.
2626
2727
input: A tensor of shape `[M, M]`.
2828
output: A scalar, equal to the determinant of the input.
29-
T: The type of values in the input and output.
3029
)doc");
3130

3231
REGISTER_OP("BatchMatrixDeterminant")
@@ -42,7 +41,6 @@ for all input submatrices `[..., :, :]`.
4241
4342
input: Shape is `[..., M, M]`.
4443
output: Shape is `[...]`.
45-
T: The type of values in the input and output.
4644
)doc");
4745

4846
REGISTER_OP("MatrixInverse")
@@ -61,7 +59,6 @@ garbage result.
6159
6260
input: Shape is `[M, M]`.
6361
output: Shape is `[M, M]` containing the matrix inverse of the input.
64-
T: The type of values in the input and output.
6562
)doc");
6663

6764
REGISTER_OP("BatchMatrixInverse")
@@ -84,7 +81,6 @@ garbage result.
8481
8582
input: Shape is `[..., M, M]`.
8683
output: Shape is `[..., M, M]`.
87-
T: The type of values in the input and output.
8884
)doc");
8985

9086
REGISTER_OP("Cholesky")
@@ -103,7 +99,6 @@ input.
10399
104100
input: Shape is `[M, M]`.
105101
output: Shape is `[M, M]`.
106-
T: The type of values in the input and output.
107102
)doc");
108103

109104
REGISTER_OP("BatchCholesky")
@@ -120,7 +115,6 @@ containing the Cholesky decompositions for all input submatrices `[..., :, :]`.
120115
121116
input: Shape is `[..., M, M]`.
122117
output: Shape is `[..., M, M]`.
123-
T: The type of values in the input and output.
124118
)doc");
125119

126120
REGISTER_OP("SelfAdjointEig")
@@ -138,7 +132,6 @@ subsequent rows are eigenvectors.
138132
139133
input: Shape is `[M, M]`.
140134
output: Shape is `[M+1, M]`.
141-
T: The type of values in the input and output.
142135
)doc");
143136

144137
REGISTER_OP("BatchSelfAdjointEig")
@@ -157,7 +150,6 @@ eigenvalues, and subsequent [...,1:, :] containing the eigenvectors.
157150
158151
input: Shape is `[..., M, M]`.
159152
output: Shape is `[..., M+1, M]`.
160-
T: The type of values in the input and output.
161153
)doc");
162154

163155
REGISTER_OP("MatrixSolve")
@@ -172,7 +164,6 @@ matrix: Shape is `[M, M]`.
172164
rhs: Shape is `[M, K]`.
173165
output: Shape is `[M, K]` containing the tensor that solves
174166
matrix * output = rhs.
175-
T: The type of values in the input and output.
176167
)doc");
177168

178169
REGISTER_OP("BatchMatrixSolve")
@@ -191,7 +182,6 @@ matrix satisfies matrix[..., :, :] * output[..., :, :] = rhs[..., :, :].
191182
matrix: Shape is `[..., M, M]`.
192183
rhs: Shape is `[..., M, K]`.
193184
output: Shape is `[..., M, K]`.
194-
T: The type of values in the input and output.
195185
)doc");
196186

197187
REGISTER_OP("MatrixTriangularSolve")
@@ -218,7 +208,6 @@ matrix: Shape is `[M, M]`.
218208
rhs: Shape is `[M, K]`.
219209
output: Shape is `[M, K]`.
220210
lower: Boolean indicating whether matrix is lower or upper triangular.
221-
T: The type of values in the input and output.
222211
)doc");
223212

224213
REGISTER_OP("BatchMatrixTriangularSolve")
@@ -247,7 +236,99 @@ matrix: Shape is `[..., M, M]`.
247236
rhs: Shape is `[..., M, K]`.
248237
output: Shape is `[..., M, K]`.
249238
lower: Boolean indicating whether matrix is lower or upper triangular.
250-
T: The type of values in the input and output.
239+
)doc");
240+
241+
REGISTER_OP("MatrixSolveLs")
242+
.Input("matrix: T")
243+
.Input("rhs: T")
244+
.Input("l2_regularizer: double")
245+
.Output("output: T")
246+
.Attr("T: {float, double}")
247+
.Attr("fast: bool = True")
248+
.Doc(R"doc(
249+
Solves a linear least-squares problem.
250+
251+
Below we will use the following notation
252+
`matrix`=\\(A \in \Re^{m \times n}\\),
253+
`rhs`=\\(B \in \Re^{m \times k}\\),
254+
`output`=\\(X \in \Re^{n \times k}\\),
255+
`l2_regularizer`=\\(\lambda\\).
256+
257+
If `fast` is `True`, then the solution is computed by solving the normal
258+
equations using Cholesky decomposition. Specifically, if \\(m \ge n\\) then
259+
\\(X = (A^T A + \lambda I)^{-1} A^T B\\), which solves the least-squares
260+
problem \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||A Z - B||_F^2 +
261+
\lambda ||Z||_F^2\\). If \\(m \lt n\\) then `output` is computed as
262+
\\(X = A^T (A A^T + \lambda I)^{-1} B\\),
263+
which (for \\(\lambda = 0\\)) is the minimum-norm solution to the
264+
under-determined linear system, i.e.
265+
\\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||Z||_F^2 \\),
266+
subject to \\(A Z = B\\).
267+
Notice that the fast path is only numerically stable when \\(A\\) is
268+
numerically full rank and has a condition number
269+
\\(\mathrm{cond}(A) \lt \frac{1}{\sqrt{\epsilon_{mach}}}\\)
270+
or \\(\lambda\\) is sufficiently large.
271+
272+
If `fast` is `False` then the solution is computed using the rank revealing QR
273+
decomposition with column pivoting. This will always compute a least-squares
274+
solution that minimizes the residual norm \\(||A X - B||_F^2 \\), even when
275+
\\( A \\) is rank deficient or ill-conditioned. Notice: The current version
276+
does not compute a minimum norm solution. If `fast` is `False` then
277+
`l2_regularizer` is ignored.
278+
279+
matrix: Shape is `[M, N]`.
280+
rhs: Shape is `[M, K]`.
281+
output: Shape is `[N, K]` containing the tensor that solves
282+
`matrix * output = rhs` in the least-squares sense.
283+
)doc");
284+
285+
REGISTER_OP("BatchMatrixSolveLs")
286+
.Input("matrix: T")
287+
.Input("rhs: T")
288+
.Input("l2_regularizer: double")
289+
.Output("output: T")
290+
.Attr("T: {float, double}")
291+
.Attr("fast: bool = True")
292+
.Doc(R"doc(
293+
Solves multiple linear least-squares problems.
294+
295+
`matrix` is a tensor of shape `[..., M, N]` whose inner-most 2 dimensions
296+
form square matrices. Rhs is a tensor of shape `[..., M, K]`. The output
297+
is a tensor shape `[..., N, K]` where each output matrix solves each of
298+
the equations matrix[..., :, :] * output[..., :, :] = rhs[..., :, :] in the
299+
least squares sense.
300+
301+
Below we will use the following notation for each pair of
302+
matrix and right-hand sides in the batch:
303+
304+
`matrix`=\\(A \in \Re^{m \times n}\\),
305+
`rhs`=\\(B \in \Re^{m \times k}\\),
306+
`output`=\\(X \in \Re^{n \times k}\\),
307+
`l2_regularizer`=\\(\lambda\\).
308+
309+
If `fast` is `True`, then the solution is computed by solving the normal
310+
equations using Cholesky decomposition. Specifically, if \\(m \ge n\\) then
311+
\\(X = (A^T A + \lambda I)^{-1} A^T B\\), which solves the least-squares
312+
problem \\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||A Z - B||_F^2 +
313+
\lambda ||Z||_F^2\\). If \\(m \lt n\\) then `output` is computed as
314+
\\(X = A^T (A A^T + \lambda I)^{-1} B\\), which (for \\(\lambda = 0\\)) is the
315+
minimum-norm solution to the under-determined linear system, i.e.
316+
\\(X = \mathrm{argmin}_{Z \in \Re^{n \times k}} ||Z||_F^2 \\), subject to
317+
\\(A Z = B\\). Notice that the fast path is only numerically stable when
318+
\\(A\\) is numerically full rank and has a condition number
319+
\\(\mathrm{cond}(A) \lt \frac{1}{\sqrt{\epsilon_{mach}}}\\) or\\(\lambda\\) is
320+
sufficiently large.
321+
322+
If `fast` is `False` then the solution is computed using the rank revealing QR
323+
decomposition with column pivoting. This will always compute a least-squares
324+
solution that minimizes the residual norm \\(||A X - B||_F^2\\), even when
325+
\\(A\\) is rank deficient or ill-conditioned. Notice: The current version does
326+
not compute a minimum norm solution. If `fast` is `False` then `l2_regularizer`
327+
is ignored.
328+
329+
matrix: Shape is `[..., M, N]`.
330+
rhs: Shape is `[..., M, K]`.
331+
output: Shape is `[..., N, K]`.
251332
)doc");
252333

253334
} // namespace tensorflow

0 commit comments

Comments
 (0)