使用equinox与optax实现逻辑回归时模型无法训练的问题

我正尝试借助JAX，使用equinox和optax库实现逻辑回归，但训练过程中损失值始终不下降，模型无法学习。同一数据集下，scikit-learn和PyTorch的逻辑回归实现运行正常，相关代码也附在下方作为参考。我尝试过SGD和Adam优化器以及不同学习率，还试过零权重初始化与随机权重初始化，但结果一致，且在多个分类数据集上测试均遇到该问题。

可复现代码 #

import jax
import jax.nn as jnn
import jax.numpy as jnp
import jax.random as jrandom
import equinox as eqx
import optax

data_key,model_key = jax.random.split(jax.random.PRNGKey(0),2)

### Generating toy-data

X_train = jax.random.normal(data_key, (1000,2))
y_train = X_train[:,0]+X_train[:,1]
y_train = jnp.where(y_train>0.5,1,0)

### Using equinox and optax
print("Training using equinox and optax")

epochs = 10000             
learning_rate = 0.1
n_inputs = X_train.shape[1]

class Logistic_Regression(eqx.Module):
    weight: jax.Array
    bias: jax.Array
    def __init__(self, in_size, out_size, key):
        wkey, bkey = jax.random.split(key)
        self.weight = jax.random.normal(wkey, (out_size, in_size))
        self.bias = jax.random.normal(bkey, (out_size,))
        #self.weight = jnp.zeros((out_size, in_size))
        #self.bias = jnp.zeros((out_size,))
    def __call__(self, x):
        return jax.nn.sigmoid(self.weight @ x + self.bias)

@eqx.filter_value_and_grad
def loss_fn(model, x, y):
    pred_y = jax.vmap(model)(x) 
    return -jnp.mean(y * jnp.log(pred_y) + (1 - y) * jnp.log(1 - pred_y))

@eqx.filter_jit
def make_step(model, x, y, opt_state):
    loss, grads = loss_fn(model, x, y)
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    return loss, model, opt_state

in_size, out_size = n_inputs, 1
model = Logistic_Regression(in_size, out_size, key=model_key)
optim = optax.sgd(learning_rate)
opt_state = optim.init(model)
for epoch in range(epochs):
    loss, model, opt_state = make_step(model,X_train,y_train, opt_state)
    loss = loss.item()
    if (epoch+1)%1000 ==0:
        print(f"loss at epoch {epoch+1}:{loss}")

# The following code is implementation of Logistic regression using scikit-learn and pytorch, and it is working well. It is added just for reference


### Using scikit-learn
print("Training using scikit-learn")
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
model = LogisticRegression()
model.fit(X_train,y_train)
y_pred = model.predict(X_train)
print("Train accuracy:",accuracy_score(y_train,y_pred))

## Using pytorch
print("Training using pytorch")
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD
from torch.nn import Sequential

X_train = np.array(X_train)
y_train = np.array(y_train)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("device:",device)
torch_LR= Sequential(nn.Linear(n_inputs, 1),
                nn.Sigmoid())
torch_LR.to(device)
criterion = nn.BCELoss() # define the optimization
optimizer = SGD(torch_LR.parameters(), lr=learning_rate)

train_loss = []
for epoch in range(epochs):
    inputs, targets = torch.tensor(X_train).to(device), torch.tensor(y_train).to(device) # move the data to GPU if available
    optimizer.zero_grad() # clear the gradients
    yhat = torch_LR(inputs.float()) # compute the model output
    loss = criterion(yhat, targets.unsqueeze(1).float()) # calculate loss
    #train_loss_batch.append(loss.cpu().detach().numpy()) # store the loss
    loss.backward() # update model weights
    optimizer.step()
    if (epoch+1)%1000 ==0:
        print(f"loss at epoch {epoch+1}:{loss.cpu().detach().numpy()}")

内容的提问来源于stack exchange，提问作者prahasanam_boi