Machine learning approach for predicting set members

Question

Say I have a large training dataset containing sets of 40 items each, and each item in the set is unique (so every training input is a set $S=\{i_1, i_2, ..., i_{40}\}$), and there are more than 40 unique items that can be part of a set.

I would like to be able to predict which items are a probably member of a set, given some incomplete set. So let's take the following example:

Training data: $S_1 = \{1,2,3\}$, $S_2 = \{3,4,5\}$, $S_3 = \{6,7,8\}$
Say I then have an input $S_4 = \{3\}$, I would want the approach to give back that 1, 2, 4 and 5 are more probable set members than 7, 8. Ideally with some probability value.

I've considered the following:
Using the apriori algorithm to learn some association rules. I wasn't sure how to interpret the support or lift as a probability of set membership.

Training a Multilayer perceptron on the input (probably one-hot encoded) to learn weights corresponding to the various input items. However, if I were to simply give the 40-item sets as input and output then the network would just learn to copy the input, giving no information about possible other set members. I've thought about giving all variations of the 40 item set as input, with the 40 item set as output, but this would result in $2^{40}$ possibilities per input which would be massive.

Is there some machine learning approach or data structure that could help in this situation?

Answer 1

You could train an embedding model. Each element would be projected onto a location a vector space based on its co-occurrence with other elements. Then finding similar elements could be done with a nearest neighbor search.

Answer 2

Maybe it's a bit overkill and biased toward my own field (neural machine translation), but you could go with a neural network architecture with self-attention in a masked language model-ish (i.e. BERT) configuration

The input to the network would be a fixed-size (40) sequence of discrete symbols meaning whether the element at that position is either present in the set ($P$), absent from it ($A$) or its presence is unknown ($U$). The "vocabulary" of the input tokens would then be ${P, A, U}$. This way, the input would be a sequence of 40 symbols, e.g. $A, A, P, U, U,..., P$.

The input would be fed into an embedding layer, followed by $N$ layers of unmasked self-attention.

Finally, the last attention layer output would be projected into a representation space of size 40, which would then be applied a sigmoid function to obtain the probabilities of each of the 40 elements in the set be part of the original input.

The loss function would only be computed on the elements where the input was marked as unknown ($U$) and ignored elsewhere. The expected output at those positions would be $1$ if the element was actually present or else $0$. You could use binary cross-entropy as the loss function to optimize.

You should prepare your training data so that you mark as unknown ($U$) elements from the original set, with the unknown element ratio that you expect in your test data.

At inference time, you simply set the info you know as either $P$ or $A$ and set the unknown elements to $U$, and get the output of the network at those positions.

Answer 3

Notice that from any complete set, you can form an question-answer pair for your task: you can construct an incomplete set (the "question"), for which you know the desired complete set (the "answer"). In particular, given a complete set $S$, you can randomly choose a subset $T$ of $S$ of a particular size, and call $T$ the incomplete set. In this way, from your dataset of complete sets, you can form a very large training set of input-output pairs $(T,S)$.

Next, given this training set, I suggest training any ML model to predict $S$ from $T$. There are many possibilities for how you might instantiate the model.

For instance, let's suppose the universe (all possible elements that could appear in a set) is not too large, say $m$ items. Then, you can encode a set using a one-hot encoding. Now, you can construct a neural network whose input is the one-hot encoding of $T$ and whose output is a $m$-dimensional vector $(p_1,\dots,p_m)$ such that $p_1+\dots+p_m=1$ (this vector can be generated by a final softmax layer). You can now interpret $p_i$ as the probability that item $i$ should be included in the final set. In other words, you can form a final set by including item $i$ with probability $p_i$ and omitting it with probability $1-p_i$ (making an independent choice for each $i$). Now, out of the distribution on sets induced in this way, we want to restrict to those of size 40, i.e., condition on the event that the resulting set has size 40. Given a target set $S$, it is possible to compute explicitly the probability of obtaining that particular set: it is the coefficient of $x^{40}$ in the polynomial $(p_1 x + 1-p_1) (p_2 x + 1-p_2) \cdots (p_m x + 1 - p_m)$, which can be computed relatively efficiently. It is also a differentiable function of $p_1,\dots,p_m$. Thus, you can define your loss function to be the negative log of this probability, and you can use stochastic gradient descent to train your network and find parameters that minimize this loss. As an optimization, you might force the output set to automatically include the input $T$, and use softmax only over the elements not in $T$.

Answer 4

You can take a purely count-based approach, commonly called maximum likelihood estimation (MLE).

Go through the sets and count the frequency of co-occurrence. Then find the most frequent items.

Here is a straight-forward (but not very optimized) solution in Python:

from collections import defaultdict
from itertools   import permutations 

sets = [{1, 2, 3}, {3, 4, 5}, {6, 7, 8}, {3, 4, 5}]

counts = defaultdict(lambda: defaultdict(int))

# Create the frequency count of co-occurance of pairs
for s in sets:
    perms = permutations(s, 2)
    for p in perms:
        counts[p[0]][p[1]] += 1

# For a given element, order by co-occurance frequency 
key = 3
counts_ordered = dict(sorted(counts[key].items(), 
                      key=lambda x: x[1],
                      reverse=True))

# Find probability
denominator = sum(counts_ordered.values())
print(f"For element {key}:") 
for k,v in counts_ordered.items():
    print(f"The probability of {k} co-occuring is {v/denominator:.3}.")
print("All other probabilities are zero.")

For element 3:
The probability of 4 co-occuring is 0.333.
The probability of 5 co-occuring is 0.333.
The probability of 1 co-occuring is 0.167.
The probability of 2 co-occuring is 0.167.
All other probabilities are zero.