Embodiments described herein adopts a k nearest neighbor (kNN) mechanism over a model's hidden representations to identify training examples closest to a given test example. Specifically, a training set of sequences and a test sequence are received, each of which is mapped to a respective hidden representation vector using a base model. A set of indices for each sequence index that minimizes a distance between the respective hidden state vector and a test hidden state vector is then determined A weighted k-nearest neighbor probability score can then be computed from the set of indices to generate a probability distribution over labels for the test sequence.