Coupled NMF for Partially Linked Features

Why Couple Only Selected Features?

Two omics matrices may measure the same samples but contain feature sets that overlap only partly. For example, a known correspondence may connect selected proteins to selected transcripts, while other features have no counterpart in the second assay. Independent nonnegative matrix factorizations do not use that correspondence; forcing all features to be identical would impose more than the data justify. Coupled NMF provides a middle ground: retain separate feature loadings, but penalize disagreement between loadings that a known map says should correspond.

This formulation assumes the rows of both matrices refer to the same \(n\) samples in the same order. It uses one shared nonnegative sample-score matrix and separate nonnegative feature-loading matrices. The earlier material on shared and modality-specific components distinguishes shared from view-specific structure; here the focus is the additional, explicit relationship between selected feature loadings.

Definition: Let \(X\in\mathbb{R}_+^{n\times p}\) and \(Y\in\mathbb{R}_+^{n\times q}\) be nonnegative data matrices for the same \(n\) samples. At rank \(r\), write \(X\approx WH_X^T\) and \(Y\approx WH_Y^T\), where \(W\in\mathbb{R}_+^{n\times r}\), \(H_X\in\mathbb{R}_+^{p\times r}\), and \(H_Y\in\mathbb{R}_+^{q\times r}\). A fixed nonnegative mapping matrix \(M\in\mathbb{R}_+^{p\times q_L}\) maps the \(p\) features of \(X\) to a selected set of \(q_L\) linked features in \(Y\). The mapped loadings are \(M^TH_X\in\mathbb{R}_+^{q_L\times r}\).

Let \(S\in\mathbb{R}^{q\times q_L}\) select the linked columns of \(Y\), so \(H_{Y,L}=S^TH_Y\) contains only their loading rows. The map \(M\) must be defined in the direction used above: each column describes how one linked \(Y\)-feature is represented using the \(X\)-features. A column may encode a one-to-one match or a weighted combination. Its entries should be nonnegative if the map is intended to preserve nonnegative loadings. The map is fixed and known; it is not estimated as part of this objective.

The Coupled Objective

For a squared-Frobenius reconstruction loss, a basic soft-coupling objective is

\mathcal{L}(W,H_X,H_Y) =\frac{1}{2}\|X-WH_X^T\|_F^2 +\frac{1}{2}\|Y-WH_Y^T\|_F^2 +\frac{\lambda}{2}\|H_{Y,L}-M^TH_X\|_F^2, \qquad W,H_X,H_Y\geq 0.

The first two terms measure reconstruction error in the two assays. The final term is the coupling penalty, and \(\lambda\geq 0\) controls its strength. With \(\lambda=0\), the loading relationship has no effect on the fit. Increasing \(\lambda\) makes disagreement between linked loadings more costly, but does not require exact equality at a finite value. Rows of \(H_Y\) for unlinked features do not appear in the coupling term; they remain free to help reconstruct \(Y\).

The scale convention in this objective matters: each squared loss has a factor of \(1/2\), which simplifies derivatives. In applications, the reconstruction terms may also be weighted to reflect different block scales or priorities. Such weights must be specified rather than assumed to be interchangeable with \(\lambda\). Rescaling one matrix or changing the normalization of \(M\) changes the relative numerical influence of the coupling penalty.

Worked Example: Calculating the Coupled Objective

Suppose two hypothetical matrices have two shared samples. At rank one, let \(X=\begin{bmatrix}2&1\\3&2\end{bmatrix}\), \(Y=\begin{bmatrix}1.5&1&4\\3&4&7\end{bmatrix}\), \(W=\begin{bmatrix}1\\2\end{bmatrix}\), \(H_X=\begin{bmatrix}2\\1\end{bmatrix}\), and \(H_Y=\begin{bmatrix}1.5\\2\\4\end{bmatrix}\). Assume only the first two features of \(Y\) are linked and \(M\) is the \(2\times2\) identity matrix. Take \(\lambda=0.4\). All factors and data are nonnegative, the sample order agrees, and the map relates the two linked feature loadings directly.

The reconstructions are \(\widehat X=WH_X^T=\begin{bmatrix}2&1\\4&2\end{bmatrix}\) and \(\widehat Y=WH_Y^T=\begin{bmatrix}1.5&2&4\\3&4&8\end{bmatrix}\). For \(X\), the only nonzero residual is \(-1\), so \(\|X-\widehat X\|_F^2=1\). For \(Y\), two residuals are \(-1\), so \(\|Y-\widehat Y\|_F^2=2\).

