Handle unknown nucleotides

[eric-czech]

I noticed in a closer look here that tokens other than ACTG aren't handled:

biofoundation/biofoundation/data.py

Lines 45 to 58 in e8ff2fe

	def transform_reflogprob_clm(
	example: dict[str, Any],
	tokenizer: PreTrainedTokenizerBase,
	) -> dict[str, Any]:
	pos = example["pos"]
	assert example["seq"][pos] in NUCLEOTIDES
	input_ids = tokenizer(example["seq"], return_tensors="pt")["input_ids"][0]
	ref = input_ids[pos].item()
	# Create 4 copies of the input sequence
	new_input_ids = input_ids.unsqueeze(0).repeat(len(NUCLEOTIDES), 1)
	for i, nuc in enumerate(NUCLEOTIDES):
	new_input_ids[i, pos] = tokenizer.encode(nuc)[0]
	ref = NUCLEOTIDES.index(example["seq"][pos])
	return dict(input_ids=new_input_ids, ref=ref)

The line NUCLEOTIDES.index(example["seq"][pos]) should fail on an unknown token. Is that handled some other way @gonzalobenegas?

Also side note: the line ref = input_ids[pos].item() is unnecessary (that ref value is not used).

[gonzalobenegas]

Thanks, Eric. From a domain logic viewpoint, I'm not sure computing the likelihood of the unknown token is an operation we should be supporting.

I think the model shouldn't fail if it encounters [UNK] as part of the context, but IMO they are not meaningful center tokens for this kind of eval. To start with they are often contiguous runs of N so they are trivially predicted, but more generally they reflect limitations of sequencing/assembly rather than actual biology.

My general expectations for a DNA model are only that it handles A, C, G, T.
N is often taken care in an ad hoc way rather than at the same level of a nucleotide:

• Kallisto replaces N with random nucleotides.
• Enformer: (A = [1,0,0,0], C = [0,1,0,0], G = [0,0,1,0], T = [0,0,0,1], N = [0,0,0,0]).

[eric-czech]

I'm not sure computing the likelihood of the unknown token is an operation we should be supporting.

Makes sense to me. A good action item then would be to fail a little more gracefully (i.e. not with a vague <index> not in list error).

Kallisto replaces N with random nucleotides.

Oh wow, good to know!

[eric-czech]

I'm not sure computing the likelihood of the unknown token is an operation we should be supporting.
Makes sense to me.

Thinking about that more for a second, I'm not so sure it makes as much sense with plant DNA. We run into unknown nucleotides pretty frequently [1] and kuleshov-group/cross-species-single-nucleotide-annotation definitely contains them. I'm not sure if any ended up in https://huggingface.co/datasets/plantcad/evolutionary-constraint-example. Either way, there are two clear cases that come up with compute_reflogprob_clm:

1. The target nucleotide is unknown
2. Some other nucelotide is uknown but still used in the sequence log prob calculation

Perhaps an error should be thrown for any N in any sequence, at any position rather than what I implied in #3 (comment).

[gonzalobenegas]

I think it's totally fine that N's are part of the context, this will happen many times if you are close to complex repeat regions, especially as you increase the context size. But I don't ever see the center N's as the focus of the question: is this conserved/functional? The answer is trivially no and I've always filtered them out in the past.

[gonzalobenegas]

In the case of https://huggingface.co/datasets/plantcad/evolutionary-constraint-example, there are 19 N's on the negative set.

[eric-czech]

this will happen many times

Is there any reason not to account for it in the joint probability for a sequence then? E.g. in this reduction all the log probs are summed, so every new token will necessarily make the joint probability go down (or stay the same). Given that N's are a kind of meaningless catch all with super low probability most of the time, I don't see how regions with many of them wouldn't highly bias the scores for those examples way downward.

To be more concrete, why not use reduction="mean" only across valid nucelotides rather than reduce="sum" over all of them?

[gonzalobenegas]

I think if you would be calculating e.g. perplexity on an entire sequence then yes, I'd exclude those log probs from the aggregation (I generally filter N's from training dataset beforehand - might not be a good idea if the context size is larger though). However, when calculating P(center | context) as in the evolutionary constraint task, when normalizing across the seq-level log probs of the 4 possible sequences (each with a different center nucleotide), the log probs of the context will in a sense cancel out.

[eric-czech]

the log probs of the context will in a sense cancel out.

That makes sense to me for the 4 alt sequences created from one example. It seems untrue to me across examples though, and that should matter for the final rankings no?

I generally filter N's from training dataset beforehand - might not be a good idea if the context size is larger though

👍

[gonzalobenegas]

I think if ultimately you are calculating P(center | context), equivalent to getting logprobs at center masked position in a masked LM, it wouldn't require modifications based on N's in the context?