Regress linear regression coefficient with neural networks TLDR:
to a NN you give a dataset and it predicts the $\beta$ that linear regression would have found if solved on that dataset
Context:
I have a dataset of datasets (synthetical) with the same number of feature, so it's $D = R^{m\times n\times d}$ where:

*

*$m$ is the number of datasets

*$n$ is the number of samples in each datasets

*$d$ the dimensionality of the regression (+1 since I add a row of 1 for the biases)

So usually the usual linear regression would be done on $D_0$, finding $d$ coefficients.
What I tried:

*

*as encoder, I've stacked some bidirectional LSTM

*at each LSTM step it gets the pair $(x_i, y_i)$ of the current dataset being processed, where $x_i \in R^d$ and $y_i \in R$

*the last LSTM returns a single $512$ long vector, that should represent the dataset

*the $512$ vector is then fed to some Dense layers to convert that to a $d-$dimensional vector (that should be the $\beta$ coefficients of the linear regression.)

At this point, the NN should return the coefficients.
Now, the whole point of the experiment is to regress them without knowing the actual coefficients, so the "output of the NN" is unknown, however, my idea is to give a "feedback" to the NN on how good the coefficients are, thus minimizing the following error:
$$
L_\theta(X,Y) = \frac{1}{m}\sum_{(x,y) \in (xs,ys)}\frac{1}{n} \sum_{(x_i,y_i) \in (x,y)} (x_i^Tf_\theta(x,y) - y_i)^2
\\[15mm]
\text{Given: }xs \in R^{m\times n\times d}, \,\,\,\,\, ys \in R^{m\times d}\\
$$
However, the NN is not giving very nice results, and stacking more LSTMs or Dense layers, or making them wider, does not seem to improve in any way the results...
So at this point I think I'm missing some theoretical basis on why this should not work.
In case you want to test it on your PC, this is the full code you need:
import sklearn.decomposition
import tensorflow as tf
import tensorflow.keras as keras
import numpy as np
from tensorflow.keras import layers
import keras.api._v2.keras as K
import matplotlib.pyplot as plt

#
# Generate data
#
COEFF_W = 3
DATA_W = 100
DATASET_SIZE = 1000
xs = []
ys = []
targets = []
for i in range(DATASET_SIZE):
    coeffs = np.random.normal(0, 100, COEFF_W)
    targets.append(coeffs)
    data = []
    for _ in range(COEFF_W-1):
        min_ = 0
        max_ = 1
        x = min_ + np.random.rand(DATA_W) * max_
        data.append(x)
    data = np.array(data).T
    data = np.concatenate((data, np.ones((DATA_W,1))), axis=1)
    xs.append(data)
    ys.append(data @ coeffs + np.random.normal(0,np.random.rand() * 5, size = DATA_W))

xs = np.array(xs)
ys = np.array(ys)
targets = np.array(targets)

#
# Create the model
#
class LearnLinearRegression(K.Model):
    def __init__(self):
        super().__init__()
        self.lstms = [
            K.layers.Bidirectional(K.layers.LSTM(512, return_sequences=True)),
            K.layers.Bidirectional(layers.LSTM(512)),
        ]
        self.denses_final = [
            K.layers.Dense(units=2048, activation=tf.nn.leaky_relu),
            K.layers.Dense(units=2048, activation=tf.nn.leaky_relu),
        ]
        self.dense_coeffs = K.layers.Dense(COEFF_W, activation="linear")

    def train_step(self, data):
        x,y = data
        xy = tf.concat((x,y[:,:,tf.newaxis]), axis=-1)
        with tf.GradientTape() as tape:
            c = self.call(xy)
            preds = tf.reduce_sum(x * c[:,None,:], axis=-1)
            loss = tf.reduce_mean(tf.square(preds - y))
        grads = tape.gradient(loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        return {
            "loss" : loss
        }
    def test_step(self, data):
        x,y = data
        xy = tf.concat((x,y[:,:,tf.newaxis]), axis=-1)
        c = self.call(xy, training=False)
        preds = tf.reduce_sum(x * c[:,None,:], axis=-1)
        loss = tf.reduce_mean(tf.square(preds - y))
        return {"loss" : loss}

    def call(self, inputs, training=True):
        x = inputs
        for lstm in self.lstms: x = lstm(x)
        for dense in self.denses_final: x = dense(x)
        coeffs = self.dense_coeffs(x)
        return coeffs

linear_regression_coefficients_regressor = LearnLinearRegression()
linear_regression_coefficients_regressor.compile(optimizer=K.optimizers.Adam(0.0001))
hist = linear_regression_coefficients_regressor.fit(
    xs, ys, epochs=100, batch_size=200, validation_split=0.1,
    callbacks=[K.callbacks.EarlyStopping(patience=10)]
)

Now if you have Plotly installed you can use the following code to check the predictions:
from  sklearn.linear_model import LinearRegression
import plotly.express as px
import plotly.graph_objects as go

# NN predictions
coeffs = linear_regression_coefficients_regressor.predict(tf.concat((xs,ys[:,:,tf.newaxis]), axis=-1), verbose = 0)
predictions = np.sum(xs * coeffs[:,None,:], axis=-1)[![enter image description here][1]][1]
mses = np.mean(np.square(predictions - ys), axis=0)
argmin_mses = np.argmin(mses)

# Closed form predictions
clf = LinearRegression()
clf.fit(xs[argmin_mses,:,:-1], ys[argmin_mses])
clf_pref = clf.predict(xs[argmin_mses,:,:-1])

#
# Result
#

import plotly.express as px
import plotly.graph_objects as go

fig = go.Figure(data=[
    go.Scatter3d(x=xs[argmin_mses,:,0],y=xs[argmin_mses,:,1],z=predictions[argmin_mses],
                                    name="neural net",
                                   mode='markers',
                                   marker=dict(
                                       size=5,
                                       color=0,                # set color to an array/list of desired values
                                       colorscale='Viridis',   # choose a colorscale
                                       opacity=0.8
                                   )),
    go.Scatter3d(x=xs[argmin_mses,:,0],y=xs[argmin_mses,:,1],z=ys[argmin_mses],
                 mode='markers',
                 name="data",
                 marker=dict(
                     size=5,
                     color=1,                # set color to an array/list of desired values
                     colorscale='Viridis',   # choose a colorscale
                     opacity=0.8
                 )),
    go.Scatter3d(x=xs[argmin_mses,:,0],y=xs[argmin_mses,:,1],z=clf_pref,
                 mode='markers',
                 name="linear reg",
                 marker=dict(
                     size=5,
                     color=2,                # set color to an array/list of desired values
                     colorscale='Viridis',   # choose a colorscale
                     opacity=0.8,
                 ))
],)
fig.show(renderer="chrome")

With the fitted NN, this is the result...


This should be doable (in theory).
The universal approximation theorems say that any decent function, such as a linear transformation, can be approximated as well as is desired, given a large enough neural network.
If you give the linear regression features ($X$) and outcomes ($y$) as features to a neural network and the $\hat\beta$ as the outcome for the neural network, some network should be able to figure out that $\hat\beta=(X^TX)^{-1}X^Ty$.
(It’s also worth noting that a linear regression can be seen as a neural network with no hidden layer and the identity as an activation function.)
Why anyone would want to do this is not clear to me, but math says that it can be done.