The linked \(Y\)-loadings are \((1.5,2)^T\), while the mapped \(X\)-loadings are \(M^TH_X=(2,1)^T\). Their difference is \((-0.5,1)^T\), with squared norm \(0.25+1=1.25\). Therefore \(\mathcal{L}=0.5(1)+0.5(2)+(0.4/2)(1.25)=1.75\). The coupling contributes \(0.25\) loss units to this objective, in addition to \(1.5\) units of reconstruction loss.

Deriving the Coupling Contribution

Set \(E=H_{Y,L}-M^TH_X\). The coupling part of the objective is \(\mathcal{L}_c=(\lambda/2)\|E\|_F^2\). Since \(\|E\|_F^2=\operatorname{tr}(E^TE)\), its differential is

d\mathcal{L}_c =\lambda\operatorname{tr}(E^T dE) =\lambda\operatorname{tr}\!\left(E^T(dH_{Y,L}-M^TdH_X)\right).

Using the Frobenius inner product to collect the coefficients of each change gives the gradients

\nabla_{H_{Y,L}}\mathcal{L}_c=\lambda E, \qquad \nabla_{H_X}\mathcal{L}_c=-\lambda ME.

The first gradient pushes the linked \(Y\)-loadings toward the mapped \(X\)-loadings. The second pushes \(H_X\) in the direction that makes its mapped loadings closer to \(H_{Y,L}\). Unlinked rows of \(H_Y\) have zero coupling gradient. If a solver uses the full \(H_Y\), the linked gradient is placed in the selected rows and zeros in the other rows.

This derivation also explains why the penalty is a soft relationship rather than a relabeling of features: each side can move, and the reconstruction terms may favor some remaining discrepancy. For fixed \(M\), the penalty is a convex quadratic in either loading matrix when the other is held fixed. The full factorization is still nonconvex because \(W\) and the loading matrices are jointly estimated through products.

Worked Example: Evaluating the Coupling Gradients

Continue with the rank-one loadings \(H_X=(2,1)^T\) and linked loadings \(H_{Y,L}=(1.5,2)^T\), but use the hypothetical map \(M=\begin{bmatrix}1&0.5\\0&1\end{bmatrix}\) and \(\lambda=0.4\). The map is nonnegative and has two linked target features. It predicts \(M^TH_X=(2,2)^T\), so \(E=(-0.5,0)^T\).

The gradient for the linked \(Y\)-loadings is \(\lambda E=0.4(-0.5,0)^T=(-0.2,0)^T\). For the \(X\)-loadings, \(ME=\begin{bmatrix}1&0.5\\0&1\end{bmatrix}(-0.5,0)^T=(-0.5,0)^T\), so \(\nabla_{H_X}\mathcal{L}_c=-0.4(-0.5,0)^T=(0.2,0)^T\).

These are gradients of the coupling term alone, not complete optimization updates. In an unconstrained negative-gradient direction, the first linked \(Y\)-loading would increase and the first \(X\)-loading would decrease. Reconstruction gradients and nonnegativity constraints also affect an actual solver step. The coupling loss here is \((0.4/2)((-0.5)^2+0^2)=0.05\).

Worked Example: A Many-to-One Map and an Unlinked Feature

Suppose a hypothetical \(X\) block has two features and \(Y\) has two features, only the first of which is linked. Let \(r=1\), \(H_X=(2,4)^T\), \(H_Y=(3,5)^T\), and \(M=\begin{bmatrix}0.25\\0.75\end{bmatrix}\). The single column of \(M\) assigns nonnegative weights summing to one to the two \(X\)-features. The first \(Y\)-feature is the linked target; the second is explicitly unlinked.

The mapped loading is \(M^TH_X=0.25(2)+0.75(4)=3.5\), while the linked \(Y\)-loading is \(3\). Thus the squared discrepancy is \((3-3.5)^2=0.25\). At \(\lambda=2\), the coupling contribution is \((2/2)(0.25)=0.25\).

The loading \(5\) for the unlinked \(Y\)-feature contributes nothing to this penalty. It can still affect the reconstruction of the second \(Y\)-feature through \(WH_Y^T\). This distinction is essential: treating the map as if it linked every feature would add an unjustified penalty.

