Phase 1a.2: Merge 3-level hierarchy into dataset-causal baseline

[research-developer]

Summary

Merges the 3-level hierarchy implementation from dataset-causal-3level into the dataset-causal baseline branch. This is Phase 1a.2 of the comprehensive branch merge strategy (second sequential merge in Phase 1a).

Changes

• Enables 3-level hierarchy support for causal reasoning domain
• Brings in all main branch improvements that were missing from dataset-causal
• 93 files changed, 32,557 insertions, 53 deletions
• Includes: NSM-33/34 infrastructure, Modal setup, physics metrics, adaptive training, checkpoint management

Key Additions

Infrastructure:

• GitHub Actions workflows (Claude Code Review)
• Modal.com deployment scripts and documentation
• Makefile for automation
• Physics-based metrics and PID control
• Checkpoint management utilities

Domain Files:

• Knowledge graph dataset (from main)
• Planning dataset (from main)
• Preflight check system
• 6-level chiral architecture files

Documentation:

• Comprehensive notes on NSM-33, NSM-34, Phase 1.5
• Modal best practices and setup guides
• Physics isomorphism analysis

Validation

• ✅ All 27 causal dataset tests passing (pytest tests/data/test_causal_dataset.py)
• ✅ Clean fast-forward merge (no conflicts)
• ✅ Coverage: 97% on nsm/data/causal_dataset.py

Testing

pytest tests/data/test_causal_dataset.py -v
# 27 passed, 1 warning in 3.81s

References

• PRD: .claude/specs/merge-main-and-3level-branches/01-product-requirements.md
• Architecture: .claude/specs/merge-main-and-3level-branches/02-system-architecture.md
• Pre-merge tags: pre-merge-dataset-causal-20251024, pre-merge-dataset-causal-3level-20251024
• Merge strategy: Phase 1a - Sequential 3-level → 2-level baseline merges
• Previous: PR Phase 1a.1: Merge 3-level hierarchy into dataset-planning baseline #16 (Planning domain merge)

Next Steps

After approval:

1. Merge this PR into dataset-causal
2. Proceed to Phase 1b: Parallel merges of main → all exploration branches
   • Phase 1b.1: main → dataset-planning
   • Phase 1b.2: main → dataset-causal
   • Phase 1b.3: main → dataset-kg-3level

🤖 Generated with Claude Code

Co-Authored-By: Claude noreply@anthropic.com

[claude]

PR Review: Phase 1a.2 - Merge 3-Level Chiral Architecture to Causal Branch

Overview

This PR merges the 3-level hierarchy implementation from dataset-causal-3level into dataset-causal baseline. This is a massive merge (93 files, 32,557 insertions) bringing in complete Phase 1.5 chiral architecture, infrastructure improvements, and comprehensive documentation.

---

✅ Strengths

1. Excellent Architecture Design

• Chiral Dual-Trifold: The implementation in nsm/models/chiral.py (721 lines) is well-structured with clear separation between:
  • ChiralHingeExchange: Simple, validated fusion-based exchange mechanism
  • MinimalChiralModel: 3-level proof-of-concept
  • FullChiralModel: Complete 6-level implementation
• Mathematical rigor: Proper normalization/denormalization for scale-invariant exchange (lines 451-495)
• Empirically validated: Fusion variant won over attention (51.26% acc, 29.60% balance vs 53.10% acc, 87.48% collapse)

2. Comprehensive Preflight Checks (nsm/evaluation/preflight_checks.py)

• Addresses NSM-31 root causes systematically:
  • Dataset balance checking (prevents class collapse)
  • Cycle loss weight validation (prevents gradient dominance)
  • Learning rate safety checks
  • PyG extension verification
• Clear error messages citing NSM-31 analysis
• Well-tested with 13 unit tests in tests/test_preflight_checks.py

3. Robust Loss Function (nsm/training/chiral_loss.py)

• ChiralCompositeLoss: Properly weighted multi-component loss
  • Main task + auxiliary (0.3×) + cycle (0.01×) + optional diversity
  • Focal loss support for class imbalance
  • compute_class_balance_metrics: Monitoring for collapse detection
• Aligns with NSM-32 specifications

4. Excellent Documentation

• Comprehensive notes on NSM-31, NSM-32, Phase 1.5 decisions
• Modal deployment guides and best practices
• Physics isomorphism analysis
• Clear commit messages with rationale

---

