Learning affinity among features

A batch of semantic objects in the image (lesions in CT scans) are represented in feature space, $X_{B times C}$. I want to represent the whole batch in a single vector, $1 times C$, in order to predict the image class from this vector. I want to find an affinity for each feature. This affinity matrix is $A_{C_2times C}$, and I take a matrix-matrix product: $X_2 =A^{T}X$. Then, I take matrix-vector product of $X_2$ and affinity vector $a_{C_2 times 1}$ to get $x_3= X_2 a$ which has $C$ dimensions and hence the representation I want.

My question is: is this a legit method of obtaining an affinity among objects? I plotted $X_2$ after 1 and 100 epochs, it looks good. Overall accuracy is also good, around $90%$, but I have doubts about the method-never did anything like this before.

In the first step, for each feature $mathbf{x}_j$ I learn $C_2$ affinities, so I map each feature into the affinity space. In the next step, I scale these affinities with vector $a$ and map the features back into the feature space. Is this OK?

<img src="https://i.stack.imgur.com/1R8RG.jpg" alt="$X_2$ matrix” />