Assumptions, Checks, and Alternatives

The formulation is meaningful only if its structural assumptions match the data and the mapping. Check that sample rows really correspond and are ordered consistently; a shared \(W\) otherwise forces unrelated observations to share scores. Check the orientation and dimensions of \(M\), and verify that its columns represent the intended linked target features. The map may encode weighted correspondences, but those weights need a substantive interpretation. In particular, normalizing columns to sum to one is sensible for some averaging maps, not a universal requirement.

The penalty is measured in squared loading units, so the relative size of the two blocks and the scale of their loadings affect the useful range of \(\lambda\). Feature scaling, factor scale indeterminacy, and rescaling \(M\) can all change the numerical balance. The factorization has the usual scale ambiguity: multiplying a column of \(W\) by a positive constant and dividing the corresponding columns of both loading matrices by that constant preserves the reconstructions, while it changes the coupling penalty. Consequently, factor normalization or other scale conventions can matter when comparing coupling strengths.

Useful checks include inspecting the linked residual matrix \(E\), calculating reconstruction and coupling losses separately, confirming that unlinked rows receive no coupling contribution, and trying a range of \(\lambda\) values. A very large coupling loss may indicate a poor map, incompatible feature scales, or insufficient flexibility in the chosen rank; it is not by itself proof that the biological correspondence is false. Conversely, a small penalty can result from a small \(\lambda\), rather than close agreement. These checks diagnose the specification; they do not establish that a feature relationship is biologically causal.

The main alternatives impose different assumptions. Independent NMF fits each block without using the feature map. A hard constraint sets \(H_{Y,L}=M^TH_X\) exactly, which can be appropriate only when that equality is required by design. A shared-loading model is a special, more restrictive choice when the feature spaces line up in a way that makes direct equality meaningful. Soft coupling is useful when correspondences are informative but should not override block-specific reconstruction. Its limitations remain: results depend on the chosen map, rank, scaling, and penalty strength, and the shared-\(W\) version requires aligned samples.

Common Pitfalls

Wrong reasoning: “The matrices have the same number of features, so their loading rows can be equated.” Correction: Equal dimensions do not establish feature correspondence. Specify a justified map and couple only the linked features.
Wrong reasoning: “A feature absent from the map should be mapped to zero and penalized against zero.” Correction: If it is unlinked, omit it from the coupling term. A zero map column inside a penalty would instead pull its target loading toward zero.
Wrong reasoning: “A larger \(\lambda\) always gives a better integration.” Correction: It gives the coupling discrepancy greater weight; excessive weight can degrade reconstruction or force an unsuitable relationship.
Wrong reasoning: “The coupling gradient is the whole update.” Correction: It is only one part of the gradient. Reconstruction terms, nonnegativity, and the chosen optimizer also determine the update.

Key takeaway: Coupled NMF adds a penalty on the difference between selected \(Y\)-loadings and the mapped \(X\)-loadings. Correct dimensions, sample alignment, link selection, and scale interpretation are as important as the algebra.

Check Your Understanding

Answer each question, then compare with the concise solutions below.

If \(H_X\) has dimensions \(p\times r\) and there are \(q_L\) linked target features, what dimensions must \(M\) have for \(M^TH_X\) to be defined?
Let \(E=H_{Y,L}-M^TH_X\). State the two coupling gradients.
For \(E=(0.2,-0.1)^T\) and \(\lambda=3\), calculate the coupling loss.
Why should an unlinked feature of \(Y\) not automatically be represented by a zero column of \(M\) inside the penalty?
Name two assumptions that should be checked before fitting this shared-\(W\) formulation.

Solutions. (1) \(M\) must be \(p\times q_L\), so \(M^T H_X\) is \(q_L\times r\). (2) The gradients are \(\lambda E\) with respect to \(H_{Y,L}\), and \(-\lambda ME\) with respect to \(H_X\). (3) The squared norm is \(0.2^2+(-0.1)^2=0.05\); the loss is \((3/2)(0.05)=0.075\). (4) A zero map column in the penalty would make the discrepancy equal to that feature's loading and penalize it toward zero, although the feature is supposed to be unconstrained by the map. (5) The sample rows must represent the same samples in the same order, and the dimensions, direction, and substantive meaning of the feature map must be correct.