Financial Statement Analysis with Mechanistic Interpretability

Large language model for predicting earnings direction from financial statements, with comprehensive attention analysis and mechanistic interpretability tools.

Quick Start

# Install dependencies (using UV - 10-100× faster)
curl -LsSf https://astral.sh/uv/install.sh | sh
uv pip install -e .

# Run inference with automatic appendix and visualization generation
python src/inference/run_inference.py --mode ray --num-gpus 4 --num-samples 100

# Outputs:
# - predictions_with_full_matrix_*.json (predictions + attention)
# - appendices/ (accuracy tables, attention analysis)
# - visualizations/ (heatmaps)

---

Inference Engines

The project supports multiple inference engines for different analysis needs:

1. attention_full (Default - Full Attention Matrix)

Best for: Comprehensive attention analysis, research, publication-quality results

python src/inference/run_inference.py --engine attention_full --mode ray --num-gpus 4

Features:

• ✅ Extracts complete attention matrix from last layer
• ✅ Word-level attention aggregation
• ✅ Multi-token keyword detection (handles "de" + "crease")
• ✅ Compatible with appendix/visualization generation
• ⚠️ Higher memory usage (~3-5GB per GPU)

---

2. attention_rollout (Advanced - Multi-Layer Analysis)

Best for: Deep mechanistic interpretability, layer-wise attention patterns

python src/inference/run_inference.py --engine attention_rollout --mode ray --num-gpus 4

Features:

• ✅ Rolls out attention across multiple layers
• ✅ Configurable layer selection (early/mid/late or custom)
• ✅ Head fusion strategies (mean/max/min)
• ✅ Residual connection handling
• ✅ Compatible with appendix/visualization generation
• ⚠️ Slower than full_matrix (processes multiple layers)

Configuration (config/attention_config.yaml):

rollout:
  enabled: true
  layers: ["early", "mid", "late"]  # or [0, 5, 10, 15, 20, 25]
  head_fusion: "mean"
  include_residuals: true
  residual_alpha: 0.5

---

3. confidence_only (Lightweight - Fast Predictions)

Best for: Quick testing, confidence scoring, large-scale prediction runs

python src/inference/run_inference.py --engine confidence_only --mode ray --num-gpus 4

Features:

• ✅ Fast inference (no attention extraction)
• ✅ Confidence scores (log probabilities)
• ✅ Low memory usage (~2GB per GPU)
• ❌ No attention analysis
• ❌ Appendix/visualization will be skipped

---

Running Inference

Command Structure

python src/inference/run_inference.py [OPTIONS]

Key Arguments

Argument	Options	Default	Description
--mode	cpu, gpu, ray	ray	Execution mode
--engine	attention_full, attention_rollout, confidence_only	attention_full	Inference engine
--num-gpus	Integer (2-8)	4	Number of GPUs for Ray mode
--num-samples	Integer	All data	Limit number of samples
--use-full-dataset	Flag	False	Use complete dataset
--checkpoint-every	Integer	50	Checkpoint frequency
--output-dir	Path	outputs/	Output directory

Execution Modes

1. Ray Mode (Multi-GPU - Recommended)

When to use: Production runs, large datasets (>1000 samples), automatic appendix/viz generation

# Standard production run
python src/inference/run_inference.py --mode ray --num-gpus 4

# Quick test with 10 samples
python src/inference/run_inference.py --mode ray --num-gpus 2 --num-samples 10

# Full dataset with checkpointing
python src/inference/run_inference.py --mode ray --num-gpus 8 --use-full-dataset --checkpoint-every 100

Automatically generates:

1. Predictions JSON
2. Appendix A (accuracy tables)
3. Appendix B (attention analysis)
4. Appendix C (example predictions)
5. Heatmap visualizations (4 images)

---

2. Single GPU Mode

When to use: Testing, debugging, single GPU systems

python src/inference/run_inference.py --mode gpu --num-samples 100

---

3. CPU Mode

When to use: Quick debugging, no GPU available

python src/inference/run_inference.py --mode cpu --num-samples 5

---

Configuration Files

config/model_config.yaml - Model and Inference Settings

# Model selection
model_name: "Qwen/Qwen2.5-1.5B-Instruct"
model_family: "qwen"

