Memory-efficient GLM fitting with convolutions and multiple epochs

[BalzaniEdoardo]

Background

Currently, NeMoS handles convolutions by pre-computing a full design matrix X from basis convolutions, then fitting via loglike(inv_link(X @ w)). This has two efficiency issues:

1. Memory: For long time series, the design matrix X can be very large
2. Multiple epoch handling: When fitting discontinuous trials of varying lengths, we currently loop over epochs and convolve one at a time, which is inefficient when convolution happens inside the optimization loop

Proposed solution: Generalized linear operators

Instead of hard-coding the dot product, generalize the GLM objective to:

loglike(inv_link(linear_operator(x, w)))

Where linear_operator can flexibly combine different operations (convolution, dot product, etc.) while maintaining linearity.

Example: For a mixed design with one convolution feature and other regular features:

# Conceptually:
L(x, w) = convolve(x[:, slice_x_1], w[slice_w_1]) + x[:, slice_x_2] @ w[slice_w_2]

# Where:
# - len(slice_w_1) determines convolution kernel size (unrelated to len(slice_x_1))
# - len(slice_x_2) must equal len(slice_w_2) for dot product compatibility

Example 2: Convolution with basis

# B is a basis function kernel of shape (window_size, num_bases)
# Instead of convolving each basis separately and forming design matrix:
#   X = convolve(x, B[:, 0]) | convolve(x, B[:, 1]) | ... (large matrix)
# We compute: convolve(x, B @ w) directly

lin_operator = lambda x, w: convolve(x, B @ w)
L(x, w) = lin_operator(x[:, slice_x_1], w[slice_w_1]) + x[:, slice_x_2] @ w[slice_w_2]

Advantages:

• Memory: O(n log n) convolution without forming large design matrix
• Speed: FFT-based convolution is faster than matrix multiplication
• Flexibility: Can mix convolution with other linear operations

Challenge: Multiple epochs with varying lengths

People fit discontinuous time series where each epoch represents a trial. Trials may vary in duration. Currently using pynapple to chunk time series and convolving one trial at a time in a loop - fine for one-time convolution, but inefficient inside optimization.

Proposed strategy (discussed with Sarah):

• Group epochs by similar length
• Pad within groups to enable vectorized convolution
• Loop over small number of groups instead of all epochs

Open questions:

• How many groups is optimal?
• What's a good criterion for "similar enough" length? (e.g., within 10%? fixed bin sizes?)
• Trade-off between padding overhead vs. number of groups

Design considerations

Interface challenge: How do users specify custom linear operators intuitively?

Compatibility: Need to maintain current flexibility where:

• GLM is agnostic about design matrix
• Basis objects handle their own transformations
• Users can provide arbitrary designs

Software engineering needs:

1. Refactor GLM internals to support pluggable linear operators
2. Implement efficient multi-epoch batching/grouping
3. Design user-friendly API for specifying operators

[BalzaniEdoardo]

@pgree, I tried to summarize our discussion here, let me know if that is what you had in mind. As you may tell from the issue description, there is a bit of software engineering to be done before we can include the feature into NeMoS but that all sounds promising.

@sjvenditto, once we have a general way to handle epochs efficiently, we can generalize our loss function. Internally, that should not be too hard: we need a general BaseRegressor._predict method which takes a linear operator as input.

The hard part is the interface design: how do we map which coefficients are used for a convolve operation vs. a dot product? And how do we check that the coefficient structure and the operator are compatible?

For example, if we have:

L(x, w) = convolve(x[:, 0], w[:k]) + x[:, 1:] @ w[k:]

We need some way to specify that:

• First k coefficients go to convolution with column 0
• Remaining coefficients go to dot product with columns 1:

Do you have thoughts on how users should specify this mapping? Should it be implicit (inferred from basis structure) or explicit (user provides a mapping)?