⚠️ Issues & Recommendations

Critical Issues

1. Potential Gradient Instability in 6-Level Model (nsm/models/chiral.py:663-686)

Location: FullChiralModel.forward() cycle loss computation

# Line 665-671: Potential issue with gradient flow
x_l1_reconstructed_from_l3 = torch.zeros_like(x_l1)
x_l3_to_l2 = torch.zeros(num_l2_nodes, self.node_features, device=x_l1.device)
x_l3_to_l2[perm_l3] = x_l3_refined  # ⚠️ In-place assignment may break gradients

x_l2_to_l1 = torch.zeros_like(x_l1)
x_l2_to_l1[perm_l2] = self.reconstruct_l1_from_l3(x_l3_to_l2)  # ⚠️ Another in-place

Problem:

• In-place tensor assignments with indexing can cause gradient computation issues
• torch.zeros_like() followed by indexed assignment doesn't preserve computational graph properly
• May lead to vanishing gradients in upper trifold cycle loss

Recommendation:

# Use scatter with proper gradient tracking
x_l3_to_l2 = torch.zeros(num_l2_nodes, self.node_features, device=x_l1.device)
x_l3_to_l2.scatter_(0, perm_l3.unsqueeze(1).expand(-1, self.node_features), x_l3_refined)
# Or use masked operations that maintain gradients

2. Size Alignment Algorithm is Lossy (nsm/models/chiral.py:497-530)

Location: FullChiralModel._align_sizes()

# Lines 524-528: Simple nearest neighbor interpolation
indices = (torch.arange(num_large, device=x_small.device).float() * (num_small / num_large)).long()
indices = torch.clamp(indices, 0, num_small - 1)
x_aligned = x_small[indices]  # ⚠️ Information loss from naive upsampling

Problems:

1. Information loss: Nearest neighbor repeats values, doesn't interpolate
2. Non-learnable: Fixed mapping doesn't adapt to data distribution
3. Gradient issues: Repeated indexing can cause gradient accumulation issues
4. Unused parameter: perm_large is not used despite being passed in

Recommendations:

