Refactor DAG shortest path algorithm using Viterbi - #777
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This is an excellent pull request that refactors the shortest path algorithm to use the Viterbi algorithm, resulting in a significant performance improvement and much cleaner, more maintainable code. The new logic is well-encapsulated within the Viterbi class. I have two suggestions: one critical fix to prevent a crash when no valid path is found, and one high-priority comment to remove some dead code left over from the refactoring.
|
This in an impressive PR! I once had the 2nd ed. of Jurafsky and Martin on my bookshelf for a class, but I regret I had never grokked most of the book :|. I think I'll want to spend some time on the code. Before that, one quick observation: you pointed out that the grid is already in topological order, and since the relaxing step in Cormen has the same runtime ( On my mid-range Intel Linux laptop, the stress test in debug mode had roughly the same run time (8000 us) as yours, of which ~1000 us (12%) was spent on the actual relaxing, and the rest was taken up by the graph-building and the topological search, which now seems totatlly unnecessary. I can't but wonder if just by simplifying that step we'd get the same result. I also wonder what your numbers will be like if you run your PR in release mode. On my laptop the original code ran at 4x faster (~2100 us), and even less time was spent on the relaxing step (150 us, or 6%), which means the prep+sorting step was even more relatively expensive. |
I completely agree that we should take our time with this. Since this is the core algorithm, it’s crucial to ensure the implementation is robus. Your assessment was spot on! I ran the experiments in Release mode as you suggested, and the results confirm your observations. The overhead of building the graph structure is indeed the primary bottleneck.
I was not sure how the original graph-building step could be simplified within the existing architecture for relaxing, so I only removed the topological sort and performed the relax step directly. Since topological sort is not the primary cause of the latency, even after removing it from the original code, it still takes ~700 us. This proves that bypassing the explicit graph construction provides the most significant performance gain. ViterbiReadingGrid::walk (All Steps)int64_t start = GetEpochNowInMicroseconds();
Viterbi viterbi(spans_, readings_.size());
viterbi.ForwardPass(result);
viterbi.BackwardPass(result);
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
1. Forward Pass (Relax)int64_t start = GetEpochNowInMicroseconds();
viterbi.ForwardPass(result);
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
2. Backward Passint64_t start = GetEpochNowInMicroseconds();
viterbi.BackwardPass(result);
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
Original CodeReadingGrid::walk (All Steps)
1. Preparationint64_t start = GetEpochNowInMicroseconds();
std::vector<VertexSpan> vspans(spans_.size(), VertexSpan());
size_t vertices = 0;
size_t edges = 0;
for (size_t i = 0, len = spans_.size(); i < len; ++i) {
const ReadingGrid::Span& span = spans_[i];
for (size_t j = 1, maxSpanLen = span.maxLength(); j <= maxSpanLen; ++j) {
NodePtr p = span.nodeOf(j);
if (p != nullptr) {
vspans[i].emplace_back(std::move(p));
++vertices;
}
}
}
result.vertices = vertices;
Vertex terminal(std::make_shared<ReadingGrid::Node>(
"_TERMINAL_", 0, std::vector<LanguageModel::Unigram>()));
for (size_t i = 0, vspansLen = vspans.size(); i < vspansLen; ++i) {
for (Vertex& v : vspans[i]) {
size_t nextVertexPos = i + v.node->spanningLength();
if (nextVertexPos == vspansLen) {
v.edges.push_back(&terminal);
continue;
}
for (Vertex& nv : vspans[nextVertexPos]) {
v.edges.push_back(&nv);
++edges;
}
}
}
result.edges = edges;
Vertex root(std::make_shared<ReadingGrid::Node>(
"_ROOT_", 0, std::vector<LanguageModel::Unigram>()));
root.distance = 0;
for (Vertex& v : vspans[0]) {
root.edges.push_back(&v);
}
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
2. TopologicalSortint64_t start = GetEpochNowInMicroseconds();
std::vector<Vertex*> ordered = TopologicalSort(&root);
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
3. Relaxint64_t start = GetEpochNowInMicroseconds();
for (auto it = ordered.rbegin(), rend = ordered.rend(); it != rend; ++it) {
Vertex* u = *it;
for (Vertex* v : u->edges) {
Relax(u, v);
}
}
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
Original Code & Remove Topological Sortint64_t start = GetEpochNowInMicroseconds();
std::vector<VertexSpan> vspans(spans_.size(), VertexSpan());
size_t vertices = 0;
size_t edges = 0;
for (size_t i = 0, len = spans_.size(); i < len; ++i) {
const ReadingGrid::Span& span = spans_[i];
for (size_t j = 1, maxSpanLen = span.maxLength(); j <= maxSpanLen; ++j) {
NodePtr p = span.nodeOf(j);
if (p != nullptr) {
vspans[i].emplace_back(std::move(p));
++vertices;
}
}
}
result.vertices = vertices;
Vertex terminal(std::make_shared<ReadingGrid::Node>(
"_TERMINAL_", 0, std::vector<LanguageModel::Unigram>()));
for (size_t i = 0, vspansLen = vspans.size(); i < vspansLen; ++i) {
for (Vertex& v : vspans[i]) {
size_t nextVertexPos = i + v.node->spanningLength();
if (nextVertexPos == vspansLen) {
v.edges.push_back(&terminal);
continue;
}
for (Vertex& nv : vspans[nextVertexPos]) {
v.edges.push_back(&nv);
++edges;
}
}
}
result.edges = edges;
Vertex root(std::make_shared<ReadingGrid::Node>(
"_ROOT_", 0, std::vector<LanguageModel::Unigram>()));
root.distance = 0;
for (Vertex& v : vspans[0]) {
root.edges.push_back(&v);
}
// 1. Relax edges from root
for (Vertex& v : vspans[0]) {
Relax(&root, &v);
}
// 2. Relax edges span by span
for (size_t i = 0; i < vspans.size(); ++i) {
for (Vertex& u : vspans[i]) {
for (Vertex* v : u.edges) {
Relax(&u, v);
}
}
}
std::vector<NodePtr> walked;
size_t totalReadingLen = 0;
Vertex* it = &terminal;
while (it->prev != nullptr) {
walked.push_back(it->prev->node);
it = it->prev;
totalReadingLen += it->node->spanningLength();
}
assert(totalReadingLen == readings_.size());
assert(walked.size() >= 2);
result.totalReadings = totalReadingLen;
result.nodes = std::vector<NodePtr>(walked.rbegin() + 1, walked.rend());
result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
|
There was a problem hiding this comment.
@ChiahongHong thank you very much! I've reviewed your code as well as the references you provided. The insight that the grid is always in topological order is very helpful, and this is an elegant solution. I've left some comments.
As an aside, I have tried to improve the original code by doing one step further than you did in the follow-up—instead of building the vertex graph, which is indeed the most expensive step, I have tried to bake the score and the prev fields into Node itself (a very unsightly thing in itself), and in the end the code would be "just" 10% slower than your PR (so 550 μs vs 500 μs in release builds).
The problem with the original approach is that during the relaxing step, each node still has to consider every incoming node, when in fact we are able to compute the top incoming node at each span location already (which is what your Viterbi implementation does!) because the grid is already in topological order. So while both algorithms run at O(|V| + |E|) times, the numbers of effective |V| and |E| each has to consider are different (as you also pointed out in the PR description), and that really explains the 10% difference.
While the original code already achieves sub-millisecond walks, this PR does it even faster with less code and less data: a really elegant solution!
|
@lukhnos Thank you for the detailed feedback! I have removed the As I’m following your recommended approach and don’t have strong preferences about the structure, please let me know if there are any other areas that need refinement or if I’ve missed anything. |
lukhnos
left a comment
lukhnos
left a comment
There was a problem hiding this comment.
Thanks again for your contribution!
Replace the DAG shortest-path algorithm documentation with the Viterbi implementation merged in #777. The new section documents the linear lattice forward pass, relaxation, and backward path reconstruction. - Replace 4-step DAG section with 2-step Viterbi (forward + backward) - Update code reference table: remove TopologicalSort/Relax rows - Update references: Jurafsky & Martin + vene.ro lattice Viterbi - Fix TOC link to match new heading - Bump version to 1.3 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Extract the Viterbi walk algorithm from ReadingGrid::walk() into a dedicated WalkStrategy class, enabling strategy-pattern extensibility. Changes: - Add WalkStrategy base class with ViterbiStrategy implementation - Use forward-pass DP (matching post-#777 algorithm) instead of explicit DAG construction + topological sort - Add fixedSpans support: blocked[] array constrains the walk to respect user-selected spans - Convert Span storage from fixed array to vector for dynamic span lengths based on language model maxKeyLength() - Add fixSpan()/clearFixedSpans() to ReadingGrid for structural override support - ReadingGrid::walk() delegates to the configured WalkStrategy Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Describe the forward-pass Viterbi DP walk (PR #777) with verified file/line references against current master, an O(|V| + |E|) complexity analysis matching the implementation comment, and measured stress-test numbers (vertices/edges from WalkResult). Replace the dropped WalkStrategy/fixedSpans walk-integration design with the actual candidate-override mechanism (overrideCandidate plus re-walk), and align the contextual user model section with the shipped ContextualUserModel design (PR #780): two-level interpolated Kneser-Ney with per-reading continuation normalization, wall-clock decay with a 5400-second half-life, LRU capacity bound, TSV persistence, and implicit base-LM fallback via empty suggestions. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2 participants
Recently, I’ve been reading the "NLP bible," Speech and Language Processing, which explains that for HMMs, the most standard and efficient algorithm is Viterbi.
Hidden Markov Models - https://web.stanford.edu/~jurafsky/slp3/A.pdf
Viterbi
The Viterbi algorithm is a DP approach used to find the most likely sequence of hidden states
↑↑↑ This looks a bit complex, so let’s consider a simpler version instead
iby relaxing all possible spans ending at that positionTopological Sort Removed
Also, I realized that
ReadingGriditself is a linear lattice, so there is no need to perform a topological sort. It is already inherently topologically ordered:iand ends ati + spanLenspanLen >= 1, an edge always points to a higher index (strictly forward)Performance
Before
After