As in the existing work [1]}, [2]},
we learn \([W|a,z]{\vec{\psi }_{\theta _{W}}(w)}\)  in Stage 1 and \(h(a,w)\)  in Stage 2, but in addition to the weights \(\vec{u}\)  and \(\vec{V}\) , we also learn the parameters of the feature maps.
Specifically, in Stage 1, we learn \(\vec{V}\)  and parameters \(\theta _{A(1)}, \theta _Z\)  by minimizing the following empirical loss:

L1(V, A(1), Z) = 1m i=1m W(wi) - V(A(1)(ai) Z(zi)) 2 + 1 V2.

Note that \(\hat{\mathcal {L}}_1\)  is an empirical estimation of \(\mathcal {L}_1\)  in (REF ) with adaptive feature maps.
Although \(\hat{\mathcal {L}}_1\)  depends on \(\theta _{W}\) , at this stage, we do not optimize \(\theta _{W}\)  with \(\hat{\mathcal {L}}_1\)  as \({\vec{\psi }_{\theta _{W}}(w)}\)  is the “target variable” in Stage 1. Given the minimizers \((\hat{\vec{V}}, \hat{\theta }_{A(1)}, \hat{\theta }_{Z}) = \hat{\mathcal {L}}_1\) , we learn weights \(\vec{u}\)  and parameters \(\theta _W, \theta _{A(2)}\)  by minimizing the empirical stage 2 loss

L2(u, W, A(2)) = 1n i=1n (yi - u(A(2)(ai) V (A(1)(ai) Z(zi))))2 + 2 u2.

Again, the loss \(\hat{\mathcal {L}}_2\)  is an empirical estimate of \(\mathcal {L}_2\)  in (REF ) with adaptive feature maps.
Although the expression of \(\hat{\mathcal {L}}_2\)  does not explicitly contain \(\theta _W\) ,
it implicitly depends on \(\theta _W\)  as \((\hat{\vec{V}}, \hat{\theta }_{A(1)}, \hat{\theta }_{Z})\)  is the solution of a minimization problem involving \(\theta _W\) . This implicit dependency makes it challenging to update \(\theta _W\) , as we cannot directly obtain its gradient. One possible solution is to use the implicit gradient method [3]}, but this approach might be too computationally expensive. Instead, we use the method proposed in [4]}, in which we ignore the dependency of \(\theta _W\)  on parameters \(\hat{\theta }_{A(1)}, \hat{\theta }_{Z}\) , and the gradient is computed via the closed-form solution of \(\hat{V}\) .