1. Use learned upsampling:

   self.align_layer = nn.Linear(node_features, node_features)
   x_aligned = self.align_layer(x_small).repeat_interleave(num_large // num_small, dim=0)

2. Or use pooling indices properly:

   # Leverage perm_large to reverse pooling operation
   x_aligned = torch.zeros(num_large, dim, device=x_small.device)
   x_aligned[perm_large[:x_small.size(0)]] = x_small

3. Numerical Stability Issues in Normalization (nsm/models/chiral.py:464-473)

min_val = x.min(dim=0, keepdim=True)[0]
max_val = x.max(dim=0, keepdim=True)[0]
scale = max_val - min_val
scale = torch.where(scale < 1e-8, torch.ones_like(scale), scale)  # ⚠️ Problematic
x_normalized = (x - min_val) / scale

Problem: When scale < 1e-8, replacing with 1.0 causes incorrect normalization:

• If original max_val ≈ min_val ≈ 0, output becomes x_normalized = -min_val (not in [0,1])
• If max_val ≈ min_val ≈ 100, output becomes (100 - 100) / 1.0 = 0 (loses information)

Recommendation:

# Use epsilon additive instead of replacement
eps = 1e-8
scale = torch.clamp(max_val - min_val, min=eps)
x_normalized = (x - min_val) / (scale + eps)  # Safer

---

High Priority Issues

4. Missing Input Validation

Location: ChiralHingeExchange.forward() (line 74-101)

No validation that x_upper and x_lower have same shape:

def forward(self, x_upper: torch.Tensor, x_lower: torch.Tensor):
    # ⚠️ Missing: assert x_upper.shape == x_lower.shape
    lower_transformed = self.transform_lower_for_upper(x_lower)  # Will fail silently if shapes mismatch

Recommendation: Add shape validation at start of forward pass.

5. Potential Division by Zero in Class Balance Metrics (chiral_loss.py:244-268)

# Line 255: Division without checking mask.sum() > 0 first
correct = (predictions[mask] == cls).sum().item()
total = mask.sum().item()
accuracy = correct / total  # ⚠️ Division by zero if total == 0

Fix: Already has if mask.sum() > 0 check, but should use it consistently (line 253).

6. Focal Loss Implementation Issue (chiral_loss.py:58-91)

# Line 79: Creates one-hot without dimension check
targets_one_hot = F.one_hot(targets, num_classes=logits.size(1)).float()

Problem: If targets contains values ≥ num_classes, will raise runtime error.

Recommendation: Add bounds checking:

assert targets.max() < logits.size(1), f"Invalid target {targets.max()} for {logits.size(1)} classes"

---

Medium Priority Issues

7. Inefficient Memory Allocation (chiral.py:665-671)

Multiple torch.zeros_like() allocations in hot path (forward). Consider pre-allocating buffers.

8. Hardcoded Magic Numbers

• Line 144: num_bases = max(1, num_relations // 4) - why 4? Should be configurable
• Line 660: (logits_l1 + logits_l2 + logits_l3) / 3 - hardcoded averaging, consider learned weights

9. Missing Type Hints

preflight_checks.py has good type hints, but chiral.py is missing return type hints on some methods (e.g., _align_sizes).

---

Low Priority / Style Issues

10. Inconsistent Docstring Style

• chiral.py: Uses Google style (good)
• Some functions: Missing "Raises" section where exceptions are raised

11. Overly Long Forward Method

FullChiralModel.forward() is 181 lines (532-713). Consider breaking into:

• _forward_upper_trifold()
• _forward_lower_trifold()
• _compute_hinge_exchanges()
• _compute_cycle_losses()

---

🧪 Test Coverage Assessment

Good Coverage:

• ✅ Preflight checks: 13 unit tests covering all validation paths
• ✅ Test data provided: 27 causal dataset tests passing

Missing Tests:

• ❌ No unit tests for ChiralHingeExchange: Should test shape handling, gradient flow
• ❌ No unit tests for MinimalChiralModel: Should test cycle loss computation
• ❌ No unit tests for FullChiralModel: Critical component with 180+ line forward pass
• ❌ No unit tests for ChiralCompositeLoss: Should test focal loss, diversity loss
• ❌ No gradient flow tests: Given NSM-31 issues, should verify gradient magnitudes

Recommendation: Add tests before merging, or create follow-up issue (NSM-XX) for comprehensive chiral architecture testing.

---

🔒 Security Concerns

✅ No security issues found:

• No external data loading without validation
• No eval/exec usage
• No credential handling
• Modal scripts properly use environment variables

---

⚡ Performance Considerations

Potential Bottlenecks:

1. Triple hinge exchange with normalization (lines 615-646)

   • 6 normalization operations per forward pass
   • Consider caching if repeated in training loop

2. Size alignment (lines 600, 610, 638)

   • Called 3 times per forward pass
   • Allocates new tensors each time

3. Multiple global pooling operations (lines 650-652)

   • Consider batching if possible

Estimated Compute:

• 6-level model: ~180K parameters (reasonable)
• 3-level model: ~44K parameters (very efficient)

---

📝 Documentation Quality

Excellent:

• ✅ Comprehensive docstrings with mathematical foundations
• ✅ Clear commit messages with Claude Code attribution
• ✅ Extensive notes/ directory with design decisions
• ✅ Modal deployment guides

Could Improve:

• Add performance benchmarks (training time, memory usage)
• Add migration guide from 2-level to 3-level architecture
• Document expected accuracy ranges for each domain

---

🎯 Final Recommendation

Decision: ⚠️ CONDITIONAL APPROVAL - Recommend addressing critical issues before merge

Must Fix Before Merge:

1. ✅ Gradient flow issue in cycle loss computation (Critical NSM-17: Implement R-GCN with Confidence-Weighted Message Passing #1)
2. ✅ Size alignment lossy interpolation (Critical NSM-15: Confidence Propagation Base #2)
3. ✅ Numerical stability in normalization (Critical NSM-16: Invertible Coupling Layers (RealNVP) #3)

Should Fix Soon (can be follow-up PR):

4. Add comprehensive unit tests for chiral components
5. Add shape validation in hinge exchange
6. Add gradient flow validation tests

Nice to Have:

7. Refactor long forward methods
8. Add learned weights for ensemble prediction
9. Add performance benchmarks

---

💡 Positive Observations

1. Empirical validation approach: Testing attention vs fusion in parallel was excellent
2. Failure analysis: NSM-31 documentation shows thorough debugging
3. Preflight checks: Proactive prevention of known failure modes
4. Code quality: Clean structure, good naming, comprehensive docstrings
5. Reproducibility: Modal scripts + git tags enable exact replication

---

📊 Merge Checklist

• ✅ All 27 causal dataset tests passing
• ✅ Clean fast-forward merge (no conflicts)
• ✅ Coverage: 97% on causal_dataset.py
• ✅ Documentation comprehensive
• ⚠️ Critical code issues need addressing
• ❌ Missing unit tests for new chiral components
• ✅ Follows project conventions (CLAUDE.md compliant)

---

Overall Assessment: This is high-quality research code with excellent architectural decisions and documentation. The critical issues are fixable and shouldn't delay merge significantly. The empirical validation approach (testing multiple variants) is exemplary.

Estimated time to address critical issues: 2-3 hours

Great work on Phase 1.5! The chiral architecture is a clever solution to the class collapse problem. 🎉

[claude]

Code Review: PR #17 - Phase 1a.2 Merge (dataset-causal-3level → dataset-causal)

Summary

This is a large-scale merge PR (93 files, 32,557 additions, 53 deletions) bringing 3-level hierarchy support and main branch improvements into the dataset-causal baseline. The PR represents Phase 1a.2 of the branch merge strategy.

Overall Assessment: ✅ Approve with minor recommendations

The code demonstrates strong architectural design, comprehensive testing (97% coverage on causal dataset), and follows NSM project conventions. All 27 causal dataset tests are passing.

---

🎯 Strengths

1. Architecture Quality

• nsm/models/hierarchical.py (lines 357-725): The NSMModel class elegantly implements both 2-level and 3-level hierarchies with backward compatibility
• Dual-pass architecture with learned fusion (lines 605-653) is well-designed with multiple fusion modes (equal, learned, abstract_only, concrete_only)
• Proper separation of concerns between WHY/WHAT operations
• Clear adherence to adjoint functor principles (WHY ⊣ WHAT)

2. Dataset Implementation

• nsm/data/knowledge_graph_dataset.py: Rich, realistic knowledge graph generation with 50+ predicates
• Proper hierarchical level separation (L1: facts, L2: types/categories)
• Multi-hop query generation and type consistency checking (lines 605-683)
• Good entity diversity with biographical, geographic, creative, and type triples

3. Testing & Validation

• Excellent test coverage (97% on causal_dataset.py)
• Physics-inspired metrics (test_physics_metrics.py) show innovative collapse detection
• Comprehensive validation suite included
• All 27 causal dataset tests passing

4. Infrastructure

• GitHub Actions workflow properly configured
• Modal.com deployment infrastructure well-documented
• Makefile provides good automation
• Checkpoint management utilities included

---

⚠️ Issues & Recommendations

Critical: Security Concern

File: .env.local (lines 4-11)

export NSM_REPO_ROOT="/Users/preston/Projects/NSM"
export NSM_WORKTREE_ROOT="/Users/preston/Projects"

Issue: This file contains hardcoded local paths specific to a developer machine and should NOT be committed to version control.

Recommendation:

1. Add .env.local to .gitignore immediately
2. Create .env.example template instead:

   # NSM Project Environment Configuration
   export NSM_REPO_ROOT="/path/to/your/nsm/repo"
   export NSM_BASELINES_FILE="${NSM_REPO_ROOT}/baselines.jsonl"
   export NSM_WORKTREE_ROOT="/path/to/your/worktrees"

3. Remove .env.local from this PR

High Priority: Code Quality

1. Edge Type Placeholder (hierarchical.py:202-208)

# Need edge types for abstract level - reindex from original
# For now, use placeholder edge types (all same type)
edge_type_abstract = torch.zeros(
    edge_index_abstract.size(1),
    dtype=torch.long,
    device=edge_index_abstract.device
)

Issue: Using placeholder edge types (all zeros) defeats the purpose of R-GCN's relational message passing. This loses semantic information during abstraction.

Recommendation:

• Implement proper edge type preservation during pooling (track which L1 edge types correspond to pooled edges)
• Or document why this is acceptable for current phase (if intentional simplification)
• Add TODO with Linear issue reference if deferring

2. Duplicate Edge Type Logic (hierarchical.py:558-565)

Same placeholder pattern appears again. Consider extracting to helper method:

def _infer_edge_types(self, edge_index: Tensor, device: torch.device) -> Tensor:
    """Infer or preserve edge types during pooling operations."""
    if edge_index.size(1) > 0:
        # TODO(NSM-XX): Implement proper edge type preservation
        return torch.zeros(edge_index.size(1), dtype=torch.long, device=device)
    return torch.tensor([], dtype=torch.long, device=device)

3. Magic Numbers in Cycle Loss Weighting (hierarchical.py:602)

cycle_loss = 0.7 * cycle_loss_l1 + 0.3 * cycle_loss_l2

Issue: Hardcoded weights without justification.

Recommendation:

• Move to constructor parameter with these as defaults
• Document rationale (why 70/30 split?)
• Consider making learnable or adaptive

Medium Priority: Performance

1. Inefficient Multi-hop Query Generation (knowledge_graph_dataset.py:605-653)

for _ in range(num_queries):
    entities_with_edges = [e for e in graph.keys() if len(graph[e]) > 0]  # Rebuilt each iteration!

Issue: List comprehension regenerated in every loop iteration.

Recommendation:

entities_with_edges = [e for e in graph.keys() if len(graph[e]) > 0]
for _ in range(min(num_queries, len(entities_with_edges))):
    if not entities_with_edges:
        break
    # ... rest of logic

2. Repeated Type Lookups

knowledge_graph_dataset.py uses self.entity_types.items() repeatedly. Consider caching or precomputing frequently accessed structures.

Low Priority: Code Style

1. Inconsistent Docstring Coverage

• ChiralHingeExchange (chiral.py:30-45) has excellent docstrings
• Some utility functions lack examples
• Consider adding "Raises" section where applicable

2. Type Hint Completeness

Most functions have good type hints, but a few are missing:

• Return types for some helper methods
• Consider using TypedDict for complex return dictionaries

---

🔍 Test Coverage Assessment

Excellent Coverage

✅ 27/27 causal dataset tests passing
✅ Physics metrics well-tested
✅ Preflight checks validated
✅ Data utilities covered

Gaps to Address (Future Work)

⚠️ No tests found for:

• NSMModel dual-pass mode fusion variants
• KnowledgeGraphTripleDataset.get_multi_hop_queries()
• Edge type preservation during pooling
• Learned fusion attention mechanism

Recommendation: Add integration tests for dual-pass variants in Phase 1b.

---

📊 Performance Considerations

1. Memory: 32k+ lines added. Ensure CI/CD has adequate memory (>=8GB as per CLAUDE.md)
2. 3-Level Hierarchy: Two sequential WHY/WHAT operations may slow inference. Consider profiling with torch.profiler on realistic batch sizes.
3. Modal Deployment: Infrastructure looks solid, but ensure timeout configurations account for larger model.

---

🔐 Security Assessment

1. Critical: .env.local exposure (see above)
2. GitHub Actions: Properly configured with minimal permissions
3. Secrets Management: CLAUDE_CODE_OAUTH_TOKEN correctly handled via secrets
4. No SQL injection risks: All data generation is synthetic
5. No external API calls without validation

---

📝 Documentation Quality

Excellent:

• MODAL_BEST_PRACTICES.md (918 lines) is comprehensive
• NSM-31-PARALLEL-EXPLORATION.md provides good context
• Physics isomorphism analysis well-documented

Minor Gap:

• Consider adding migration guide for users updating from 2-level to 3-level hierarchy

---

✅ Merge Readiness Checklist

• Tests passing (27/27)
• Code quality high (minor issues noted)
• Documentation comprehensive
• Architecture sound
• Security issue resolved (.env.local)
• Follows CLAUDE.md conventions
• No breaking changes to public API

---

🚀 Recommendations Before Merge

Must Fix

1. Remove .env.local or add to .gitignore + provide template

Should Fix (or create follow-up issues)

2. Implement proper edge type preservation during pooling (or document deferral)
3. Extract magic numbers to named constants/parameters
4. Optimize multi-hop query generation

Nice to Have

5. Add integration tests for dual-pass fusion modes
6. Profile 3-level hierarchy performance
7. Add migration guide documentation

---

🎓 Final Verdict

This is a high-quality merge that successfully brings 3-level hierarchy support to the causal dataset branch. The architecture is sound, testing is comprehensive, and the code follows project conventions.

Recommendation: ✅ Approve after fixing .env.local security issue

The critical security concern is easily remedied and should not block the merge once resolved. All other issues are minor and can be addressed in follow-up PRs or tracked as separate Linear issues for Phase 1b.

Excellent work on maintaining code quality across such a large merge! 🎉

---

Review generated by Claude Code via GitHub Actions