Stateful forward does not result in identical embeddings with different sequence lengths

[felicitas215]

Hello, thank you for developing Evo2. Creating DNA foundation models trained on such an extensive dataset is truly impressive.

I recently tried to use your 7B parameter model (evo2-7b) for inference on the first chromosome of Arabidopsis thaliana (~30 Mb). I have my own code to feed the DNA into Evo2 in smaller blocks and extract the embeddings from this layer: 'blocks.28.mlp.l3'. I'm currently only interested in the embeddings and not the final output. When I test different sequence lengths for example 2560 bp vs. 92160 bp, I notice that no matter if I use StripedHyena's stateless_forward or stateful_forward I get different results when comparing the embeddings created with these two sequence lengths/block sizes. For the stateful function I would have assumed that the embeddings for the entire chromosome are nearly identical, no matter which sequence length is used. I tested the similarity with Pearson's correlation and instead of values around 0.99 I find Pearson's correlation values ranging from 0.5 to 0.99 (mean: 0.91). I'm initializing the inference parameters like so:

from evo2 import Evo2
model = Evo2('evo2_7b')
inference_params = model.model.initialize_inference_params(max_seqlen=1048576)

After the data processing I use:

# StripedHyena forward function
model.forward(input_ids, inference_params_dict=inference_params)

If I overlooked how statefulness can be achieved in this repository or vortex, I apologize and would kindly ask you to point me to the right tutorial or code.

[sjy-eyujun]

Hello, I encountered a similar problem during my calculations. Could you please tell me how you segmented the DNA sequence? My sequence is also very long, but I'm unsure if using a fixed value to cut it will cause bias.

[jstjohn]

How are you comparing equality? For example are you only comparing the first 2560 positions in both cases, and when you extend the input length are you using the same start position? I would hope that the embeddings on the first 2560 positions match assuming you use the same DNA for those first 2560 positions regardless of how much more DNA you have in future positions, if that makes sense.

These models generate embeddings and outputs using some combination of all prior positions when generating the output for some position, so just double checking what your comparison is when doing the extension. This may be expected (and desirable) behavior depending on what you are doing.

[felicitas215]

I use an entire chromosome and break it into sequences of a specific length (see table). I'm aware that cutting the sequence like this can cause problems, especially when certain motifs one wants to find later are separated into different sequences. This setup was just used as an initial test.

	1	2
forward funtion	stateful	stateful
sequence length	21_384	92_160

Comparison (rowwise pearson correlation)

• the first 21384 embeddings are similar enough (0.982-0.998, mean 0.995) (see correlation_1 down below)
• the next 64152 embeddings (the next three sequences from 1 compared to their equivalent positions in the first sequence from 2, see correlation_2 below) are dissimilar (0.742-0.998, mean 0.899)

Code

def rowwise_pearson(A, B):
    A32 = A.astype(np.float32)
    B32 = B.astype(np.float32)

    A_mean = A32.mean(axis=1, keepdims=True)
    B_mean = B32.mean(axis=1, keepdims=True)

    A_centered = A32 - A_mean
    B_centered = B32 - B_mean

    num = np.sum(A_centered * B_centered, axis=1)
    denom = np.sqrt(np.sum(A_centered**2, axis=1) * 
                    np.sum(B_centered**2, axis=1)) + 1e-8

    return num / denom

 # shapes: embeddings_array_21k (1550, 21384, 4096), embeddings_array_92k (360, 92160, 4096)
correlation_1 = rowwise_pearson(embeddings_array_21k[0], embeddings_array_92k[0][:21384])
correlation_2 = rowwise_pearson(embeddings_array_21k[1:4].reshape((64152, 4096)), embeddings_array_92k[0][21384:85536])

I apoligize if I misunderstood the stateful_forward function and it's use case. I thought it meant that even when using shorter sequences, the model would retain the information from the previous sequence to a certain degree, so that the example computation (correlation_2) would show an average pearson correlation of 0.97-0.99. For my use case (calculating embeddings for entire chromosomes of varying lengths) this would be more desirable than an average pearson correlation of 0.899 or lower.