# Inference engine
inference:
  engine: "attention_full"  # Options: attention_full, attention_rollout, confidence_only
  mode: "ray"

# GPU settings
gpu:
  num_gpus: 4
  device_map: "auto"
  torch_dtype: "bfloat16"

# Inference parameters
inference_params:
  max_new_tokens: 100
  temperature: 0.0
  top_p: 0.95
  do_sample: false

config/attention_config.yaml - Attention Extraction Settings

# Basic extraction
extraction:
  layers_to_extract: "last_1"  # Options: "last_1", "last_6", "all", or [0, 11, 23]
  save_raw_attention: false

# Rollout configuration (for attention_rollout engine)
rollout:
  enabled: true
  layers: ["early", "mid", "late"]  # or specific indices [0, 5, 10, ...]
  head_fusion: "mean"  # Options: mean, max, min
  include_residuals: true
  residual_alpha: 0.5

# Visualization
visualization:
  colormap: "viridis"
  default_figsize: [12, 8]

config/prompts.yaml - Prompt Templates

financial_analysis: |
  Below is the Balance Sheet:
  {balance_income_sheet}

  Solve this problem: Will earnings increase, stay the same, or decrease?

---

Output Structure

After running inference, you'll get:

outputs/
├── predictions_with_full_matrix_TIMESTAMP.json    # Main predictions file
│
├── appendices/                                     # Research appendices
│   ├── appendix_a_accuracy.txt                    # Time series accuracy
│   ├── appendix_b_with_attention.txt              # Attention analysis
│   └── appendix_c_examples.txt                    # Example predictions
│
└── visualizations/                                 # Heatmaps
    ├── balance_sheet_correct_heatmap.png
    ├── balance_sheet_incorrect_heatmap.png
    ├── income_statement_correct_heatmap.png
    └── income_statement_incorrect_heatmap.png

---

Common Use Cases

1. Production Run (Full Dataset)

python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 8 \
  --use-full-dataset \
  --engine attention_full \
  --checkpoint-every 100

Output: Complete predictions + appendices + visualizations

---

2. Quick Testing (10 Samples)

python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 2 \
  --num-samples 10 \
  --engine attention_full

Output: Fast test with all outputs

---

3. Attention Rollout Analysis

python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 4 \
  --num-samples 500 \
  --engine attention_rollout

Output: Multi-layer attention analysis

---

4. Fast Predictions Only (No Attention)

python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 4 \
  --num-samples 1000 \
  --engine confidence_only

Output: Predictions with confidence scores (no appendices/viz)

---

5. Custom Output Directory

python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 4 \
  --output-dir /custom/path/results/

Output: All files saved to custom directory

---

Evaluation

After inference completes, evaluation runs automatically in Ray mode.

For manual evaluation:

python src/inference/evaluate.py outputs/predictions_*.json

Metrics calculated:

• Accuracy (overall and per class)
• F1 Score (macro/weighted)
• Precision and Recall
• Confusion Matrix

---

Appendix Generation

Automatic (Ray Mode)

Appendices generate automatically after Ray inference.

Manual (Any Mode)

python test_appendix_pipeline.py outputs/predictions_*.json

Generates:

• Appendix A, B, C
• Heatmap visualizations

---

Visualization

Automatic (Ray Mode with Attention Engines)

Visualizations generate automatically when using attention_full or attention_rollout.

Manual Attention Visualization

python visualize_attention.py outputs/predictions_*.json

Generates: Bar charts of top words by attention score

---

HPC Cluster Usage

SLURM Job Submission

# Multi-GPU job (2-8 H100s)
sbatch jobs/submit_inference_multi_gpu.slurm

# Single-GPU job
sbatch jobs/submit_inference.slurm

Customizing SLURM Scripts

Edit jobs/submit_inference_multi_gpu.slurm:

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=32
#SBATCH --gres=gpu:H100:4        # Change GPU count
#SBATCH --time=24:00:00
#SBATCH --mem=128GB

# Modify inference command
python src/inference/run_inference.py \
  --mode ray \
  --num-gpus 4 \              # Match SLURM GPU count
  --num-samples 5000          # Adjust sample size