Comment by cooljoseph

Last year I also had a rather similar idea, but dropped it before I went very far in working on it. I wonder if you and I had similar ideas?

1. Start with an LLM, so that your model understands natural language.

2. Replace RoPE with a tree embedding scheme, and causal attention with a sparse attention on the graph structure. (You could use full attention... but it's cheaper to use graph attention.)

3. Chop off the final unembedding layer, replacing it with a projection down to two scalars, one for logits and one for confidence.

4. Each option of a choice is represented by a number of tokens in leaf position; average these tokens' logit outputs to get the option's logit. Average all of the confidences from all of the options to get the choice's confidence.

5. Train the logits by KL divergence from a true distribution (or NLL on samples from a true distribution).

6. Train the confidences on a subset of the data in which you know the entire true distribution.

The hardest part is getting real world data for workflows, but I wildly speculate that you can get by with only ~50,000 documents if you first adapt domains using synthetic data.