result of populationGLM depends on number of features

[RobertoDF]

Hi,

According to what´s written here https://nemos.readthedocs.io/en/latest/how_to_guide/plot_03_population_glm.html :
"For an unregularized, Lasso, Ridge, or group-Lasso GLM, fitting a GLM one neuron at the time, or fitting jointly the neural population is equivalent. "

however in the following model

model = nmo.glm.PopulationGLM(     regularizer="Ridge",
    regularizer_strength=0.0001,
    # lowering the tolerance means that the solution will be closer to the optimum
    # (at the cost of increasing execution time)
    solver_kwargs=dict(tol=10**-12),

)

the result of predict method changes if I increase the number of neurons.

Environemnt:

• OS: Linux
• Python version 3.13.5
• JAX version 0.7.1
• NeMoS version 0.2.4 dev & 0.2.3

[sjvenditto]

Hi Roberto,

What this means is that fitting a population GLM will give the same coefficients and predictions for each neuron as if you had fitted each of them independently. However, the shape of the output of predict (and the coefficients, etc) will change since it is fitting more neurons, where a single fitted neuron would have shape (n_timesteps,) and a population fit would have shape (n_timesteps, n_neurons). What you should see is that the column of a single neuron in the population fit matches the prediction of that single neuron fitted by itself.

Does this answer your question? Or are you seeing a mismatch between the per-neuron results of the population GLM and single GLM?

[RobertoDF]

I meant that the prediction for the first neuron changes if I'm using 200 vs 300 neurons. Specifically if I plot predicted activity over time it becomes flatter, with a much smaller standard deviation. For this reason for now I'm batching neurons in 100s.

[sjvenditto]

Would you be able to share some code that reproduces this problem? I am unable to reproduce it using randomly generated data in the following script:

import jax 
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import matplotlib.pyplot as plt

import numpy as np

import nemos as nmo

np.random.seed(123)

n_features = 5
n_neurons = 300
n_samples = 500

# random design array. Shape (n_time_points, n_features).
X = 0.5*np.random.normal(size=(n_samples, n_features))

# log-rates & weights
b_true = np.zeros((n_neurons, ))
w_true = np.random.uniform(size=(n_features, n_neurons))

# generate counts (spikes will be (n_samples, n_features)
rate = jnp.exp(jnp.dot(X, w_true) + b_true)
spikes = np.random.poisson(rate)

# fit 300 neurons
model1 = nmo.glm.PopulationGLM(regularizer="Ridge",
    regularizer_strength=0.0001,
    solver_kwargs=dict(tol=10**-12),)
model1.fit(X, spikes)
y1 = model1.predict(X)

# fit first 200 neurons
model2 = nmo.glm.PopulationGLM(regularizer="Ridge",
    regularizer_strength=0.0001,
    solver_kwargs=dict(tol=10**-12),)
model2.fit(X, spikes[:, :200])
y2 = model2.predict(X)

# fit first neuron
model3 = nmo.glm.GLM(regularizer="Ridge",
    regularizer_strength=0.0001,
    solver_kwargs=dict(tol=10**-12),)
model3.fit(X, spikes[:, 0])
y3 = model3.predict(X)

# plot
plt.plot(y1[:50, 0], label='model1')
plt.plot(y2[:50, 0], label='model2')
plt.plot(y3[:50], label='model3')
plt.legend()

One possibility could be the numerical precision, which Jax defaults to 32-bit. You can try enabling 64-bit precision to see if that helps the issue

[RobertoDF]

Thanjks for proividing the code!!

you get some differences if you push it.

import jax 
jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import matplotlib.pyplot as plt

import numpy as np

import nemos as nmo

np.random.seed(123)

n_features = 130
n_neurons = 1300
n_samples = 7000

# random design array. Shape (n_time_points, n_features).
X = 0.5*np.random.normal(size=(n_samples, n_features))

# log-rates & weights
b_true = np.zeros((n_neurons, ))
w_true = np.random.uniform(size=(n_features, n_neurons))

# generate counts (spikes will be (n_samples, n_features)
rate = jnp.exp(jnp.dot(X, w_true) + b_true)
spikes = np.random.poisson(rate)


model1 = nmo.glm.PopulationGLM(regularizer="Ridge",
    regularizer_strength=0.0001,
    solver_kwargs=dict(tol=10**-12),)
model1.fit(X, spikes)
y1 = model1.predict(X)


model2 = nmo.glm.PopulationGLM(regularizer="Ridge",
    regularizer_strength=0.0001,
    solver_kwargs=dict(tol=10**-12),)
model2.fit(X, spikes[:, :200])
y2 = model2.predict(X)



# plot
plt.plot(y1[:20, 0], label='model1')
plt.plot(y2[:20, 0], label='model2')

plt.legend()

[RobertoDF]

[sjvenditto]

Yeah, it looks like these differences are still arising from accumulating precision error which scales with the number of neurons. This is because all the neurons are being fit jointly in a single solver for the PopulationGLM, which sums across them to compute the likelihood, thus compounding the error in numerical precision. This, however, provides a faster fit than fitting all the neurons independently. We can look into providing a vmapped option for the PopulationGLM to reduce the accumulating precision error at the cost of being slower, but ideally the fit differences are minor and wouldn't lead to differences in model interpretation. For now, fitting the neurons in smaller batches (as you are now) will also mitigate this issue. Thanks for bringing this to our attention!