Let $X\sim N(0, I_d)$ be a $d$-dimensional Gaussian random vector. Let $W_1, \ldots, W_k \in \mathbb{R}^d$ be $k$ fixed vectors in general positions. It is clear that $w_i^\top X, \ldots, w_k^\top X$ are jointly Gaussian random variables. Let $$ Y = \max _{i \in [k]} W_i^\top x ,$$ my question is how to compute $\mathbb{E} [ Y\cdot X]$.

My idea is to use a smooth approximation of the max function and the Stein's identity. Let $$ f(y, \alpha ) = \alpha^{-1} \log \left [ \exp(\alpha y_1) + \ldots + \exp( \alpha y_k ) \right ],$$ it is known that $ \left | \max_i y_i - f(y, \alpha) \right | \leq \log k / \alpha$. Then by Stein's identity, we consider \begin{align} \mathbb{E} \left [ f(WX , \alpha) \cdot X \right ] = \mathbb{E} [ \nabla_{X} f(W X , \alpha) ] = \mathbb{E} \left [ \frac{ \sum_{i \in [k]} \exp ( \alpha \cdot W_i ^\top X ) \cdot W_i }{ \sum_{i \in [k]} \exp(\alpha \cdot W_i ^\top X )} \right ]. \end{align} If I naively take $\alpha \rightarrow +\infty$, I would get \begin{align} &\lim_{\alpha\rightarrow +\infty} \mathbb{E} \left [ f(WX , \alpha) \cdot X \right ] = \mathbb{E} \left[ \sum_{j\in[k]} I\left\{ j =\arg\max _{i\in[k]} W_i^\top X \right \} \cdot W_j \right ] \\ &\quad = \sum_{j\in [k]} \mathbb{P}\left (j =\arg\max _{i\in[k]} W_i^\top X \right ) \cdot W_i. \end{align} I wonder whether my naive derivation gets the correct answer and whether such derivation could be made rigorous.